Repository: TileDB-Inc/TileDB-Py

Issue:
Support for orthogonal indexing
I ran this on TileDB-py 0.3.0 with TileDB 1.4.0.

It seems like the slicing operation from an SparseArray or a DenseArray would only support queries such as:

a[1, 30:40]

However, if I have a list of values to query (say 1, 4, 10), it would fail:

a[[1, 4, 10], 30:40]

The only way I can think of doing this is to iterate through the list and do a slice like:

a[1, 30:40]
a[4, 30:40]
a[10, 30:40]

Is there a more efficient to do this?  I'd imagine doing this in python would be pretty slow.  If I have more dimensions it can get more complicated.  For instance, if I know the exact index for the things I want to query say for a 3d array:

a[[1, 4, 10], [20, 30, 40], 30:40]

Then I have to get the cartesian  production of [1, 4, 10] and [20, 30, 40] and slice through the rest.

In zarr for instance, there's a special function to do this  (https://zarr.readthedocs.io/en/stable/tutorial.html#orthogonal-indexing).  Do you see this is a viable option to implement in TileDB?


Discussion:
@will133  you are right in that orthogonal indexing is not supported (currently). In core TIleDB this issue is tracked by https://github.com/TileDB-Inc/TileDB/issues/675.  We could implement a solution similar to Zarr, but this would degenerate into multiple different subqueries to TileDB.  Pushing this to TileDB would enable us to do it in a more efficient way, but there is no technical reason why we could not at least match Zarr here in the HL Python API (in Zarr the performance is probably acceptable, but not great).

It seems like the multi-array reading in the C++ library is implemented and marked as resolved.  Do you know what's the plan of implementing this on the python side?
Hi @will133, we do plan to support incorporate these changes in the next version of TileDB / TileDB-PY.  This feature was a lot of work!  (it was developed under the constraint of supporting 10M+ range queries at once on a 100 TB uncompressed array).
Sounds impressive!  Thanks for all the hard work!
The proposed API will follow the spec. in the following NEP https://www.numpy.org/neps/nep-0021-advanced-indexing.html#nep-21-simplified-and-explicit-advanced-indexing, as Numpy only supports vectorized indexing.  TileDB does not currently support vectorized indexing (so only `Array.oindex` will be provided for now).  This is similar to the API provided in Zarr, XArray.

Resolving patch:
diff --git a/tiledb/highlevel.py b/tiledb/highlevel.py
--- a/tiledb/highlevel.py
+++ b/tiledb/highlevel.py
@@ -111,4 +111,4 @@ def array_exists(uri, isdense=False, issparse=False):
         rval = True
 
     a.close()
-    return rval
+    return rval
\ No newline at end of file
diff --git a/tiledb/multirange_indexing.py b/tiledb/multirange_indexing.py
new file mode 100644
--- /dev/null
+++ b/tiledb/multirange_indexing.py
@@ -0,0 +1,121 @@
+import tiledb
+from tiledb import Array, ArraySchema
+import os, numpy as np
+import sys
+
+try:
+    from tiledb.libtiledb import multi_index
+except:
+    from tiledb.indexing import multi_index
+
+def _index_as_tuple(idx):
+    """Forces scalar index objects to a tuple representation"""
+    if isinstance(idx, tuple):
+        return idx
+    return (idx,)
+
+def mr_dense_result_shape(ranges, base_shape = None):
+    # assumptions: len(ranges) matches number of dims
+    if base_shape is not None:
+        assert len(ranges) == len(base_shape), "internal error: mismatched shapes"
+
+    new_shape = list()
+    for i,rr in enumerate(ranges):
+        if rr is not ():
+            m = list(map(lambda y: abs(y[1] - y[0]) + 1, rr))
+            new_shape.append(np.sum(m))
+        else:
+            if base_shape is None:
+                raise ValueError("Missing required base_shape for whole-dimension slices")
+            # empty range covers dimension
+            new_shape.append(base_shape[i])
+
+    return tuple(new_shape)
+
+def mr_dense_result_numel(ranges):
+    return np.prod(mr_dense_result_shape(ranges))
+
+def sel_to_subranges(dim_sel):
+    if isinstance(dim_sel, list):
+        dim_sel = tuple(dim_sel)
+    elif not isinstance(dim_sel, tuple):
+        dim_sel = (dim_sel,)
+
+    subranges = list()
+    for range in dim_sel:
+        if np.isscalar(range):
+            subranges.append( (range, range) )
+        elif isinstance(range, slice):
+            if range.step is not None:
+                raise ValueError("Stepped slice ranges are not supported")
+            elif range.start is None and range.stop is None:
+                # ':' full slice
+                pass
+            else:
+                subranges.append( (range.start, range.stop) )
+        elif isinstance(range, list):
+            for el in range:
+                subranges.append( (el, el) )
+        else:
+            raise TypeError("Unsupported selection ")
+
+    return tuple(subranges)
+
+
+class MultiRangeIndexer(object):
+    """
+    Implements multi-range / outer / orthogonal indexing.
+
+    """
+    # for cython
+    # comment out for Python 2 :/
+    #array: Array
+    #schema: ArraySchema
+    #def __init__(self, array: Array, query = None):
+
+    def __init__(self, array, query = None):
+        self.array = array
+        # TODO remove
+        if hasattr(array, 'schema'):
+            self.schema = array.schema
+        self.query = query
+
+    def getitem_ranges(self, idx):
+        dom = self.schema.domain
+        ndim = dom.ndim
+        idx = _index_as_tuple(idx)
+
+        ranges = list()
+        for i,sel in enumerate(idx):
+            subranges = sel_to_subranges(sel)
+            ranges.append(subranges)
+
+        # extend the list to ndim
+        if len(ranges) < ndim:
+            ranges.extend([ tuple() for _ in range(ndim-len(ranges))])
+
+        rval = tuple(ranges)
+        return rval
+
+    def __getitem__(self, idx):
+        # implements multi-range / outer / orthogonal indexing
+        ranges = self.getitem_ranges(idx)
+
+        dom = self.schema.domain
+        attr_names = tuple(self.schema.attr(i).name for i in range(self.schema.nattr))
+
+        # TODO order
+        result_dict = multi_index(
+            self.array,
+            attr_names,
+            ranges
+        )
+
+        if self.schema.sparse:
+            return result_dict
+        else:
+            result_shape = mr_dense_result_shape(ranges, self.schema.shape)
+            for arr in result_dict.values():
+                # TODO check/test layout
+                arr.shape = result_shape
+            return result_dict
\ No newline at end of file


Test patch:
diff --git a/tiledb/tests/common.py b/tiledb/tests/common.py
--- a/tiledb/tests/common.py
+++ b/tiledb/tests/common.py
@@ -46,6 +46,14 @@ def assert_subarrays_equal(a, b):
     for a_el, b_el in zip(a.flat, b.flat):
         assert_array_equal(a_el, b_el)
 
+def assert_all_arrays_equal(*arrays):
+    # TODO this should display raise in the calling location if possible
+    assert len(arrays) % 2 == 0, \
+           "Expected even number of arrays"
+
+    for a1,a2 in zip(arrays[0::2], arrays[1::2]):
+        assert_array_equal(a1, a2)
+
 # python 2 vs 3 compatibility
 if sys.hexversion >= 0x3000000:
     getchr = chr
@@ -83,3 +91,23 @@ def rand_int_sequential(size, dtype=np.uint64):
     arr = np.random.randint(dtype_max(dtype), size=size, dtype=dtype)
     return np.sort(arr)
 
+def intspace(start, stop, num=50, dtype=np.int64):
+    """
+    Return evenly spaced values over range ensuring that stop is
+    always the maximum (will not overflow with int dtype as linspace)
+    :param start:
+    :param stop:
+    :param num:
+    :param dtype:
+    :return:
+    """
+    rval = np.zeros(num, dtype=dtype)
+    step = (stop-start)//num
+    nextval = start
+
+    for i in range(num):
+        rval[i] = nextval
+        nextval += step
+
+    rval[-1] = stop
+    return rval
diff --git a/tiledb/tests/test_multi_index.py b/tiledb/tests/test_multi_index.py
new file mode 100644
--- /dev/null
+++ b/tiledb/tests/test_multi_index.py
@@ -0,0 +1,447 @@
+"""
+TODO
+- # implement mock of expected behavior in pure numpy w/ test function
+
+
+- implement read function and tests (single [x], multi-attribute [ ])
+- implement custom indexer
+- implement oindex...
+"""
+
+import tiledb
+from tiledb.multirange_indexing import *
+import os, numpy as np
+import sys
+
+from numpy.testing import assert_array_equal
+import unittest
+from tiledb.tests.common import *
+
+def make_1d_dense(ctx, path, attr_name=''):
+    a_orig = np.arange(36)
+
+    dom = tiledb.Domain(tiledb.Dim(domain=(0, 35), tile=35, dtype=np.uint64, ctx=ctx),
+                        ctx=ctx)
+    att = tiledb.Attr(name=attr_name, dtype=np.int64, ctx=ctx)
+    schema = tiledb.ArraySchema(domain=dom, attrs=(att,), sparse=False, ctx=ctx)
+    tiledb.DenseArray.create(path, schema)
+
+    with tiledb.DenseArray(path, 'w') as A:
+        A[:] = a_orig
+
+
+def make_2d_dense(ctx, path, attr_name=''):
+    a_orig = np.arange(1,37).reshape(9, 4)
+
+    dom = tiledb.Domain(tiledb.Dim(domain=(0, 8), tile=9, dtype=np.uint64, ctx=ctx),
+                        tiledb.Dim(domain=(0, 3), tile=4, dtype=np.uint64, ctx=ctx),
+                        ctx=ctx)
+    att = tiledb.Attr(name=attr_name, dtype=np.int64, ctx=ctx)
+    schema = tiledb.ArraySchema(domain=dom, attrs=(att,), sparse=False, ctx=ctx)
+    tiledb.DenseArray.create(path, schema)
+
+    with tiledb.DenseArray(path, 'w') as A:
+        A[:] = a_orig
+
+class TestMultiRange(DiskTestCase):
+
+    def test_mr_1d_1dim_ranges(self):
+        path = self.path('test_mr_1d_1dim_ranges')
+        attr_name = 'a'
+
+        ctx = tiledb.Ctx()
+        make_1d_dense(ctx, path, attr_name=attr_name)
+
+        expected = np. array([0],
+                             dtype=np.uint64)
+
+
+        with tiledb.DenseArray(path) as A:
+            ranges = ( ((0, 0),), )
+            expected = np.array([0], dtype=np.int64)
+            a = tiledb.libtiledb.multi_index(A, (attr_name,), ranges)[attr_name]
+            assert_array_equal(a, expected)
+            self.assertEqual(a.dtype, expected.dtype)
+
+            #===
+        #with tiledb.DenseArray(path) as A:
+            ranges2 = ( ((1, 1), (5,8)), )
+            expected2 = np.array([1,5,6,7,8], dtype=np.int64)
+            a2 = tiledb.libtiledb.multi_index(A, (attr_name,), ranges2)[attr_name]
+            assert_array_equal(a2, expected2)
+            self.assertEqual(a2.dtype, expected2.dtype)
+
+    def test_mr_2d_1dim_ranges(self):
+        path = self.path('test_mr_1dim_ranges')
+        attr_name = 'a'
+
+        ctx = tiledb.Ctx()
+        make_2d_dense(ctx, path, attr_name=attr_name)
+
+        expected = np. array([ 1,  2,  3,  4,
+                               21, 22, 23, 24,
+                               25, 26, 27, 28,
+                               29, 30, 31, 32, 33,
+                               34, 35, 36],
+                             dtype=np.uint64)
+
+        ranges = (
+            ( (0, 0), (5,8), ),
+        )
+
+        with tiledb.DenseArray(path) as A:
+            a = tiledb.libtiledb.multi_index(A, (attr_name,), ranges)[attr_name]
+
+            assert_array_equal(a, expected)
+
+    def test_mr_2d_2dim_ranges(self):
+        ctx = tiledb.Ctx()
+        path = self.path('test_mr_2dim_ranges')
+        attr_name = 'a'
+
+        make_2d_dense(ctx, path, attr_name=attr_name)
+
+        expected = np.array([1, 2, 3, 4,
+                             5, 6, 7, 8,
+                             9, 10, 11, 12,
+                             13, 14, 15, 16,
+                             17, 18, 19, 20])
+
+        ranges = (
+            ( (0,4), ),
+            ( (0,3), )
+        )
+
+        with tiledb.DenseArray(path) as A:
+            a = tiledb.libtiledb.multi_index(A, (attr_name,), ranges)[attr_name]
+            assert_array_equal(a, expected)
+
+
+    def test_shape_funcs(self):
+        #-----------
+        range1el = ( ((1,1),), )
+        self.assertEqual(
+            mr_dense_result_shape(range1el),
+            (1,))
+        self.assertEqual(
+            mr_dense_result_numel(range1el),
+            1)
+
+        #-----------
+        range1d = tuple([((1, 2), (4, 4))])
+        self.assertEqual(
+            mr_dense_result_shape(range1d),
+            (3,))
+        self.assertEqual(
+            mr_dense_result_numel(range1d),
+            3
+        )
+
+        #-----------
+        range2d1 = (
+            ( (3,6), (7,7), (10,12) ),
+            ( (5,7), ),
+        )
+
+        self.assertEqual(
+            mr_dense_result_shape(range2d1),
+            (8,3)
+        )
+        self.assertEqual(
+            mr_dense_result_numel(range2d1),
+            24
+        )
+
+        #-----------
+        range2d2 = (
+            [(3,6), (7,7), (10,12)],
+            [(5,7), (10,10)]
+        )
+
+    def test_3d(self):
+        range3d1 = (
+            ( (2,4), ),
+            ( (3,6), ),
+            ( (1,4), (5,9) )
+        )
+
+        #self.assertEqual()
+
+    def test_sel_to_ranges(self):
+        class Obj(object): pass
+        class IBI(object):
+            def __getitem__(self, idx):
+                return idx
+
+        def make_arr(ndim):
+            arr = Obj()
+            arr.schema = Obj()
+            arr.schema.domain = Obj()
+            arr.schema.domain.ndim = ndim
+            return arr
+
+        ibi = IBI()
+        # ndim = 1
+        arr = make_arr(1)
+        m = MultiRangeIndexer(arr)
+        self.assertEqual(
+            m.getitem_ranges( ibi[[1]] ),
+            (((1, 1),),)
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[[1, 2]] ),
+            (((1, 1), (2, 2)),)
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[slice(1, 2)], ),
+            (((1, 2),),)
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[1:2, 3:5] ),
+            (((1, 2),), ((3, 5),),)
+        )
+
+        # ndim = 2
+        arr2 = make_arr(2)
+        m = MultiRangeIndexer(arr2)
+
+        self.assertEqual(
+            m.getitem_ranges( ibi[[1]] ),
+            (((1, 1),), ())
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[slice(1, 33)] ),
+            (((1, 33),), ())
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[ [1, 2], [1], slice(1, 3) ] ),
+            (((1, 1), (2, 2)), ((1, 1),), ((1, 3),))
+        )
+
+        # ndim = 3
+        arr3 = make_arr(3)
+        m = MultiRangeIndexer(arr3)
+
+        self.assertEqual(
+            m.getitem_ranges( ibi[1, 2, 3] ),
+            (
+                ((1, 1),), ((2, 2,),), ((3, 3),)
+            )
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[1, 2] ),
+            (
+                (((1, 1),), ((2, 2,),), ())
+            )
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[1:2, 3:4]),
+            (((1, 2),), ((3, 4),), ())
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[1:2, 3:4, 5:6] ),
+            (((1, 2),), ((3, 4),), ((5, 6),))
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[[1], [2], [5, 6]] ),
+            (((1, 1),), ((2, 2),), ((5, 5), (6, 6),))
+        )
+        self.assertEqual(
+            m.getitem_ranges( ibi[1, [slice(3, 6), 8], slice(4, 6)] ),
+            (
+                ((1, 1),),
+                ((3, 6), (8, 8)),
+                ((4, 6),)
+            )
+        )
+
+    def test_mr_1d_dense_int(self):
+        attr_name = ''
+        ctx = tiledb.Ctx()
+        path = self.path('multi_index_1d')
+
+        dom = tiledb.Domain(tiledb.Dim(domain=(-10, 10), tile=9,
+                                       dtype=np.int64, ctx=ctx),
+                            ctx=ctx)
+        att = tiledb.Attr(name=attr_name, dtype=np.float32, ctx=ctx)
+        schema = tiledb.ArraySchema(domain=dom, attrs=(att,), ctx=ctx)
+        tiledb.DenseArray.create(path, schema)
+
+        orig_array = np.random.rand(schema.domain.dim(0).size).astype(np.float32)
+        with tiledb.open(path, 'w') as A:
+            A[:] = orig_array
+
+        with tiledb.open(path) as A:
+            # stepped ranges are not supported
+            with self.assertRaises(ValueError):
+                A.multi_index[ 1::2 ]
+
+            assert_array_equal(
+                orig_array[ [0,-1] ],
+                A.multi_index[ [-10,10] ][attr_name]
+            )
+            self.assertEqual(
+                orig_array[0],
+                A.multi_index[-10][attr_name]
+            )
+
+    def test_mr_1d_sparse_double(self):
+        attr_name = ''
+        ctx = tiledb.Ctx()
+        path = self.path('mr_1d_sparse_double')
+
+        dom = tiledb.Domain(tiledb.Dim(domain=(0, 30), tile=10,
+                                       dtype=np.float64, ctx=ctx),
+                            ctx=ctx)
+        att = tiledb.Attr(name=attr_name, dtype=np.float64, ctx=ctx)
+        schema = tiledb.ArraySchema(domain=dom, sparse=True, attrs=(att,), ctx=ctx)
+        tiledb.SparseArray.create(path, schema)
+
+        coords = np.linspace(0,30, num=31)
+        orig_array = np.random.rand(coords.size)
+
+        with tiledb.open(path, 'w') as A:
+            A[coords] = orig_array
+
+        with tiledb.open(path) as A:
+            assert_array_equal(
+                orig_array[ [0] ],
+                A.multi_index[ [0] ][attr_name]
+            )
+            assert_array_equal(
+                orig_array[-1],
+                A.multi_index[30][attr_name]
+            )
+
+            res = A.multi_index[ slice(0,5) ]
+            assert_array_equal(
+                orig_array[0:6],
+                res[attr_name]
+            )
+            assert_array_equal(
+                coords[0:6],
+                res['coords'].astype(np.float64)
+            )
+
+    def test_mr_2d_sparse_domain_utypes(self):
+        attr_name = 'foo'
+        ctx = tiledb.Ctx()
+
+        types = (np.uint8, np.uint16, np.uint32, np.uint64)
+
+        for dtype in types:
+            min = 0
+            max = np.iinfo(dtype).max - 1
+            path = self.path('multi_index_2d_sparse_' + str(dtype.__name__))
+
+            dom = tiledb.Domain(tiledb.Dim(domain=(min, max), tile=1,
+                                           dtype=dtype, ctx=ctx),
+                                ctx=ctx)
+
+            att = tiledb.Attr(name=attr_name, dtype=dtype, ctx=ctx)
+            schema = tiledb.ArraySchema(domain=dom, sparse=True, attrs=(att,), ctx=ctx)
+            tiledb.SparseArray.create(path, schema)
+
+            coords = intspace(min, max, num=100, dtype=dtype)
+
+            with tiledb.open(path, 'w') as A:
+                A[coords] = coords
+
+            with tiledb.open(path) as A:
+
+                res = A.multi_index[slice(coords[0], coords[-1])]
+                assert_array_equal(
+                    res[attr_name],
+                    coords
+                )
+                assert_array_equal(
+                    res['coords'].astype(dtype),
+                    coords
+                )
+
+                res = A.multi_index[coords[0]]
+                assert_array_equal(
+                    res[attr_name],
+                    coords[0]
+                )
+                assert_array_equal(
+                    res['coords'].astype(dtype),
+                    coords[0]
+                )
+
+                res = A.multi_index[coords[-1]]
+                assert_array_equal(
+                    res[attr_name],
+                    coords[-1]
+                )
+                assert_array_equal(
+                    res['coords'].astype(dtype),
+                    coords[-1]
+                )
+
+                midpoint = len(coords)//2
+                start = midpoint-20
+                stop = midpoint+20
+                srange = slice(coords[start], coords[stop])
+                res = A.multi_index[ srange ]
+                assert_array_equal(
+                    res[attr_name],
+                    coords[start:stop+1])
+                assert_array_equal(
+                    res['coords'].astype(dtype),
+                    coords[start:stop+1])
+
+    def test_mr_2d_sparse_float(self):
+        attr_name = ''
+        ctx = tiledb.Ctx()
+        path = self.path('mr_2d_sparse_float')
+
+        dom = tiledb.Domain(tiledb.Dim(domain=(0, 10), tile=1,
+                                       dtype=np.float32, ctx=ctx),
+                            tiledb.Dim(domain=(0, 10), tile=1,
+                                       dtype=np.float32, ctx=ctx),
+                            ctx=ctx)
+        att = tiledb.Attr(name=attr_name, dtype=np.float64, ctx=ctx)
+        schema = tiledb.ArraySchema(domain=dom, attrs=(att,), sparse=True, ctx=ctx)
+        tiledb.SparseArray.create(path, schema)
+
+        orig_array = np.random.rand(11,11)
+        d1 = np.linspace(0, 10, num=11, dtype=np.float32)
+        d2 = np.linspace(0, 10, num=11, dtype=np.float32)
+        coords_d1,coords_d2 = np.meshgrid(d1,d2,indexing='ij')
+        coords_d1,coords_d2 = coords_d1.flatten(), coords_d2.flatten()
+        coords = np.column_stack([coords_d1, coords_d2]).reshape(11,11,2)
+
+        with tiledb.open(path, 'w') as A:
+            A[coords_d1,coords_d2] = orig_array
+
+        with tiledb.open(path) as A:
+            res = A.multi_index[[0], :]
+            assert_array_equal(
+                orig_array[[0], :].squeeze(),
+                res[attr_name]
+            )
+            assert_array_equal(
+                coords[[0], :].squeeze(),
+                res['coords'].view('f4').reshape(-1,2)
+            )
+
+            res = A.multi_index[10, :]
+            assert_array_equal(
+                orig_array[[-1], :].squeeze(),
+                res[attr_name]
+            )
+            assert_array_equal(
+                coords[[-1], :].squeeze(),
+                res['coords'].view('f4').reshape(-1,2)
+            )
+
+            res = A.multi_index[ [ slice(0,2), [5]] ]
+            assert_array_equal(
+                np.vstack([orig_array[0:3,:], orig_array[5,:]]).flatten(),
+                res[attr_name]
+            )
+            assert_array_equal(
+                np.vstack([coords[0:3,:].reshape(-1,2), coords[5,:]]),
+                res['coords'].view('f4').reshape(-1,2)
+            )