Fix .df roundtrip for str/bytes dimensions and tighten StringDtype handling

- _update_df_from_meta: casting to the zero-width '<U0'/'|S0' dtypes stored
  in the metadata breaks on pandas 3 (str columns turn into object, bytes
  into fixed-width '|S<n>' that cannot be used as an index). Replace it with
  an explicit bytes<->str conversion restoring the type the dimension was
  written with, and skip no-op casts so that dropping `copy=False` (removed
  in pandas 3) does not introduce extra copies on pandas 2.
- ColumnInfo: distinguish the two StringDtype variants by `na_value` (per the
  pandas migration guide) and reuse from_dtype in from_values instead of
  duplicating the mapping.
- _df_to_np_arrays: restore unconditional nullmap creation; the data scan
  guard was unnecessary and the explicit vectorized nullmap is cheaper than
  the per-element fallback in the write path.
- Add a str/bytes dimension .df roundtrip regression test covering both the
  arrow and non-arrow read paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: kounelisagis

tiledb/dataframe_.py

@@ -2,7 +2,7 @@
 import json
 import os
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any, List, Optional, Union
 
 import numpy as np
@@ -162,15 +162,13 @@ def from_values(cls, array_like, varlen_types=()):
                     f"{inferred_dtype} inferred dtype not supported (column {array_like.name})"
                 )
         elif hasattr(array_like, "dtype") and isinstance(array_like.dtype, StringDtype):
-            # Explicit pd.StringDtype() (name="string") is always nullable;
-            # auto-inferred str (name="str") depends on data
-            explicit = array_like.dtype.name == "string"
-            return cls(
-                np.dtype(np.str_),
-                repr="string" if explicit else None,
-                var=True,
-                nullable=explicit or bool(array_like.isna().any()),
-            )
+            info = cls.from_dtype(array_like.dtype, array_like.name)
+            # the NaN-backed "str" dtype holds missing values without being a
+            # nullable dtype; write it as a nullable attribute only when
+            # missing values are actually present
+            if not info.nullable and array_like.isna().any():
+                info = replace(info, nullable=True)
+            return info
         elif hasattr(array_like, "dtype") and isinstance(
             array_like.dtype, CategoricalDtype
         ):
@@ -200,6 +198,7 @@ def from_categorical(cls, cat, dtype, column_name):
 
     @classmethod
     def from_dtype(cls, dtype, column_name, varlen_types=()):
+        from pandas import NA, StringDtype
         from pandas.api import types as pd_types
 
         if isinstance(dtype, str) and dtype == "ascii":
@@ -211,13 +210,17 @@ def from_dtype(cls, dtype, column_name, varlen_types=()):
         dtype = pd_types.pandas_dtype(dtype)
         # Note: be careful if you rearrange the order of the following checks
 
-        # pandas StringDtype (auto-inferred 'str' and explicit 'string')
-        from pandas import StringDtype
-
+        # pandas string types: the NA-backed "string" dtype maps to a nullable
+        # attribute, while the NaN-backed "str" dtype (the default for strings
+        # since pandas 3) behaves like plain str/object columns
         if isinstance(dtype, StringDtype):
-            repr_val = "string" if dtype.name == "string" else None
-            nullable = dtype.name == "string"
-            return cls(np.dtype(np.str_), repr=repr_val, var=True, nullable=nullable)
+            nullable = dtype.na_value is NA
+            return cls(
+                np.dtype(np.str_),
+                repr="string" if nullable else None,
+                var=True,
+                nullable=nullable,
+            )
 
         # extension types
         if pd_types.is_extension_array_dtype(dtype):
@@ -520,8 +523,8 @@ def _df_to_np_arrays(df, column_infos, fillna):
         if not column_info.var:
             to_numpy_kwargs.update(dtype=column_info.dtype)
 
-        if column_info.nullable and column.isna().any():
-            # Only create nullmap if data actually has nulls
+        if column_info.nullable:
+            # use default 0/empty for the dtype
             to_numpy_kwargs.update(na_value=column_info.dtype.type())
             nullmaps[name] = (~column.isna()).to_numpy(dtype=np.uint8)
 

tiledb/multirange_indexing.py

@@ -889,17 +889,22 @@ def _update_df_from_meta(
             if name in df:
                 col_dtypes[name] = dtype
 
-    if col_dtypes:
-        # '<U0' is stored in __pandas_index_dims metadata for var-length string
-        # dimensions (str(np.dtype(np.str_)) == '<U0>'). Applying astype('<U0')
-        # was a no-op on pandas 2 but on pandas 3 it forces StringDtype back to
-        # object, breaking the roundtrip. The string data already has the correct
-        # dtype from pandas' own inference, so we skip it here.
-        col_dtypes = {
-            name: dtype for name, dtype in col_dtypes.items() if dtype != "<U0"
-        }
-        if col_dtypes:
-            df = df.astype(col_dtypes)
+    for name, dtype in col_dtypes.items():
+        # str/bytes dimensions are always written as ASCII bytes and, depending
+        # on the read path, come back as either bytes or str; restore the type
+        # they were written with. astype with the zero-width '<U0'/'|S0' dtypes
+        # stored in the metadata used to do this on pandas < 3, but on
+        # pandas >= 3 it produces object or fixed-width bytes columns instead.
+        if dtype == "<U0":
+            if len(df) and isinstance(df[name].iat[0], bytes):
+                df[name] = df[name].str.decode("utf-8")
+        elif dtype == "|S0":
+            if len(df) and not isinstance(df[name].iat[0], bytes):
+                df[name] = df[name].str.encode("utf-8")
+        # skip columns that already have their target dtype: casting them
+        # would needlessly copy the data on pandas < 3
+        elif str(df[name].dtype) != dtype:
+            df[name] = df[name].astype(dtype)
 
     if index_col:
         if index_col is not True:

tiledb/tests/test_pandas_dataframe.py

@@ -626,6 +626,19 @@ def test_dataframe_index_to_sparse_dims(self):
                 res_df = pd.DataFrame(res, index=index)
                 tm.assert_frame_equal(new_df, res_df, check_like=True)
 
+    @pytest.mark.parametrize("index_data", [["aa", "b"], [b"aa", b"b"]])
+    def test_dataframe_str_dim_df_roundtrip(self, index_data):
+        # both str and bytes dimensions are stored as ASCII bytes; on read the
+        # type they were written with must be restored, on both read paths
+        uri = self.path("df_str_dim_df_roundtrip")
+
+        df = pd.DataFrame({"data": [1.0, 2.0]}, index=pd.Index(index_data, name="idx"))
+        tiledb.from_pandas(uri, df, sparse=True)
+
+        with tiledb.open(uri) as A:
+            tm.assert_frame_equal(A.df[:], df)
+            tm.assert_frame_equal(A.query(use_arrow=False).df[:], df)
+
     def test_dataframe_set_index_dims(self):
         uri = self.path("df_set_index_dims")