try @thetorpedodog advice re _flags

apis/python/src/tiledbsoma/__init__.py

@@ -134,19 +134,6 @@
  tiledbsoma_stats_json,
 )
 
-# Temporary for https://github.com/single-cell-data/TileDB-SOMA/issues/2407
-_new_shape_feature_flag = os.getenv("SOMA_PY_NEW_SHAPE") is not None
-
-
-def _new_shape_feature_flag_enabled() -> bool:
- """
- This is temporary only and will be removed once
- https://github.com/single-cell-data/TileDB-SOMA/issues/2407
- is complete.
- """
- return _new_shape_feature_flag
-
-
 # Load native libraries. On wheel builds, we may have a shared library
 # already linked. In this case, we can import directly
 try:

apis/python/src/tiledbsoma/_dataframe.py

@@ -14,7 +14,7 @@
 from somacore import options
 from typing_extensions import Self
 
-from tiledbsoma import _new_shape_feature_flag_enabled
+from tiledbsoma._flags import _new_shape_feature_flag_enabled
 
 from . import _arrow_types, _util
 from . import pytiledbsoma as clib

apis/python/src/tiledbsoma/_flags.py

@@ -0,0 +1,284 @@
+# Copyright (c) 2021-2023 The Chan Zuckerberg Initiative Foundation
+# Copyright (c) 2021-2023 TileDB, Inc.
+#
+# Licensed under the MIT License.
+
+"""Conversion to/from Arrow and TileDB type systems. Must be capable
+of representing full type semantics, and correctly performing a
+round trip conversion (e.g., T == to_arrow(to_tiledb(T)))
+
+Most primitive types are simple -- e.g., uint8. Of particular challenge
+are datetime/timestamps as TileDB has no distinction between a "datetime" and
+a "timedelta". The best Arrow match is TimestampType, as long as that
+TimestampType instance does NOT have a timezone set.
+
+Because of our round-trip requirement, all other Arrow temporal types
+are unsupported (even though they are just int64 under the covers).
+
+We auto-promote Arrow's string and binary to large_string and large_binary,
+respectively, as this is what TileDB stores -- a sequence of bytes preceded
+by a 64-bit (not 32-bit) length int.
+
+DataFrame-specific note: currently (as of 2.14), TileDB does not support
+Unicode array dimensions. All Arrow string types used in a DataFrame index
+columns (i.e., TileDB dimension) are coerced to ASCII. This equirement for
+ASCII-only dimensions will be relaxed in a future release. Unicode/UTF-8 is
+fully supported in SOMA DataFrame non-indexed columns.
+"""
+
+from typing import Any, Dict, Union
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+import pyarrow as pa
+
+_ARROW_TO_TDB_ATTR: Dict[Any, Union[str, TypeError]] = {
+ pa.string(): "U1",
+ pa.large_string(): "U1",
+ pa.binary(): "bytes",
+ pa.large_binary(): "bytes",
+ pa.timestamp("s"): "datetime64[s]",
+ pa.timestamp("ms"): "datetime64[ms]",
+ pa.timestamp("us"): "datetime64[us]",
+ pa.timestamp("ns"): "datetime64[ns]",
+ #
+ # Unsupported types in TileDB type system
+ pa.float16(): TypeError("float16 - unsupported type (use float32)"),
+ pa.date32(): TypeError("32-bit date - unsupported type (use TimestampType)"),
+ pa.date64(): TypeError("64-bit date - unsupported type (use TimestampType)"),
+}
+"""Dict of types unsupported by to_pandas_dtype, which require overrides for
+use in TileDB Attributes (aka DataFrame non-indexe columns).
+
+If the value is an instance of Exception, it will be raised.
+
+IMPORTANT: ALL non-primitive types supported by TileDB must be in this table.
+"""
+
+_PYARROW_TO_CARROW: Dict[pa.DataType, str] = {
+ pa.bool_(): "b",
+ pa.int8(): "c",
+ pa.int16(): "s",
+ pa.int32(): "i",
+ pa.int64(): "l",
+ pa.uint8(): "C",
+ pa.uint16(): "S",
+ pa.uint32(): "I",
+ pa.uint64(): "L",
+ pa.float32(): "f",
+ pa.float64(): "g",
+ pa.timestamp("s"): "tss:",
+ pa.timestamp("ms"): "tsm:",
+ pa.timestamp("us"): "tsu:",
+ pa.timestamp("ns"): "tsn:",
+}
+
+# Same as _ARROW_TO_TDB_ATTR, but used for DataFrame indexed columns, aka TileDB Dimensions.
+# Any type system differences from the base-case Attr should be added here.
+_ARROW_TO_TDB_DIM: Dict[Any, Union[str, TypeError]] = _ARROW_TO_TDB_ATTR.copy()
+"""Same as _ARROW_TO_TDB_ATTR, but used for DataFrame indexed columns, aka TileDB Dimensions.
+Any type system differences from the base-case Attr should be added here.
+"""
+_ARROW_TO_TDB_DIM.update(
+ {
+ pa.string(): "ascii", # TODO: temporary work-around until Dimension UTF8 support is available.
+ pa.large_string(): "ascii", # TODO: temporary work-around until Dimension UTF8 support is available.
+ }
+)
+
+
+def tiledb_type_from_arrow_type(
+ t: pa.DataType, is_indexed_column: bool = False
+) -> npt.DTypeLike:
+ """Given an Arrow type, return the corresponding TileDB type as a NumPy dtype.
+ Building block for Arrow-to-TileDB schema translation.
+
+ TileDB currently has different Unicode handling for dimensions and attributes.
+ Set the ``is_dimension`` parameter to True for indexed-column (AKA dimension)
+ rules, which currently requires all strings to be ASCII.
+
+ If type is unsupported, with raise a TypeError exception.
+
+ Args:
+ t:
+ Arrow DataType instance, e.g., pyarrow.int8().
+ is_indexed_column:
+ Use TileDB dimension type conversion rules.
+
+ Returns:
+ The numpy dtype corresponding to the ``t`` parameter.
+
+ Raises:
+ TypeError: if the type is unsupported.
+ """
+ if pa.types.is_dictionary(t):
+ t = t.index_type
+
+ arrow_to_tdb = _ARROW_TO_TDB_DIM if is_indexed_column else _ARROW_TO_TDB_ATTR
+ if t in arrow_to_tdb:
+ arrow_type = arrow_to_tdb[t]
+ if isinstance(arrow_type, Exception):
+ raise arrow_type
+ if arrow_type in ["ascii", "bytes"]:
+ return arrow_type
+ return np.dtype(arrow_type)
+
+ if not pa.types.is_primitive(t):
+ raise TypeError(f"Type {str(t)} - unsupported type")
+ if pa.types.is_timestamp(t):
+ raise TypeError("TimeStampType - unsupported type (timezone not supported)")
+ if pa.types.is_time32(t):
+ raise TypeError("Time64Type - unsupported type (use TimestampType)")
+ if pa.types.is_time64(t):
+ raise TypeError("Time32Type - unsupported type (use TimestampType)")
+ if pa.types.is_duration(t):
+ raise TypeError("DurationType - unsupported type (use TimestampType)")
+
+ # else lets try the default conversion path
+ try:
+ # Must force into a dtype to catch places where the Pandas type
+ # system has extra information that can't be expressed
+ return np.dtype(t.to_pandas_dtype())
+ except NotImplementedError as exc:
+ raise TypeError("Unsupported Arrow type") from exc
+
+
+def arrow_type_from_tiledb_dtype(
+ tiledb_dtype: npt.DTypeLike, bytes_are_ascii: bool = True
+) -> pa.DataType:
+ """Maps a TileDB dtype (``'bytes'``, ``'ascii'``, or an ``np.dtype``) to an Arrow type. Note that
+ when we read tiledb schema off storage, ``ascii`` and ``bytes`` both have ``dtype`` of `"S"`
+ which is equal to ``bytes`` -- so, the caller should disambgiuate.
+ """
+ if tiledb_dtype == "bytes":
+ if bytes_are_ascii:
+ return pa.large_string()
+ else:
+ return pa.large_binary()
+ elif tiledb_dtype == "ascii" or tiledb_dtype == np.dtype(str):
+ return pa.large_string()
+ else:
+ return pa.from_numpy_dtype(tiledb_dtype)
+
+
+def is_string_dtypelike(dtype: npt.DTypeLike) -> bool:
+ # Much of this (including the null-check) is to make the type-checker happy,
+ # as npt.DTypeLike is a complex union including 'str' and None.
+ if dtype is None:
+ return False
+ if dtype == "str":
+ return True
+ if isinstance(dtype, np.dtype):
+ return is_string_dtype(dtype)
+ return False
+
+
+def is_string_dtype(dtype: Any) -> bool:
+ return dtype.name in ["object", "string", "str32", "str64"]
+
+
+def df_to_arrow(df: pd.DataFrame) -> pa.Table:
+ """
+ Handle special cases where pa.Table.from_pandas is not sufficient.
+ """
+ nullable_fields = set()
+ # Not for name, col in df.items() since we need df[k] on the left-hand sides
+ for key in df:
+ # Make attributes nullable. Context:
+ # * df_to_arrow is _solely_ for use of tiledbsoma.io
+ # o Anyone calling the SOMA API directly has user-provided Arrow
+ # schema which must be respected
+ # o Anyone calling tiledbsoma.io -- including from_h5ad/from_anndata,
+ # and update_obs/update_var -- does not provide an Arrow schema
+ # explicitly. We compute an Arrow schema for them here.
+ # * Even when the _initial_ data is all non-null down a particular
+ # string column, there are two ways a _subsequent_ write can provide
+ # nulls: append-mode ingest, or, update_obs/update_var wherein the new
+ # data has nulls even when the data used at schema-create time was
+ # non-null.
+ # * We have no way of knowing at initial ingest time whether or not
+ # users will later be appending, or updating, with null data.
+ # * Note that Arrow has a per-field nullable flag in its schema metadata
+ # -- and so do TileDB array schemas.
+ #
+ # Note in particular this is for the use of tiledbsoma.io:
+ #
+ # * In the tiledbsoma API (e.g. DataFrame.create) the user passes an
+ # Arrow schema and we respect it as-is. They specify nullability, or
+ # not, as they wish.
+ # * In tiledbsoma.io, the user-provided inputs are AnnData objects.
+ # We compute the Arrow schema _for_ them. And we must accommodate
+ # reasonable/predictable needs.
+
+ nullable_fields.add(key)
+
+ # Handle special cases for all null columns where the dtype is "object"
+ # or "category" and must be explicitly casted to the correct pandas
+ # extension dtype.
+ #
+ # Note: with
+ # anndata.obs['new_col'] = pd.Series(data=np.nan, dtype=np.dtype(str))
+ # the dtype comes in to us via `tiledbsoma.io.from_anndata` not
+ # as `pd.StringDtype()` but rather as `object`.
+ if df[key].isnull().all():
+ if df[key].dtype.name == "object":
+ df[key] = pd.Series([None] * df.shape[0], dtype=pd.StringDtype())
+ elif df[key].dtype.name == "category":
+ df[key] = pd.Series([None] * df.shape[0], dtype=pd.CategoricalDtype())
+
+ # For categoricals, it's possible to get
+ # TypeError: Object of type bool_ is not JSON serializable
+ # deep within library functions. Debugging reveals that this happens when
+ # the df[key].values.ordered is of type np.bool_ rather than Python bool.
+ # So, we cast and reconstruct.
+ for key in df:
+ column = df[key]
+ if isinstance(column.dtype, pd.CategoricalDtype):
+ if hasattr(column.values, "categories"):
+ categories = column.values.categories
+
+ if hasattr(column.values, "ordered"):
+ ordered = bool(column.values.ordered)
+
+ df[key] = pd.Categorical(
+ values=column, categories=categories, ordered=ordered
+ )
+
+ arrow_table = pa.Table.from_pandas(df)
+
+ md = arrow_table.schema.metadata
+ md.update(dict.fromkeys(nullable_fields, "nullable"))
+ arrow_table = arrow_table.replace_schema_metadata(md)
+
+ # For tiledbsoma.io (for which this method exists) _any_ dataset can be appended to
+ # later on. This means that on fresh ingest we must use a larger bit-width than
+ # the bare minimum necessary.
+ new_map = {}
+ for field in arrow_table.schema:
+ if pa.types.is_dictionary(field.type):
+ old_index_type = field.type.index_type
+ new_index_type = (
+ pa.int32()
+ if old_index_type in [pa.int8(), pa.int16()]
+ else old_index_type
+ )
+ new_map[field.name] = pa.dictionary(
+ new_index_type,
+ field.type.value_type,
+ field.type.ordered,
+ )
+ else:
+ new_map[field.name] = field.type
+ new_schema = pa.schema(new_map, metadata=arrow_table.schema.metadata)
+
+ arrow_table = pa.Table.from_pandas(df, schema=new_schema)
+
+ return arrow_table
+
+
+def pyarrow_to_carrow_type(pa_type: pa.DataType) -> str:
+ try:
+ return _PYARROW_TO_CARROW[pa_type]
+ except KeyError:
+ raise TypeError(f"Invalid pyarrow type {pa_type}") from None

apis/python/src/tiledbsoma/_sparse_nd_array.py

@@ -26,7 +26,7 @@
 from somacore.options import PlatformConfig
 from typing_extensions import Self
 
-from tiledbsoma import _new_shape_feature_flag_enabled
+from tiledbsoma._flags import _new_shape_feature_flag_enabled
 
 from . import _util
 

apis/python/tests/test_dataframe.py

@@ -92,7 +92,7 @@ def test_dataframe(tmp_path, arrow_schema):
 
  # More to come on https://github.com/single-cell-data/TileDB-SOMA/issues/2407
  assert (
- sdf.tiledbsoma_has_upgraded_domain == soma._new_shape_feature_flag_enabled()
+ sdf.tiledbsoma_has_upgraded_domain == soma._flags._new_shape_feature_flag_enabled()
  )
 
  with pytest.raises(AttributeError):

apis/python/tests/test_shape.py

@@ -48,7 +48,7 @@ def test_sparse_nd_array_basics(
  # More to come on https://github.com/single-cell-data/TileDB-SOMA/issues/2407
  assert (
  snda.tiledbsoma_has_upgraded_shape
- == tiledbsoma._new_shape_feature_flag_enabled()
+ == tiledbsoma._flags._new_shape_feature_flag_enabled()
  )
 
  # Before current-domain support: shape is maxshape.
@@ -58,7 +58,7 @@ def test_sparse_nd_array_basics(
  # involving R compatibility, and leaving room for a single tile
  # capacity, etc ... we could check for some magic value but it suffices
  # to check that it's over 2 billion.)
- if tiledbsoma._new_shape_feature_flag_enabled():
+ if tiledbsoma._flags._new_shape_feature_flag_enabled():
  for e in snda.maxshape:
  assert e > 2_000_000_000
  else:
@@ -89,7 +89,7 @@ def test_sparse_nd_array_basics(
  with tiledbsoma.SparseNDArray.open(uri) as snda:
  assert snda.shape == arg_shape
  # This will change with current-domain support
- if tiledbsoma._new_shape_feature_flag_enabled():
+ if tiledbsoma._flags._new_shape_feature_flag_enabled():
  for e in snda.maxshape:
  assert e > 2_000_000_000
  else:
@@ -113,7 +113,7 @@ def test_sparse_nd_array_basics(
  with tiledbsoma.SparseNDArray.open(uri) as snda:
  assert snda.shape == arg_shape
 
- if tiledbsoma._new_shape_feature_flag_enabled():
+ if tiledbsoma._flags._new_shape_feature_flag_enabled():
 
  # Test resize down
  new_shape = tuple([arg_shape[i] - 50 for i in range(ndim)])
@@ -267,11 +267,11 @@ def test_dataframe_basics(tmp_path, soma_joinid_domain, index_column_names):
  has_sjid_dim = "soma_joinid" in index_column_names
  if has_sjid_dim:
  assert sdf._maybe_soma_joinid_shape == 1 + soma_joinid_domain[1]
- if not tiledbsoma._new_shape_feature_flag_enabled():
+ if not tiledbsoma._flags._new_shape_feature_flag_enabled():
  assert sdf._maybe_soma_joinid_maxshape == 1 + soma_joinid_domain[1]
  else:
  assert sdf._maybe_soma_joinid_shape is None
- if not tiledbsoma._new_shape_feature_flag_enabled():
+ if not tiledbsoma._flags._new_shape_feature_flag_enabled():
  assert sdf._maybe_soma_joinid_maxshape is None
 
  assert len(sdf.non_empty_domain()) == len(index_column_names)

apis/python/tests/test_sparse_nd_array.py

@@ -1088,7 +1088,7 @@ def test_tile_extents(tmp_path):
  ).close()
 
  with tiledb.open(uri) as A:
- if soma._new_shape_feature_flag_enabled():
+ if soma._flags._new_shape_feature_flag_enabled():
  assert A.schema.domain.dim(0).tile == 2048
  assert A.schema.domain.dim(1).tile == 2048
  else: