diff --git a/tiledb/dense_array.py b/tiledb/dense_array.py index 1722480943..be470efde8 100644 --- a/tiledb/dense_array.py +++ b/tiledb/dense_array.py @@ -77,8 +77,7 @@ def __getitem__(self, selection): """ if self.view_attr: - result = self.subarray(selection, attrs=(self.view_attr,)) - return result[self.view_attr] + return self.subarray(selection) result = self.subarray(selection) for i in range(self.schema.nattr): @@ -269,6 +268,8 @@ def subarray(self, selection, attrs=None, cond=None, coords=False, order=None): attr = self.schema.attr(0) if attr.isanon: return out[attr._internal_name] + if self.view_attr is not None: + return out[self.view_attr] return out def _read_dense_subarray( diff --git a/tiledb/sparse_array.py b/tiledb/sparse_array.py index 255dc2e72f..a2d5fe0072 100644 --- a/tiledb/sparse_array.py +++ b/tiledb/sparse_array.py @@ -1,3 +1,4 @@ +import warnings from collections import OrderedDict import numpy as np @@ -292,6 +293,9 @@ def __getitem__(self, selection): >>> # A[5.0:579.9] """ + if self.view_attr is not None: + return self.subarray(selection) + result = self.subarray(selection) for i in range(self.schema.nattr): attr = self.schema.attr(i) @@ -506,7 +510,11 @@ def subarray(self, selection, coords=True, attrs=None, cond=None, order=None): attr_names = list() - if attrs is None: + if self.view_attr is not None: + if attrs is not None: + warnings.warn("view_attr is set, ignoring attrs parameter", UserWarning) + attr_names.extend(self.view_attr) + elif attrs is None: attr_names.extend( self.schema.attr(i)._internal_name for i in range(self.schema.nattr) ) diff --git a/tiledb/tests/test_libtiledb.py b/tiledb/tests/test_libtiledb.py index 4f5e7655bb..2447b1e9fe 100644 --- a/tiledb/tests/test_libtiledb.py +++ b/tiledb/tests/test_libtiledb.py @@ -23,6 +23,7 @@ from .common import ( DiskTestCase, assert_captured, + assert_dict_arrays_equal, assert_subarrays_equal, assert_unordered_equal, fx_sparse_cell_order, # noqa: F401 @@ -923,8 +924,8 @@ def assert_ts(timestamp, result): assert_ts((timestamps[2], None), A * 3) assert_ts((timestamps[2], None), A * 3) - def test_open_attr(self): - uri = self.path("test_open_attr") + def test_open_attr_dense(self): + uri = self.path("test_open_attr_dense") schema = tiledb.ArraySchema( domain=tiledb.Domain( tiledb.Dim(name="dim0", dtype=np.uint32, domain=(1, 4)) @@ -949,6 +950,48 @@ def test_open_attr(self): assert_array_equal(A[:], np.array((1, 2, 3, 4))) assert list(A.multi_index[:].keys()) == ["x"] + with tiledb.open(uri, attr="x") as A: + q = A.query(cond="x <= 3") + expected = np.array([1, 2, 3, schema.attr("x").fill[0]]) + assert_array_equal(q[:], expected) + + def test_open_attr_sparse(self): + uri = self.path("test_open_attr_sparse") + schema = tiledb.ArraySchema( + domain=tiledb.Domain( + tiledb.Dim(name="dim0", dtype=np.uint32, domain=(1, 4)) + ), + attrs=( + tiledb.Attr(name="x", dtype=np.int32), + tiledb.Attr(name="y", dtype=np.int32), + ), + sparse=True, + ) + tiledb.Array.create(uri, schema) + + with tiledb.open(uri, mode="w") as A: + A[[1, 2, 3, 4]] = {"x": np.array((1, 2, 3, 4)), "y": np.array((5, 6, 7, 8))} + + with self.assertRaises(KeyError): + tiledb.open(uri, attr="z") + + with self.assertRaises(KeyError): + tiledb.open(uri, attr="dim0") + + with tiledb.open(uri, attr="x") as A: + expected = OrderedDict( + [("dim0", np.array([1, 2, 3, 4])), ("x", np.array([1, 2, 3, 4]))] + ) + assert_dict_arrays_equal(A[:], expected) + assert list(A.multi_index[:].keys()) == ["dim0", "x"] + + with tiledb.open(uri, attr="x") as A: + q = A.query(cond="x <= 3") + expected = OrderedDict( + [("dim0", np.array([1, 2, 3])), ("x", np.array([1, 2, 3]))] + ) + assert_dict_arrays_equal(q[:], expected) + def test_ncell_attributes(self): dom = tiledb.Domain(tiledb.Dim(domain=(0, 9), tile=10, dtype=int)) attr = tiledb.Attr(dtype=[("", np.int32), ("", np.int32), ("", np.int32)])