Dense Array [:] Returns Nonempty Domain Only (#1199)

Author: nguyenv

HISTORY.md

@@ -3,6 +3,9 @@
 ## Improvements
 * `setup.py` retrieves core version by using `ctypes` to call `tiledb_version` rather than parsing `tiledb_version.h` [#1191](https://github.com/TileDB-Inc/TileDB-Py/pull/1191)
 
+## API Changes
+* Querying dense array with `[:]` returns shape that matches nonempty domain, consistent with `.df[:]` and .multi_index[:]` [#1199](https://github.com/TileDB-Inc/TileDB-Py/pull/1199)
+
 # TileDB-Py 0.16.1 Release Notes
 
 ## TileDB Embedded updates:

tiledb/libtiledb.pyx

@@ -2525,7 +2525,7 @@ def replace_scalars_slice(dom: Domain, idx: tuple):
     return tuple(new_idx), tuple(drop_axes)
 
 
-def index_domain_subarray(array: Array, dom: Domain, idx: tuple):
+def index_domain_subarray(array: Array, dom: Domain, idx: tuple, read_array: bool = False):
     """
     Return a numpy array representation of the tiledb subarray buffer
     for a given domain and tuple of index slices
@@ -2543,9 +2543,17 @@ def index_domain_subarray(array: Array, dom: Domain, idx: tuple):
 
         if np.issubdtype(dim_dtype, np.unicode_) or np.issubdtype(dim_dtype, np.bytes_):
             ned = array.nonempty_domain()
-            (dim_lb, dim_ub) = ned[r] if ned else (None, None)
+            (dim_lb, dim_ub) = ned[r] if ned is not None else (None, None)
         else:
-            (dim_lb, dim_ub) = dim.domain
+            if read_array:
+                ned = array.nonempty_domain()
+                (dim_lb, dim_ub) = (
+                    np.array(ned[r], dim_dtype)
+                    if ned is not None
+                    else dim.domain
+                )
+            else:
+                (dim_lb, dim_ub) = dim.domain
 
 
         dim_slice = idx[r]
@@ -4437,7 +4445,7 @@ cdef class DenseArrayImpl(Array):
         selection = index_as_tuple(selection)
         idx = replace_ellipsis(self.schema.domain.ndim, selection)
         idx, drop_axes = replace_scalars_slice(self.schema.domain, idx)
-        subarray = index_domain_subarray(self, self.schema.domain, idx)
+        subarray = index_domain_subarray(self, self.schema.domain, idx, True)
         # Note: we included dims (coords) above to match existing semantics
         out = self._read_dense_subarray(subarray, attr_names, attr_cond, layout,
                                         coords)

tiledb/tests/test_libtiledb.py

@@ -1801,6 +1801,21 @@ def test_reopen_dense_array(self):
 
         T.close()
 
+    def test_writing_partial_dense_array(self):
+        uri = self.path("test_writing_partial_dense_array")
+
+        dom = tiledb.Domain(tiledb.Dim(domain=(0, 1000), tile=1000, dtype=np.int64))
+        attr = tiledb.Attr(name="rows", dtype=np.int64)
+        schema = tiledb.ArraySchema(domain=dom, sparse=False, attrs=[attr])
+        tiledb.Array.create(uri, schema)
+
+        with tiledb.open(uri, "w") as A:
+            A[0:1000] = np.arange(1000)
+
+        with tiledb.open(uri) as A:
+            ned = A.nonempty_domain()[0]
+            assert len(A[:]["rows"]) == ned[1] - ned[0] + 1
+
 
 class TestVarlen(DiskTestCase):
     def test_varlen_write_bytes(self):