Repository: TileDB-Inc/TileDB-Py
Pull request #2168 (2025-03-04 09:12:02)

Improve the logic for using `Arrow` when it is not explicitly requested
This PR expands the cases where a Pandas DataFrame can be retrieved from a TileDB `Array` or `Query`. This improvement applies _only_ when the `use_arrow` argument is not set.  

### Logic before this PR:  
If the `pyarrow` package was available, `use_arrow` would automatically be set to `True`. This caused issues in cases involving multi-value attributes. The problem was that the user never explicitly requested Arrow, yet encountered an Arrow-related issue. Additionally, the expected `TileDBError` for multi-value attributes was never triggered because the loop `for attr in map(array.attr, query.attrs or ())` did not work as intended. The reason is that `array.attr` is a method that takes an attribute index or name as a parameter, rather than returning attributes directly.  

### Logic after this PR:  
A new method, `attr_names`, has been introduced in the `ArraySchema` and `Array` classes. Utilizing this method, we can now determine whether Arrow can be used for specific attributes. If Arrow is supported, we use it. Otherwise, we don't.  

---

[sc-41582]


Patch:
diff --git a/tiledb/array.py b/tiledb/array.py
index f7abdb87..e7221004 100644
--- a/tiledb/array.py
+++ b/tiledb/array.py
@@ -569,6 +569,11 @@ class Array:
         :raises TypeError: invalid key type"""
         return self.schema.attr(key)
 
+    @property
+    def attr_names(self):
+        """Returns a list of attribute names"""
+        return self.schema.attr_names
+
     def dim(self, dim_id):
         """Returns a :py:class:`Dim` instance given a dim index or name
 
diff --git a/tiledb/array_schema.py b/tiledb/array_schema.py
index d631df9e..f8d9c429 100644
--- a/tiledb/array_schema.py
+++ b/tiledb/array_schema.py
@@ -1,7 +1,7 @@
 import io
 import numbers
 import warnings
-from typing import Sequence, Tuple, Union
+from typing import List, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -360,6 +360,14 @@ class ArraySchema(CtxMixin, lt.ArraySchema):
             "or an integer index, not {0!r}".format(type(key))
         )
 
+    @property
+    def attr_names(self) -> List[str]:
+        """Returns a list of attribute names
+
+        :rtype: list
+        """
+        return [self.attr(i).name for i in range(self.nattr)]
+
     def dim_label(self, name: str) -> DimLabel:
         """Returns a TileDB DimensionLabel given the label name
 
diff --git a/tiledb/multirange_indexing.py b/tiledb/multirange_indexing.py
index 47181ebb..05969136 100644
--- a/tiledb/multirange_indexing.py
+++ b/tiledb/multirange_indexing.py
@@ -239,15 +239,6 @@ def getitem_ranges_with_labels(
     return dim_ranges, label_ranges
 
 
-def has_numeric_var_length_attribute(array: Array) -> bool:
-    schema = array.schema
-    for i in range(schema.nattr):
-        attr = schema.attr(i)
-        if attr.isvar and np.issubdtype(attr.dtype, np.number):
-            return True
-    return False
-
-
 class _BaseIndexer(ABC):
     """
     Implements multi-range indexing.
@@ -448,25 +439,31 @@ class DataFrameIndexer(_BaseIndexer):
         # we need to use a Query in order to get coords for a dense array
         if not query:
             query = QueryProxy(array, has_coords=True)
-        use_arrow = (
-            bool(importlib.util.find_spec("pyarrow"))
-            if use_arrow is None
-            else use_arrow
-        )
 
-        # TODO: currently there is lack of support for Arrow list types. This prevents
+        # Currently there is lack of support for Arrow list types. This prevents
         # multi-value attributes, asides from strings, from being queried properly.
-        # Until list attributes are supported in core, error with a clear message.
-        if use_arrow and any(
+        attr_names_to_check = query.attrs or array.attr_names
+
+        arrow_can_be_used = not any(
             (attr.isvar or len(attr.dtype) > 1)
             and attr.dtype not in (np.str_, np.bytes_)
-            for attr in map(array.attr, query.attrs or ())
-        ):
+            for attr in (array.attr(name) for name in attr_names_to_check)
+        )
+
+        # Until list attributes are supported in core, error with a clear message.
+        if use_arrow and not arrow_can_be_used:
             raise TileDBError(
                 "Multi-value attributes are not currently supported when use_arrow=True. "
                 "This includes all variable-length attributes and fixed-length "
                 "attributes with more than one value. Use `query(use_arrow=False)`."
             )
+
+        use_arrow = (
+            bool(importlib.util.find_spec("pyarrow"))
+            if use_arrow is None and arrow_can_be_used
+            else use_arrow
+        )
+
         super().__init__(array, query, use_arrow, preload_metadata=True)
 
     def _run_query(self) -> Union["pandas.DataFrame", "pyarrow.Table"]:
@@ -479,10 +476,6 @@ class DataFrameIndexer(_BaseIndexer):
         if self.pyquery is None:
             df = pandas.DataFrame(self._empty_results)
         elif self.use_arrow:
-            if has_numeric_var_length_attribute(self.array):
-                raise TileDBError(
-                    "Variable-length numeric attributes are not supported in Arrow. Please use Array.query(use_arrow=False) instead."
-                )
             with timing("buffer_conversion_time"):
                 table = self.pyquery._buffers_to_pa_table()
 


Test patch:
diff --git a/.github/workflows/daily-tests.yml b/.github/workflows/daily-tests.yml
index 6d6eae91..f9a86d25 100644
--- a/.github/workflows/daily-tests.yml
+++ b/.github/workflows/daily-tests.yml
@@ -15,7 +15,7 @@ jobs:
   ci2:
     uses: ./.github/workflows/daily-test-build.yml
     with:
-      libtiledb_version: '2.26.3'
+      libtiledb_version: '2.26.2'
 
   ci3:
     uses: ./.github/workflows/daily-test-build-numpy.yml
@@ -25,4 +25,4 @@ jobs:
   ci4:
     uses: ./.github/workflows/daily-test-build-numpy.yml
     with:
-      libtiledb_version: '2.26.3'
+      libtiledb_version: '2.26.2'
diff --git a/tiledb/tests/test_attribute.py b/tiledb/tests/test_attribute.py
index fd3f62a6..ff67b17c 100644
--- a/tiledb/tests/test_attribute.py
+++ b/tiledb/tests/test_attribute.py
@@ -345,17 +345,29 @@ class AttributeTest(DiskTestCase):
             assert_array_equal(A[:]["value"][2], data[2])
 
             if has_pandas():
-                with pytest.raises(tiledb.TileDBError) as exc:
-                    A.df[:]
-                assert (
-                    "Variable-length numeric attributes are not supported in Arrow"
-                    in str(exc.value)
+                import pandas as pd
+
+                tm = pd._testing
+
+                expected_df = pd.DataFrame(
+                    {
+                        "sample": [b"sample1", b"sample2", b"sample3"],
+                        "value": [
+                            np.array([1, 2], dtype="uint32"),
+                            np.array([3], dtype="uint32"),
+                            np.array([4], dtype="uint32"),
+                        ],
+                    }
                 )
 
+                tm.assert_frame_equal(A.df[:], expected_df)
+
                 with pytest.raises(tiledb.TileDBError) as exc:
                     A.query(use_arrow=True).df[:]
                 assert (
-                    "Variable-length numeric attributes are not supported in Arrow"
+                    "Multi-value attributes are not currently supported when use_arrow=True. "
+                    "This includes all variable-length attributes and fixed-length "
+                    "attributes with more than one value. Use `query(use_arrow=False)`."
                     in str(exc.value)
                 )
 
@@ -392,16 +404,28 @@ class AttributeTest(DiskTestCase):
             assert_array_equal(A[:]["value"][2], data[2])
 
             if has_pandas():
-                with pytest.raises(tiledb.TileDBError) as exc:
-                    A.df[:]
-                assert (
-                    "Variable-length numeric attributes are not supported in Arrow"
-                    in str(exc.value)
+                import pandas as pd
+
+                tm = pd._testing
+
+                expected_df = pd.DataFrame(
+                    {
+                        "sample": [b"sample1", b"sample2", b"sample3"],
+                        "value": [
+                            np.array([1.0, 2.0], dtype="float32"),
+                            np.array([np.nan], dtype="float32"),
+                            np.array([4.0], dtype="float32"),
+                        ],
+                    }
                 )
 
+                tm.assert_frame_equal(A.df[:], expected_df)
+
                 with pytest.raises(tiledb.TileDBError) as exc:
                     A.query(use_arrow=True).df[:]
                 assert (
-                    "Variable-length numeric attributes are not supported in Arrow"
+                    "Multi-value attributes are not currently supported when use_arrow=True. "
+                    "This includes all variable-length attributes and fixed-length "
+                    "attributes with more than one value. Use `query(use_arrow=False)`."
                     in str(exc.value)
                 )
diff --git a/tiledb/tests/test_multi_index.py b/tiledb/tests/test_multi_index.py
index a6ad6232..6f8c7ef1 100644
--- a/tiledb/tests/test_multi_index.py
+++ b/tiledb/tests/test_multi_index.py
@@ -755,8 +755,8 @@ class TestMultiRange(DiskTestCase):
             result = A.query(attrs=["111"])[0]
             assert_array_equal(result["111"], data_111)
 
-            with self.assertRaises(tiledb.TileDBError):
-                result = A.query(attrs=["111"]).df[0]
+            result = A.query(attrs=["111"]).df[0]
+            assert_array_equal(result["111"], data_111)
 
             result = A.query(attrs=["111"], use_arrow=False)
             assert_array_equal(result.df[0]["111"], data_111)
@@ -792,8 +792,10 @@ class TestMultiRange(DiskTestCase):
             assert_array_equal(result[1]["1s"][0], data[1])
             assert_array_equal(result[2]["1s"][0], data[2])
 
-            with self.assertRaises(tiledb.TileDBError):
-                result = A.query(attrs=["1s"]).df[0]
+            result = A.query(attrs=["1s"])
+            assert_array_equal(result.df[0]["1s"][0], data[0])
+            assert_array_equal(result.df[1]["1s"][0], data[1])
+            assert_array_equal(result.df[2]["1s"][0], data[2])
 
             result = A.query(attrs=["1s"], use_arrow=False)
             assert_array_equal(result.df[0]["1s"][0], data[0])
diff --git a/tiledb/tests/test_pandas_dataframe.py b/tiledb/tests/test_pandas_dataframe.py
index f3d0d0e5..c9076785 100644
--- a/tiledb/tests/test_pandas_dataframe.py
+++ b/tiledb/tests/test_pandas_dataframe.py
@@ -1342,13 +1342,11 @@ class TestPandasDataFrameRoundtrip(DiskTestCase):
         )
 
         with tiledb.open(uri) as A:
-            # TODO: update the test when we support Arrow lists
-            with pytest.raises(tiledb.TileDBError) as exc:
-                A.df[:]
-            assert (
-                "Variable-length numeric attributes are not supported in Arrow"
-                in str(exc.value)
-            )
+            df1 = A.df[:]
+            tm.assert_frame_equal(df, df1, check_dtype=False)
+            for array1, array2 in zip(df["data"].values, df1["data"].values):
+                self.assertEqual(array1.dtype, array2.dtype)
+                np.testing.assert_array_equal(array1, array2)
 
             df2 = A.query(use_arrow=False).df[:]
             tm.assert_frame_equal(df, df2, check_dtype=False)