Repository: TileDB-Inc/TileDB-Py
Pull request #2157 (2025-02-10 19:57:27)

Fix `NDRectangle::set_range` overloads
This PR addresses an issue where the `NDRectangle::set_range` overloads were ignored, and only the first overload (`uint64_t`) was used after casting the arguments to `uint64_t`. The `NDRectangle::range_dtype` is now utilized to determine the appropriate type (`T`) to use.


Patch:
diff --git a/tiledb/array.py b/tiledb/array.py
index 12919d91..f6034b16 100644
--- a/tiledb/array.py
+++ b/tiledb/array.py
@@ -127,6 +127,22 @@ def index_domain_subarray(array, dom, idx: tuple):
         )
 
     subarray = list()
+
+    # In the case that current domain is non-empty, we need to consider it
+    if (
+        hasattr(array.schema, "current_domain")
+        and not array.schema.current_domain.is_empty
+    ):
+        for i in range(array.schema.domain.ndim):
+            subarray.append(
+                (
+                    array.schema.current_domain.ndrectangle.range(i)[0],
+                    array.schema.current_domain.ndrectangle.range(i)[1],
+                )
+            )
+
+        return subarray
+
     for r in range(ndim):
         # extract lower and upper bounds for domain dimension extent
         dim = dom.dim(r)
diff --git a/tiledb/libtiledb/current_domain.cc b/tiledb/libtiledb/current_domain.cc
index f01896bc..2bcef500 100644
--- a/tiledb/libtiledb/current_domain.cc
+++ b/tiledb/libtiledb/current_domain.cc
@@ -22,168 +22,125 @@ void init_current_domain(py::module& m) {
 
         .def(py::init<const Context&, const Domain&>())
 
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, uint64_t, uint64_t>(
-                &NDRectangle::set_range<uint64_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, int64_t, int64_t>(
-                &NDRectangle::set_range<int64_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, uint32_t, uint32_t>(
-                &NDRectangle::set_range<uint32_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, int32_t, int32_t>(
-                &NDRectangle::set_range<int32_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, uint16_t, uint16_t>(
-                &NDRectangle::set_range<uint16_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, int16_t, int16_t>(
-                &NDRectangle::set_range<int16_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, uint8_t, uint8_t>(
-                &NDRectangle::set_range<uint8_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, int8_t, int8_t>(
-                &NDRectangle::set_range<int8_t>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, double, double>(
-                &NDRectangle::set_range<double>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<const std::string&, float, float>(
-                &NDRectangle::set_range<float>),
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
         .def(
             "_set_range",
             [](NDRectangle& ndrect,
                const std::string& dim_name,
-               const std::string& start,
-               const std::string& end) {
-                return ndrect.set_range(dim_name, start, end);
-            },
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, uint64_t, uint64_t>(
-                &NDRectangle::set_range<uint64_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, int64_t, int64_t>(
-                &NDRectangle::set_range<int64_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, uint32_t, uint32_t>(
-                &NDRectangle::set_range<uint32_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, int32_t, int32_t>(
-                &NDRectangle::set_range<int32_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, uint16_t, uint16_t>(
-                &NDRectangle::set_range<uint16_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, int16_t, int16_t>(
-                &NDRectangle::set_range<int16_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, uint8_t, uint8_t>(
-                &NDRectangle::set_range<uint8_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, int8_t, int8_t>(
-                &NDRectangle::set_range<int8_t>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, double, double>(
-                &NDRectangle::set_range<double>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
-        .def(
-            "_set_range",
-            py::overload_cast<uint32_t, float, float>(
-                &NDRectangle::set_range<float>),
-            py::arg("dim_idx"),
-            py::arg("start"),
-            py::arg("end"))
+               py::object start,
+               py::object end) {
+                const tiledb_datatype_t n_type = ndrect.range_dtype(dim_name);
+
+                if (n_type == TILEDB_UINT64) {
+                    auto start_val = start.cast<uint64_t>();
+                    auto end_val = end.cast<uint64_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_INT64) {
+                    auto start_val = start.cast<int64_t>();
+                    auto end_val = end.cast<int64_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_UINT32) {
+                    auto start_val = start.cast<uint32_t>();
+                    auto end_val = end.cast<uint32_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_INT32) {
+                    auto start_val = start.cast<int32_t>();
+                    auto end_val = end.cast<int32_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_UINT16) {
+                    auto start_val = start.cast<uint16_t>();
+                    auto end_val = end.cast<uint16_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_INT16) {
+                    auto start_val = start.cast<int16_t>();
+                    auto end_val = end.cast<int16_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_UINT8) {
+                    auto start_val = start.cast<uint8_t>();
+                    auto end_val = end.cast<uint8_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_INT8) {
+                    auto start_val = start.cast<int8_t>();
+                    auto end_val = end.cast<int8_t>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_FLOAT64) {
+                    auto start_val = start.cast<double>();
+                    auto end_val = end.cast<double>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (n_type == TILEDB_FLOAT32) {
+                    auto start_val = start.cast<float>();
+                    auto end_val = end.cast<float>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else if (
+                    n_type == TILEDB_STRING_ASCII ||
+                    n_type == TILEDB_STRING_UTF8) {
+                    auto start_val = start.cast<std::string>();
+                    auto end_val = end.cast<std::string>();
+                    ndrect.set_range(dim_name, start_val, end_val);
+                } else {
+                    TPY_ERROR_LOC(
+                        "Unsupported type for NDRectangle's set_range");
+                }
+            })
+
         .def(
             "_set_range",
             [](NDRectangle& ndrect,
                uint32_t dim_idx,
-               const std::string& start,
-               const std::string& end) {
-                return ndrect.set_range(dim_idx, start, end);
-            },
-            py::arg("dim_name"),
-            py::arg("start"),
-            py::arg("end"))
+               py::object start,
+               py::object end) {
+                const tiledb_datatype_t n_type = ndrect.range_dtype(dim_idx);
+
+                if (n_type == TILEDB_UINT64) {
+                    auto start_val = start.cast<uint64_t>();
+                    auto end_val = end.cast<uint64_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_INT64) {
+                    auto start_val = start.cast<int64_t>();
+                    auto end_val = end.cast<int64_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_UINT32) {
+                    auto start_val = start.cast<uint32_t>();
+                    auto end_val = end.cast<uint32_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_INT32) {
+                    auto start_val = start.cast<int32_t>();
+                    auto end_val = end.cast<int32_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_UINT16) {
+                    auto start_val = start.cast<uint16_t>();
+                    auto end_val = end.cast<uint16_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_INT16) {
+                    auto start_val = start.cast<int16_t>();
+                    auto end_val = end.cast<int16_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_UINT8) {
+                    auto start_val = start.cast<uint8_t>();
+                    auto end_val = end.cast<uint8_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_INT8) {
+                    auto start_val = start.cast<int8_t>();
+                    auto end_val = end.cast<int8_t>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_FLOAT64) {
+                    auto start_val = start.cast<double>();
+                    auto end_val = end.cast<double>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (n_type == TILEDB_FLOAT32) {
+                    auto start_val = start.cast<float>();
+                    auto end_val = end.cast<float>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else if (
+                    n_type == TILEDB_STRING_ASCII ||
+                    n_type == TILEDB_STRING_UTF8) {
+                    auto start_val = start.cast<std::string>();
+                    auto end_val = end.cast<std::string>();
+                    ndrect.set_range(dim_idx, start_val, end_val);
+                } else {
+                    TPY_ERROR_LOC(
+                        "Unsupported type for NDRectangle's set_range");
+                }
+            })
 
         .def(
             "_range",
diff --git a/tiledb/multirange_indexing.py b/tiledb/multirange_indexing.py
index 2db599c6..483d1fcf 100644
--- a/tiledb/multirange_indexing.py
+++ b/tiledb/multirange_indexing.py
@@ -195,6 +195,22 @@ def label_ranges_from_selection(selection):
 
 def getitem_ranges(array: Array, idx: Any) -> Sequence[Sequence[Range]]:
     ranges: List[Sequence[Range]] = [()] * array.schema.domain.ndim
+
+    # In the case that current domain is non-empty, we need to consider it
+    if (
+        hasattr(array.schema, "current_domain")
+        and not array.schema.current_domain.is_empty
+    ):
+        for i in range(array.schema.domain.ndim):
+            ranges[i] = [
+                (
+                    array.schema.current_domain.ndrectangle.range(i)[0],
+                    array.schema.current_domain.ndrectangle.range(i)[1],
+                )
+            ]
+
+        return tuple(ranges)
+
     ned = array.nonempty_domain()
     if ned is None:
         ned = [None] * array.schema.domain.ndim


Test patch:
diff --git a/tiledb/tests/test_current_domain.py b/tiledb/tests/test_current_domain.py
index 55dff8a5..780a64f0 100644
--- a/tiledb/tests/test_current_domain.py
+++ b/tiledb/tests/test_current_domain.py
@@ -7,6 +7,8 @@ import pytest
 import tiledb
 import tiledb.libtiledb as lt
 
+from .common import DiskTestCase
+
 if not (lt.version()[0] == 2 and lt.version()[1] >= 25):
     pytest.skip(
         "CurrentDomain is only available in TileDB 2.26 and later",
@@ -14,7 +16,7 @@ if not (lt.version()[0] == 2 and lt.version()[1] >= 25):
     )
 
 
-class NDRectangleTest(unittest.TestCase):
+class NDRectangleTest(DiskTestCase):
     def test_ndrectagle_standalone_string(self):
         ctx = tiledb.Ctx()
         dom = tiledb.Domain(
@@ -57,8 +59,35 @@ class NDRectangleTest(unittest.TestCase):
         self.assertEqual(ndrect.range("x"), range_one)
         self.assertEqual(ndrect.range("y"), range_two)
 
+    @pytest.mark.parametrize(
+        "dtype_, range_",
+        (
+            (np.int32, (0, 1)),
+            (np.int64, (2, 7)),
+            (np.uint32, (1, 5)),
+            (np.uint64, (5, 9)),
+            (np.float32, (0.0, 1.0)),
+            (np.float64, (2.0, 4.0)),
+            (np.dtype(bytes), (b"abc", b"def")),
+        ),
+    )
+    def test_set_range_different_types(self, dtype_, range_):
+        domain = tiledb.Domain(
+            *[
+                tiledb.Dim(
+                    name="rows",
+                    domain=(0, 9),
+                    dtype=dtype_,
+                ),
+            ]
+        )
+
+        ctx = tiledb.Ctx()
+        ndrect = tiledb.NDRectangle(ctx, domain)
+        ndrect.set_range("rows", range_[0], range_[1])
+
 
-class CurrentDomainTest(unittest.TestCase):
+class CurrentDomainTest(DiskTestCase):
     def test_current_domain_with_ndrectangle_integer(self):
         ctx = tiledb.Ctx()
         dom = tiledb.Domain(
diff --git a/tiledb/tests/test_fixes.py b/tiledb/tests/test_fixes.py
index af44eefc..f0612c5a 100644
--- a/tiledb/tests/test_fixes.py
+++ b/tiledb/tests/test_fixes.py
@@ -1,5 +1,6 @@
 import concurrent
 import concurrent.futures
+import itertools
 import json
 import os
 import subprocess
@@ -382,6 +383,62 @@ class FixesTest(DiskTestCase):
         with tiledb.DenseArray(uri) as T:
             assert_array_equal(array_data, T)
 
+    def test_sc61914_take_current_domain_into_account_indexing(self):
+        uri = self.path("test_sc61914")
+        ctx = tiledb.Ctx()
+        dom = tiledb.Domain(
+            tiledb.Dim(name="d1", domain=(0, 99), tile=20, dtype=np.int64),
+            tiledb.Dim(name="d2", domain=(0, 99), tile=20, dtype=np.int64),
+        )
+        att = tiledb.Attr(name="a", dtype=np.int64)
+        schema = tiledb.ArraySchema(sparse=False, ctx=ctx, domain=dom, attrs=(att,))
+
+        tiledb.Array.create(uri, schema)
+
+        data = np.arange(0, 10000).reshape(100, 100)
+
+        with tiledb.open(uri, "w") as A:
+            A[:] = data
+
+        with tiledb.DenseArray(uri, mode="r") as A:
+            ndrect = tiledb.NDRectangle(ctx, dom)
+            range_one = (10, 20)
+            range_two = (30, 35)
+            ndrect.set_range(0, range_one[0], range_one[1])
+            ndrect.set_range(1, range_two[0], range_two[1])
+
+            current_domain = tiledb.CurrentDomain(ctx)
+            current_domain.set_ndrectangle(ndrect)
+            A.schema.set_current_domain(current_domain)
+
+            # Define the expected results
+            d1_values = range(10, 21)
+            d2_values = range(30, 36)
+            data = [
+                (d1, d2, d1 * 100 + d2)
+                for d1, d2 in itertools.product(d1_values, d2_values)
+            ]
+            expected_df = pd.DataFrame(data, columns=["d1", "d2", "a"])
+
+            expected_array = np.array(
+                [
+                    [1030, 1031, 1032, 1033, 1034, 1035],
+                    [1130, 1131, 1132, 1133, 1134, 1135],
+                    [1230, 1231, 1232, 1233, 1234, 1235],
+                    [1330, 1331, 1332, 1333, 1334, 1335],
+                    [1430, 1431, 1432, 1433, 1434, 1435],
+                    [1530, 1531, 1532, 1533, 1534, 1535],
+                    [1630, 1631, 1632, 1633, 1634, 1635],
+                    [1730, 1731, 1732, 1733, 1734, 1735],
+                    [1830, 1831, 1832, 1833, 1834, 1835],
+                    [1930, 1931, 1932, 1933, 1934, 1935],
+                    [2030, 2031, 2032, 2033, 2034, 2035],
+                ]
+            )
+
+            assert_array_equal(A, expected_array)
+            assert_array_equal(A.df[:], expected_df)
+
 
 class SOMA919Test(DiskTestCase):
     """