Repository: TileDB-Inc/TileDB-Py
Pull request #2170 (2025-03-05 14:27:09)

Expand `.df` Array/Query accessor to allow indexing with NumPy and PyArrow arrays
This PR expands the functionality of the `.df` TileDB Array/Query accessor to also allow indexing with both PyArrow and NumPy arrays, instead of only scalar, slice, or a list of scalars or slice objects, as it currently does.

cc. @johnkerl for visibility.

---

[sc-63089]


Patch:
diff --git a/tiledb/array.py b/tiledb/array.py
index e7221004..fe773aad 100644
--- a/tiledb/array.py
+++ b/tiledb/array.py
@@ -1039,9 +1039,9 @@ class Array:
         """Retrieve data cells with multi-range, domain-inclusive indexing. Returns
         the cross-product of the ranges.
 
-        :param list selection: Per dimension, a scalar, ``slice``, or list of scalars
-            or ``slice`` objects. Scalars and ``slice`` components should match the
-            type of the underlying Dimension.
+        :param list selection: Per dimension, a scalar, ``slice``,
+            or a list/numpy array/pyarrow array of scalars or ``slice`` objects.
+            Scalars and ``slice`` components should match the type of the underlying Dimension.
         :returns: dict of {'attribute': result}. Coords are included by default for
             Sparse arrays only (use `Array.query(coords=<>)` to select).
         :raises IndexError: invalid or unsupported index selection
@@ -1093,9 +1093,9 @@ class Array:
         """Retrieve data cells as a Pandas dataframe, with multi-range,
         domain-inclusive indexing using ``multi_index``.
 
-        :param list selection: Per dimension, a scalar, ``slice``, or list of scalars
-            or ``slice`` objects. Scalars and ``slice`` components should match the
-            type of the underlying Dimension.
+        :param list selection: Per dimension, a scalar, ``slice``,
+            or a list/numpy array/pyarrow array of scalars or ``slice`` objects.
+            Scalars and ``slice`` components should match the type of the underlying Dimension.
         :returns: dict of {'attribute': result}. Coords are included by default for
             Sparse arrays only (use `Array.query(coords=<>)` to select).
         :raises IndexError: invalid or unsupported index selection
diff --git a/tiledb/multirange_indexing.py b/tiledb/multirange_indexing.py
index 05969136..93c33d92 100644
--- a/tiledb/multirange_indexing.py
+++ b/tiledb/multirange_indexing.py
@@ -40,8 +40,14 @@ if TYPE_CHECKING:
     # We don't want to import these eagerly since importing Pandas in particular
     # can add around half a second of import time even if we never use it.
     import pandas
+
+
+try:
     import pyarrow
 
+    has_pyarrow = True
+except ImportError:
+    has_pyarrow = False
 
 current_timer: ContextVar[str] = ContextVar("timer_scope")
 
@@ -112,11 +118,15 @@ def to_scalar(obj: Any) -> Scalar:
         return cast(Scalar, obj)
     if isinstance(obj, np.ndarray) and obj.ndim == 0:
         return cast(Scalar, obj[()])
+    if has_pyarrow and isinstance(obj, pyarrow.Array):
+        return to_scalar(obj.to_numpy()[()])
+    if has_pyarrow and isinstance(obj, pyarrow.Scalar):
+        return cast(Scalar, obj.as_py())
     raise ValueError(f"Cannot convert {type(obj)} to scalar")
 
 
 def iter_ranges(
-    sel: Union[Scalar, slice, Range, List[Scalar]],
+    sel: Union[Scalar, slice, Range, List[Scalar], np.ndarray, "pyarrow.Array"],
     sparse: bool,
     nonempty_domain: Optional[Range] = None,
 ) -> Iterator[Range]:
@@ -145,7 +155,9 @@ def iter_ranges(
         assert len(sel) == 2
         yield to_scalar(sel[0]), to_scalar(sel[1])
 
-    elif isinstance(sel, list):
+    elif isinstance(sel, (list, np.ndarray)) or (
+        has_pyarrow and isinstance(sel, pyarrow.Array)
+    ):
         for scalar in map(to_scalar, sel):
             yield scalar, scalar
 
@@ -178,8 +190,6 @@ def iter_label_range(sel: Union[Scalar, slice, Range, List[Scalar]]):
 
 def dim_ranges_from_selection(selection, nonempty_domain, is_sparse):
     # don't try to index nonempty_domain if None
-    if isinstance(selection, np.ndarray):
-        return selection
     selection = selection if isinstance(selection, list) else [selection]
     return tuple(
         rng for sel in selection for rng in iter_ranges(sel, is_sparse, nonempty_domain)
diff --git a/tiledb/query_condition.py b/tiledb/query_condition.py
index 0902906c..8606460b 100644
--- a/tiledb/query_condition.py
+++ b/tiledb/query_condition.py
@@ -212,7 +212,33 @@ class QueryConditionTree(ast.NodeVisitor):
                     "`in` operator syntax must be written as `variable in ['l', 'i', 's', 't']`"
                 )
 
-            variable = node.left.id
+            # For 'my_string in ["red", "yellow"]': node.left is ast.Name
+            # For 'attr(my.string) in ["red", "yellow"]': node.left is ast.Call
+            if isinstance(node.left, ast.Call):
+                if node.left.func.id != "attr":
+                    raise TileDBError(
+                        f"query condition left-hand side function call must be 'attr'; got '{node.left.func.id}'"
+                    )
+
+                if len(node.left.args) != 1:
+                    raise TileDBError(
+                        f"query condition left-hand side 'attr' function call must have one argument; got '{len(node.left.args)}'"
+                    )
+
+                arg = node.left.args[0]
+                if not isinstance(arg, ast.Constant):
+                    raise TileDBError(
+                        "query condition left-hand side 'attr' argument must be a constant"
+                    )
+                variable = arg.value
+
+            elif isinstance(node.left, ast.Name):
+                variable = node.left.id
+            else:
+                raise TileDBError(
+                    f"cannot handle query condition left-hand side of type '{type(node.left)}'"
+                )
+
             values = [self.get_value_from_node(val) for val in self.visit(rhs)]
             if len(values) == 0:
                 raise TileDBError(
@@ -230,7 +256,7 @@ class QueryConditionTree(ast.NodeVisitor):
 
             dtype = "string" if dt.kind in "SUa" else dt.name
             op = qc.TILEDB_IN if isinstance(operator, ast.In) else qc.TILEDB_NOT_IN
-            result = self.create_pyqc(dtype)(self.ctx, node.left.id, values, op)
+            result = self.create_pyqc(dtype)(self.ctx, variable, values, op)
 
         else:
             raise TileDBError(f"unrecognized operator in <<{ast.dump(node)}>>")


Test patch:
diff --git a/tiledb/tests/test_pandas_dataframe.py b/tiledb/tests/test_pandas_dataframe.py
index c9076785..c6b5a091 100644
--- a/tiledb/tests/test_pandas_dataframe.py
+++ b/tiledb/tests/test_pandas_dataframe.py
@@ -1276,6 +1276,56 @@ class TestPandasDataFrameRoundtrip(DiskTestCase):
         basic3 = make_dataframe_basic3()
         try_rt("basic3", basic3)
 
+    @pytest.mark.parametrize(
+        "dim_data, attr_data, dtype, domain",
+        [
+            (pyarrow.array([1, 2, 3]), pyarrow.array([1, 2, 3]), np.int64, (1, 3)),
+            (pyarrow.array(["a", "b", "c"]), pyarrow.array([1, 2, 3]), "ascii", None),
+        ],
+    )
+    def test_read_indexing_with_pyarrow_and_numpy_arrays(
+        self, dim_data, attr_data, dtype, domain
+    ):
+        # This test is to ensure that .df can be indexed with both PyArrow and NumPy arrays.
+        uri = self.path("read_indexing_with_pyarrow_and_numpy_arrays")
+
+        dim = (
+            tiledb.Dim(name="dim_a", dtype=dtype, domain=domain)
+            if domain
+            else tiledb.Dim(name="dim_a", dtype=dtype)
+        )
+        schema = tiledb.ArraySchema(
+            domain=tiledb.Domain(dim),
+            sparse=True,
+            attrs=[tiledb.Attr(name="rand", dtype=np.int32)],
+            allows_duplicates=True,
+        )
+        tiledb.Array.create(uri, schema)
+
+        with tiledb.open(uri, "w") as arr:
+            arr[dim_data] = attr_data
+
+        with tiledb.open(uri, "r") as arr:
+            expected_df = pd.DataFrame(
+                {"dim_a": dim_data.tolist(), "rand": attr_data.tolist()}
+            )
+
+            assert_array_equal(arr.df[:], expected_df)
+            assert_array_equal(arr.df[pyarrow.array(dim_data)], expected_df)
+            assert_array_equal(arr.df[np.array(dim_data)], expected_df)
+
+            partial_dim_data = dim_data[:2]
+            expected_partial_df = expected_df.iloc[:2]
+
+            assert_array_equal(
+                arr.df[pyarrow.array(partial_dim_data)], expected_partial_df
+            )
+            assert_array_equal(arr.df[np.array(partial_dim_data)], expected_partial_df)
+
+            expected_dict = OrderedDict(
+                [("dim_a", dim_data.tolist()), ("rand", attr_data.tolist())]
+            )
+
     def test_nullable_integers(self):
         nullable_int_dtypes = (
             pd.Int64Dtype(),
diff --git a/tiledb/tests/test_query_condition.py b/tiledb/tests/test_query_condition.py
index e41d319b..1f69539d 100644
--- a/tiledb/tests/test_query_condition.py
+++ b/tiledb/tests/test_query_condition.py
@@ -765,11 +765,16 @@ class QueryConditionTest(DiskTestCase):
         attrs = [
             tiledb.Attr(name="attr.one", dtype=np.uint32),
             tiledb.Attr(name="attr.two", dtype=np.uint32),
+            tiledb.Attr(name="at.tr.thr.ee", dtype=np.uint32),
         ]
         schema = tiledb.ArraySchema(domain=dom, attrs=attrs, sparse=True)
         tiledb.Array.create(path, schema)
         with tiledb.open(path, "w") as A:
-            A[np.arange(11)] = {"attr.one": np.arange(11), "attr.two": np.arange(11)}
+            A[np.arange(11)] = {
+                "attr.one": np.arange(11),
+                "attr.two": np.arange(11),
+                "at.tr.thr.ee": np.arange(11),
+            }
         with tiledb.open(path, "r") as A:
             with pytest.raises(tiledb.TileDBError) as exc_info:
                 A.query(cond="attr.one < 6")[:]
@@ -783,12 +788,26 @@ class QueryConditionTest(DiskTestCase):
                 "TileDBError: Unhandled dot operator in Attribute(value=Name(id='attr', ctx=Load()), attr='two', ctx=Load()) -- if your attribute name has a dot in it, e.g. `orig.ident`, please wrap it with `attr(\"...\")`, e.g. `attr(\"orig.ident\")`"
                 in str(exc_info.value)
             )
+            with pytest.raises(tiledb.TileDBError) as exc_info:
+                A.query(cond="attr.one in [1, 2, 3]")[:]
+            assert (
+                "TileDBError: cannot handle query condition left-hand side of type"
+                in str(exc_info.value)
+            )
 
             # now test with the correct syntax
             result = A.query(cond='attr("attr.one") < 6')[:]
             assert_array_equal(result["attr.one"], A[:6]["attr.one"])
+            result = A.query(cond='attr("at.tr.thr.ee") >= 6')[:]
+            assert_array_equal(result["at.tr.thr.ee"], A[6:]["at.tr.thr.ee"])
             result = A.query(cond='attr("attr.two") >= 6')[:]
             assert_array_equal(result["attr.two"], A[6:]["attr.two"])
+            result = A.query(cond='attr("at.tr.thr.ee") >= 6')[:]
+            assert_array_equal(result["at.tr.thr.ee"], A[6:]["at.tr.thr.ee"])
+            result = A.query(cond='attr("attr.one") in [1, 2, 3]')[:]
+            assert_array_equal(result["attr.one"], A[1:4]["attr.one"])
+            result = A.query(cond='attr("at.tr.thr.ee") in [1, 2, 3]')[:]
+            assert_array_equal(result["at.tr.thr.ee"], A[1:4]["at.tr.thr.ee"])
 
     @pytest.mark.skipif(not has_pandas(), reason="pandas>=1.0,<3.0 not installed")
     def test_do_not_return_attrs(self):