Support QueryCondition on Dense Arrays (#1198)

Author: nguyenv

HISTORY.md

@@ -4,7 +4,8 @@
 * `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)
+* Support `QueryCondition` for dense arrays [#1198](https://github.com/TileDB-Inc/TileDB-Py/pull/1198)
+* 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
 
@@ -17,9 +18,9 @@
 * TileDB-Py 0.16.0 includes TileDB Embedded [TileDB 2.10.0](https://github.com/TileDB-Inc/TileDB/releases/tag/2.10.0)
 
 ## API Changes
-* Addition of Filestore API [#1070](https://github.com/TileDB-Inc/TileDB-Py/pull/1070)
+* Addition of `Filestore` API [#1070](https://github.com/TileDB-Inc/TileDB-Py/pull/1070)
 * Use `bool` instead of `uint8` for Boolean dtype in `dataframe_.py` [#1154](https://github.com/TileDB-Inc/TileDB-Py/pull/1154)
-* Support QueryCondition OR operator [#1146](https://github.com/TileDB-Inc/TileDB-Py/pull/1146)
+* Support `QueryCondition` OR operator [#1146](https://github.com/TileDB-Inc/TileDB-Py/pull/1146)
 
 # TileDB-Py 0.15.6 Release Notes
 

examples/query_condition_dense.py

@@ -0,0 +1,105 @@
+# query_condition_dense.py
+#
+# LICENSE
+#
+# The MIT License
+#
+# Copyright (c) 2021 TileDB, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+
+# This example creates an array with one string-typed attribute,
+# writes sample data to the array, and then prints out a filtered
+# dataframe using the TileDB QueryCondition feature.
+
+import tiledb
+import numpy as np
+from pprint import pprint
+import tempfile
+import string
+
+uri = "query_condition_dense"
+
+
+def create_array(path):
+    # create a dense array
+    dom = tiledb.Domain(
+        tiledb.Dim(name="coords", domain=(1, 10), tile=1, dtype=np.uint32)
+    )
+    attrs = [
+        tiledb.Attr(name="attr1", dtype=np.uint64),
+        tiledb.Attr(name="attr2", dtype=np.float64),
+    ]
+    schema = tiledb.ArraySchema(domain=dom, attrs=attrs, sparse=False)
+    tiledb.Array.create(path, schema, overwrite=True)
+
+    # fill array with randomized values
+    with tiledb.open(path, "w") as arr:
+        rand = np.random.default_rng()
+        arr[:] = {
+            "attr1": rand.integers(low=0, high=10, size=10),
+            "attr2": rand.random(size=10),
+        }
+
+
+def read_array(path):
+    with tiledb.open(uri) as arr:
+        print("--- without query condition:")
+        print()
+        pprint(arr[:])
+        print()
+
+    with tiledb.open(uri) as arr:
+        qc = tiledb.QueryCondition("(2 < attr1 < 6) and (attr2 < 0.5 or attr2 > 0.85)")
+        print(f"--- with query condition {qc}:")
+
+        print(f"--- the fill value for attr1 is {arr.attr('attr1').fill}")
+        print(f"--- the fill value for attr1 is {arr.attr('attr2').fill}")
+
+        print()
+        res = arr.query(attr_cond=qc)[:]
+        pprint(res)
+
+
+if __name__ == "__main__":
+    """Example output for `python query_condition_dense.py`:
+
+    --- without query condition:
+
+    OrderedDict([('attr1', array([4, 0, 9, 7, 6, 0, 0, 5, 7, 5], dtype=uint64)),
+                ('attr2',
+                array([0.74476144, 0.47211544, 0.99054245, 0.36640416, 0.91699594,
+        0.06216043, 0.58581863, 0.00505695, 0.7486192 , 0.87649422]))])
+
+    --- with query condition QueryCondition(expression='(2 < attr1 < 6) and (attr2 < 0.5 or attr2 > 0.85)'):
+    --- the fill value for attr1 is [18446744073709551615]
+    --- the fill value for attr1 is [nan]
+
+    OrderedDict([('attr1',
+                array([18446744073709551615, 18446744073709551615, 18446744073709551615,
+        18446744073709551615, 18446744073709551615, 18446744073709551615,
+        18446744073709551615,                    5, 18446744073709551615,
+                            5], dtype=uint64)),
+                ('attr2',
+                array([       nan,        nan,        nan,        nan,        nan,
+                nan,        nan, 0.00505695,        nan, 0.87649422]))])
+    """
+    create_array(uri)
+    read_array(uri)

examples/query_condition_sparse.py

@@ -0,0 +1,96 @@
+# query_condition_sparse.py
+#
+# LICENSE
+#
+# The MIT License
+#
+# Copyright (c) 2021 TileDB, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+
+# This example creates an array with one string-typed attribute,
+# writes sample data to the array, and then prints out a filtered
+# dataframe using the TileDB QueryCondition feature.
+
+import tiledb
+import numpy as np
+from pprint import pprint
+import tempfile
+import string
+
+uri = "query_condition_sparse"
+
+
+def create_array(path):
+    # create a sparse array
+    dom = tiledb.Domain(
+        tiledb.Dim(name="coords", domain=(1, 10), tile=1, dtype=np.uint32)
+    )
+    attrs = [
+        tiledb.Attr(name="attr1", dtype=np.uint64),
+        tiledb.Attr(name="attr2", dtype=np.float64),
+    ]
+    schema = tiledb.ArraySchema(domain=dom, attrs=attrs, sparse=True)
+    tiledb.Array.create(path, schema, overwrite=True)
+
+    # fill array with randomized values
+    with tiledb.open(path, "w") as arr:
+        rand = np.random.default_rng()
+        arr[np.arange(1, 11)] = {
+            "attr1": rand.integers(low=0, high=10, size=10),
+            "attr2": rand.random(size=10),
+        }
+
+
+def read_array(path):
+    with tiledb.open(uri) as arr:
+        print("--- without query condition:")
+        print()
+        pprint(arr[:])
+        print()
+
+    with tiledb.open(uri) as arr:
+        qc = tiledb.QueryCondition("(2 < attr1 < 6) and (attr2 < 0.5 or attr2 > 0.85)")
+        print(f"--- with query condition {qc}:")
+        print()
+        res = arr.query(attr_cond=qc)[:]
+        pprint(res)
+
+
+if __name__ == "__main__":
+    """Example output for `python query_condition_sparse.py`:
+
+    --- without query condition:
+
+    OrderedDict([('attr1', array([2, 4, 4, 3, 4, 7, 5, 2, 2, 8], dtype=uint64)),
+                ('attr2',
+                array([0.62445071, 0.32415481, 0.39117764, 0.66609931, 0.48122102,
+        0.93561984, 0.70998524, 0.10322076, 0.28343041, 0.33623958])),
+                ('coords',
+                array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10], dtype=uint32))])
+
+    --- with query condition QueryCondition(expression='(2 < attr1 < 6) and (attr2 < 0.5 or attr2 > 0.85)'):
+
+    OrderedDict([('attr1', array([4, 4, 4], dtype=uint64)),
+                ('attr2', array([0.32415481, 0.39117764, 0.48122102])),
+                ('coords', array([2, 3, 5], dtype=uint32))])
+    """
+    create_array(uri)
+    read_array(uri)

tiledb/libtiledb.pyx

@@ -109,6 +109,7 @@ _tiledb_dtype_to_numpy_typeid_convert ={
     TILEDB_INT16: np.NPY_INT16,
     TILEDB_UINT16: np.NPY_UINT16,
     TILEDB_CHAR: np.NPY_STRING,
+    TILEDB_STRING_ASCII: np.NPY_STRING,
     TILEDB_STRING_UTF8: np.NPY_UNICODE,
 }
 IF LIBTILEDB_VERSION_MAJOR >= 2:
@@ -131,7 +132,7 @@ _tiledb_dtype_to_numpy_dtype_convert = {
     TILEDB_INT16: np.int16,
     TILEDB_UINT16: np.uint16,
     TILEDB_CHAR: np.dtype('S1'),
-    TILEDB_STRING_ASCII: np.bytes_,
+    TILEDB_STRING_ASCII: np.dtype('S'),
     TILEDB_STRING_UTF8: np.dtype('U1'),
 }
 IF LIBTILEDB_VERSION_MAJOR >= 2:
@@ -4102,8 +4103,6 @@ cdef class Query(object):
                     raise TileDBError(f"Selected attribute does not exist: '{name}'")
         self.attrs = attrs
         self.attr_cond = attr_cond
-        if attr_cond is not None and not array.schema.sparse:
-            raise TileDBError("QueryConditions may only be applied to sparse arrays")
 
         if order == None:
             if array.schema.sparse:

tiledb/query_condition.py

@@ -21,7 +21,18 @@
 class QueryCondition:
     """
     Class representing a TileDB query condition object for attribute filtering
-    pushdown. Set the query condition with a string representing an expression
+    pushdown.
+
+    When querying a sparse array, only the values that satisfy the given
+    condition are returned (coupled with their associated coordinates). An example
+    may be found in `examples/query_condition_sparse.py`.
+
+    For dense arrays, the given shape of the query matches the shape of the output
+    array. Values that DO NOT satisfy the given condition are filled with the
+    TileDB default fill value. Different attribute types have different default
+    fill values as outlined here (https://docs.tiledb.com/main/background/internal-mechanics/writing#default-fill-values). An example may be found in `examples/query_condition_dense.py`.
+
+    Set the query condition with a string representing an expression
     as defined by the grammar below. A more straight forward example of usage is
     given beneath.