Repository: TileDB-Inc/TileDB-Py
Pull request #2179 (2025-03-24 20:49:21)

Expand type support for built-in array indexer
This PR expands the functionality of the square bracket TileDB Array accessor to allow indexing with Python lists, PyArrow arrays, and NumPy arrays (for point queries), in addition to the currently supported scalar, slice, or tuple indexing.

cc. @johnkerl for visibility.

---

[sc-64094]


Patch:
diff --git a/tiledb/aggregation.py b/tiledb/aggregation.py
index 1ec93684..b84360bb 100644
--- a/tiledb/aggregation.py
+++ b/tiledb/aggregation.py
@@ -55,7 +55,7 @@ class Aggregation:
         dim_ranges = index_domain_subarray(array, dom, idx)
 
         subarray = Subarray(array, array.ctx)
-        subarray.add_ranges([list([x]) for x in dim_ranges])
+        subarray.add_ranges(dim_ranges)
         q.set_subarray(subarray)
 
         cond = self.query.cond
diff --git a/tiledb/array.py b/tiledb/array.py
index c861216a..79a3b891 100644
--- a/tiledb/array.py
+++ b/tiledb/array.py
@@ -142,11 +142,42 @@ def index_domain_subarray(array, dom, idx: tuple):
         else:
             (dim_lb, dim_ub) = dim.domain
 
-        dim_slice = idx[r]
-        if not isinstance(dim_slice, slice):
-            raise IndexError("invalid index type: {!r}".format(type(dim_slice)))
+        dim_idx = idx[r]
 
-        start, stop, step = dim_slice.start, dim_slice.stop, dim_slice.step
+        if isinstance(dim_idx, np.ndarray) or isinstance(dim_idx, list):
+            if not np.issubdtype(dim_dtype, np.str_) and not np.issubdtype(
+                dim_dtype, np.bytes_
+            ):
+                # if this is a list, convert to numpy array
+                if isinstance(dim_idx, list):
+                    dim_idx = np.array(dim_idx)
+                subarray.append(
+                    (dim_idx),
+                )
+            else:
+                subarray.append([(x, x) for x in dim_idx])
+            continue
+        try:
+            import pyarrow
+
+            if isinstance(dim_idx, pyarrow.Array):
+                if not np.issubdtype(dim_dtype, np.str_) and not np.issubdtype(
+                    dim_dtype, np.bytes_
+                ):
+                    # this is zero copy by default
+                    subarray.append(dim_idx.to_numpy())
+                else:
+                    # zero copy is not supported for string types
+                    subarray.append(
+                        [(x, x) for x in dim_idx.to_numpy(zero_copy_only=False)]
+                    )
+                continue
+        except ImportError:
+            pass
+        if not isinstance(dim_idx, slice):
+            raise IndexError(f"invalid index type: {type(dim_idx)!r}")
+
+        start, stop, step = dim_idx.start, dim_idx.stop, dim_idx.step
 
         # In the case that current domain is non-empty, we need to consider it
         if (
@@ -170,7 +201,7 @@ def index_domain_subarray(array, dom, idx: tuple):
                 raise tiledb.TileDBError(
                     f"Non-string range '({start},{stop})' provided for string dimension '{dim.name}'"
                 )
-            subarray.append((start, stop))
+            subarray.append([(start, stop)])
             continue
 
         if step and array.schema.sparse:
@@ -249,16 +280,16 @@ def index_domain_subarray(array, dom, idx: tuple):
             # inclusive bounds for floating point / datetime ranges
             start = dim_dtype.type(start)
             stop = dim_dtype.type(stop)
-            subarray.append((start, stop))
+            subarray.append([(start, stop)])
         elif is_datetime:
             # need to ensure that datetime ranges are in the units of dim_dtype
             # so that add_range and output shapes work correctly
             start = start.astype(dim_dtype)
             stop = stop.astype(dim_dtype)
-            subarray.append((start, stop))
+            subarray.append([(start, stop)])
         elif np.issubdtype(type(stop), np.integer):
             # normal python indexing semantics
-            subarray.append((start, int(stop) - 1))
+            subarray.append([(start, int(stop) - 1)])
         else:
             raise IndexError(
                 "domain indexing is defined for integral and floating point values"
diff --git a/tiledb/dense_array.py b/tiledb/dense_array.py
index 7eca0ef8..51489e80 100644
--- a/tiledb/dense_array.py
+++ b/tiledb/dense_array.py
@@ -36,8 +36,9 @@ class DenseArrayImpl(Array):
     def __getitem__(self, selection):
         """Retrieve data cells for an item or region of the array.
 
-        :param tuple selection: An int index, slice or tuple of integer/slice objects,
-            specifying the selected subarray region for each dimension of the DenseArray.
+        :param selection: An int index, ``slice``, tuple, list/numpy array/pyarrow array
+            of integer/``slice`` objects, specifying the selected subarray region
+            for each dimension of the DenseArray.
         :rtype: :py:class:`numpy.ndarray` or :py:class:`collections.OrderedDict`
         :returns: If the dense array has a single attribute then a Numpy array of corresponding shape/dtype \
                 is returned for that attribute.  If the array has multiple attributes, a \
@@ -276,7 +277,7 @@ class DenseArrayImpl(Array):
         idx, drop_axes = replace_scalars_slice(self.schema.domain, idx)
         dim_ranges = index_domain_subarray(self, self.schema.domain, idx)
         subarray = Subarray(self, self.ctx)
-        subarray.add_ranges([list([x]) for x in dim_ranges])
+        subarray.add_ranges(dim_ranges)
         # Note: we included dims (coords) above to match existing semantics
         out = self._read_dense_subarray(subarray, attr_names, cond, layout, coords)
         if any(s.step for s in idx):
@@ -423,7 +424,7 @@ class DenseArrayImpl(Array):
         else:
             dim_ranges = index_domain_subarray(self, domain, idx)
             subarray = Subarray(self, self.ctx)
-            subarray.add_ranges([list([x]) for x in dim_ranges])
+            subarray.add_ranges(dim_ranges)
 
         subarray_shape = subarray.shape()
         if isinstance(val, np.ndarray):
@@ -754,7 +755,7 @@ class DenseArrayImpl(Array):
         idx = tuple(slice(None) for _ in range(domain.ndim))
         range_index = index_domain_subarray(self, domain, idx)
         subarray = Subarray(self, self.ctx)
-        subarray.add_ranges([list([x]) for x in range_index])
+        subarray.add_ranges(range_index)
         out = self._read_dense_subarray(
             subarray,
             [
diff --git a/tiledb/sparse_array.py b/tiledb/sparse_array.py
index c76150eb..55692bfc 100644
--- a/tiledb/sparse_array.py
+++ b/tiledb/sparse_array.py
@@ -251,8 +251,9 @@ class SparseArrayImpl(Array):
     def __getitem__(self, selection):
         """Retrieve nonempty cell data for an item or region of the array
 
-        :param tuple selection: An int index, slice or tuple of integer/slice objects,
-            specifying the selected subarray region for each dimension of the SparseArray.
+        :param selection: An int index, ``slice``, tuple, list/numpy array/pyarrow array
+            of integer/``slice`` objects, specifying the selected subarray region
+            for each dimension of the SparseArray.
         :rtype: :py:class:`collections.OrderedDict`
         :returns: An OrderedDict is returned with dimension and attribute names as keys. \
             Nonempty attribute values are returned as Numpy 1-d arrays.
@@ -546,7 +547,7 @@ class SparseArrayImpl(Array):
         idx, drop_axes = replace_scalars_slice(dom, idx)
         dim_ranges = index_domain_subarray(self, dom, idx)
         subarray = Subarray(self, self.ctx)
-        subarray.add_ranges([list([x]) for x in dim_ranges])
+        subarray.add_ranges(dim_ranges)
         return self._read_sparse_subarray(subarray, attr_names, cond, layout)
 
     def __repr__(self):


Test patch:
diff --git a/tiledb/tests/test_pandas_dataframe.py b/tiledb/tests/test_pandas_dataframe.py
index c6b5a091..84cf88e1 100644
--- a/tiledb/tests/test_pandas_dataframe.py
+++ b/tiledb/tests/test_pandas_dataframe.py
@@ -1277,27 +1277,46 @@ class TestPandasDataFrameRoundtrip(DiskTestCase):
         try_rt("basic3", basic3)
 
     @pytest.mark.parametrize(
-        "dim_data, attr_data, dtype, domain",
+        "dim_data, attr_data, dim_dtype, attr_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),
+            (
+                pyarrow.array([1, 2, 3, 4, 5]),
+                pyarrow.array([10, 20, 30, 40, 50]),
+                np.int64,
+                np.int64,
+                (1, 5),
+            ),
+            (
+                pyarrow.array([1.1, 20.2, 300.3, 4000.4, 50000.5]),
+                pyarrow.array(["tiledb", "is", "the", "best", "db"]),
+                np.float64,
+                "U",
+                (1.1, 50000.5),
+            ),
+            (
+                pyarrow.array([b"a", b"b", b"c", b"d", b"e"]),
+                pyarrow.array([1, 2, 3, 4, 5]),
+                "ascii",
+                np.int64,
+                None,
+            ),
         ],
     )
     def test_read_indexing_with_pyarrow_and_numpy_arrays(
-        self, dim_data, attr_data, dtype, domain
+        self, dim_data, attr_data, dim_dtype, attr_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)
+            tiledb.Dim(name="dim_a", dtype=dim_dtype, domain=domain)
             if domain
-            else tiledb.Dim(name="dim_a", dtype=dtype)
+            else tiledb.Dim(name="dim_a", dtype=dim_dtype)
         )
         schema = tiledb.ArraySchema(
             domain=tiledb.Domain(dim),
             sparse=True,
-            attrs=[tiledb.Attr(name="rand", dtype=np.int32)],
+            attrs=[tiledb.Attr(name="rand", dtype=attr_dtype)],
             allows_duplicates=True,
         )
         tiledb.Array.create(uri, schema)
@@ -1306,10 +1325,20 @@ class TestPandasDataFrameRoundtrip(DiskTestCase):
             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()}
-            )
+            if dim_dtype != "ascii":
+                expected_df = pd.DataFrame(
+                    {"dim_a": dim_data.tolist(), "rand": attr_data.tolist()}
+                )
+            else:
+                # cast to str for ascii dtype
+                expected_df = pd.DataFrame(
+                    {
+                        "dim_a": [d.as_py().decode("utf-8") for d in dim_data],
+                        "rand": [d.as_py() for d in attr_data],
+                    }
+                )
 
+            # df accessor
             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)
@@ -1322,10 +1351,29 @@ class TestPandasDataFrameRoundtrip(DiskTestCase):
             )
             assert_array_equal(arr.df[np.array(partial_dim_data)], expected_partial_df)
 
+            # square brackets accessor
+            py_attr_data = [
+                attr_data[1].as_py(),
+                attr_data[3].as_py(),
+                attr_data[4].as_py(),
+            ]
+            py_dim_data = [
+                dim_data[1].as_py(),
+                dim_data[3].as_py(),
+                dim_data[4].as_py(),
+            ]
+
             expected_dict = OrderedDict(
-                [("dim_a", dim_data.tolist()), ("rand", attr_data.tolist())]
+                {
+                    "rand": py_attr_data,
+                    "dim_a": py_dim_data,
+                }
             )
 
+            assert_dict_arrays_equal(arr[py_dim_data], expected_dict)
+            assert_dict_arrays_equal(arr[np.array(py_dim_data)], expected_dict)
+            assert_dict_arrays_equal(arr[pyarrow.array(py_dim_data)], expected_dict)
+
     def test_nullable_integers(self):
         nullable_int_dtypes = (
             pd.Int64Dtype(),