Change DenseTensorSchema.key_range to return a ConstrainedPartitionsIntRange (#193)

Author: gsakkis

tests/readers/test_tensor_schema.py

@@ -15,8 +15,7 @@ def dense_uri(tmp_path_factory):
         sparse=False,
         domain=tiledb.Domain(
             tiledb.Dim(name="d0", domain=(0, 9999), dtype=np.int32, tile=123),
-            tiledb.Dim(name="d1", domain=(-2, 2), dtype=np.int32, tile=2),
-            tiledb.Dim(name="d2", domain=(1, 2), dtype=np.int32, tile=1),
+            tiledb.Dim(name="d1", domain=(-20, 19), dtype=np.int32, tile=4),
         ),
         attrs=[
             tiledb.Attr(name="af8", dtype=np.float64),
@@ -90,21 +89,31 @@ def parametrize_fields(*fields, num=3):
 @pytest.mark.parametrize(
     "key_dim,memory_budget,dim_selectors",
     [
-        ("d0", 16_000, {}),
-        ("d0", 32_000, {}),
-        ("d0", 64_000, {}),
-        ("d1", 500_000, {}),
-        ("d1", 600_000, {}),
+        ("d0", 160_000, {}),
+        ("d0", 160_000, {"d0": slice(1000, 9000)}),
+        ("d0", 160_000, {"d0": slice(None, 9000)}),
+        ("d0", 160_000, {"d0": slice(1000, None)}),
+        ("d0", 160_000, {"d1": slice(-10, 10)}),
+        ("d0", 160_000, {"d1": slice(-10, None)}),
+        ("d0", 160_000, {"d1": slice(None, 10)}),
+        ("d0", 160_000, {"d1": list(range(-10, 10, 3))}),
+        ("d0", 160_000, {"d0": slice(1000, 9000), "d1": slice(-10, 10)}),
+        ("d0", 160_000, {"d0": slice(None, 9000), "d1": slice(-10, None)}),
+        ("d0", 160_000, {"d0": slice(1000, None), "d1": slice(None, 10)}),
         ("d1", 700_000, {}),
-        ("d0", 16_000, {"d1": slice(0, 2)}),
-        ("d0", 32_000, {"d1": slice(None, 0)}),
-        ("d0", 64_000, {"d1": slice(-1, None), "d2": [1]}),
-        ("d1", 500_000, {"d0": [1, 2, 3]}),
-        ("d1", 600_000, {"d0": [1, 100, 143, 976], "d2": slice(2, 2)}),
-        ("d1", 700_000, {"d0": [1, 100, 143, 1093, 1094]}),
+        ("d1", 700_000, {"d1": slice(-10, 10)}),
+        ("d1", 700_000, {"d1": slice(None, 10)}),
+        ("d1", 700_000, {"d1": slice(-10, None)}),
+        ("d1", 700_000, {"d0": slice(1000, 9000)}),
+        ("d1", 700_000, {"d0": slice(None, 9000)}),
+        ("d1", 700_000, {"d0": slice(1000, None)}),
+        ("d1", 700_000, {"d0": list(range(100, 5000, 3))}),
+        ("d1", 700_000, {"d1": slice(-10, 10), "d0": slice(1000, 9000)}),
+        ("d1", 700_000, {"d1": slice(None, 10), "d0": slice(None, 9000)}),
+        ("d1", 700_000, {"d1": slice(-10, None), "d0": slice(1000, None)}),
     ],
 )
-@parametrize_fields("d0", "d1", "d2", "af8", "af4", "au1")
+@parametrize_fields("d0", "d1", "af8", "af4", "au1")
 def test_max_partition_weight_dense(
     dense_uri, fields, key_dim, memory_budget, dim_selectors
 ):

tests/readers/utils.py

@@ -153,7 +153,7 @@ def dim(self):
         return tiledb.Dim(
             name=self.name,
             domain=(self.min_value, self(self.size - 1)),
-            tile=np.random.randint(1, self.size + 1),
+            tile=np.random.randint(1, self.size),
             dtype=self.dtype,
         )
 

tiledb/ml/readers/_tensor_schema/dense.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from .base import TensorSchema
-from .ranges import InclusiveRange, IntRange
+from .ranges import ConstrainedPartitionsIntRange, InclusiveRange
 
 
 class DenseTensorSchema(TensorSchema[np.ndarray]):
@@ -21,19 +21,26 @@ def __init__(self, **kwargs: Any):
             )
 
     @property
-    def key_range(self) -> IntRange:
+    def key_range(self) -> ConstrainedPartitionsIntRange:
+        self._key_range: ConstrainedPartitionsIntRange
         try:
-            key_dim_slice = self._dim_selectors[0]
-        except KeyError:
+            return self._key_range
+        except AttributeError:
             key_dim_min, key_dim_max = self._ned[0]
-        else:
-            assert isinstance(key_dim_slice, slice)
-            key_dim_min = key_dim_slice.start
-            key_dim_max = key_dim_slice.stop
-            raise NotImplementedError(
-                "Key dimension slicing is not yet implemented for dense arrays"
+            key_dim_slice = self._dim_selectors.get(0)
+            if key_dim_slice is not None:
+                assert isinstance(key_dim_slice, slice)
+                min_key = key_dim_slice.start
+                max_key = key_dim_slice.stop
+            else:
+                min_key = key_dim_min
+                max_key = key_dim_max
+            key_dim_tile = self._array.dim(self.key_dim).tile
+            start_offsets = range(key_dim_min, key_dim_max + 1, key_dim_tile)
+            self._key_range = ConstrainedPartitionsIntRange(
+                min_key, max_key, start_offsets
             )
-        return IntRange(key_dim_min, key_dim_max)
+            return self._key_range
 
     def iter_tensors(
         self, key_ranges: Iterable[InclusiveRange[int, int]]

tiledb/ml/readers/types.py

@@ -30,11 +30,11 @@ class ArrayParams:
     - array: TileDB array to be accessed
     - key_dim: Name (or index) of the array key dimension. Defaults to the first dimension.
     - fields: Fields (dimensions and attributes) to be retrieved from array. Defaults to
-        all attributes of the array.
+      all attributes of the array.
     - dim_selectors: Mapping from dimension name to a slice or sequence of indices of this
-        dimension to select. Currently implemented only for non-key dimensions of dense arrays
-    - tensor_kind: kind of tensor desired. If not specified, the default tensor kind is
-        determined based on the array schema.
+      dimension to select.
+    - tensor_kind: kind of tensor desired. If not specified, it is determined based on the
+      array schema.
     """
 
     array: tiledb.Array