Fix Bug For QueryCondition With Unselected Boolean Type (#1291)

Author: Vivian Nguyen

HISTORY.md

@@ -1,3 +1,8 @@
+# In Progress
+
+## Bug Fixes
+* Fix issue where querying an array with a Boolean type when `arrow=True`, but is unselected in `.query(attr=...)`, results in an error `pyarrow.lib.ArrowInvalid: Invalid column index to set field.` []()
+
 # TileDB-Py 0.17.1 Release Notes
 
 ## API Changes

tiledb/multirange_indexing.py

@@ -349,12 +349,16 @@ def _run_query(self) -> Union[DataFrame, Table]:
             # this is a workaround to cast TILEDB_BOOL types from uint8
             # representation in Arrow to Boolean
             schema = table.schema
-            for n in range(self.array.nattr):
-                attr = self.array.attr(n)
+            for attr_or_dim in schema:
+                if not self.array.schema.has_attr(attr_or_dim.name):
+                    continue
+
+                attr = self.array.attr(attr_or_dim.name)
                 if attr.dtype == bool:
                     field_idx = schema.get_field_index(attr.name)
                     field = pyarrow.field(attr.name, pyarrow.bool_())
                     schema = schema.set(field_idx, field)
+
             table = table.cast(schema)
 
             if self.query.return_arrow:

tiledb/tests/test_query_condition.py

@@ -1,8 +1,6 @@
 import pytest
 
-import math
 import numpy as np
-from numpy.testing import assert_array_equal
 import string
 
 import tiledb
@@ -638,3 +636,32 @@ def test_dense_datetime(self):
             result = A.query(attr_cond=qc).df[:]
 
             assert all(self.filter_dense(result["dates"], dt_mask) == A[idx]["dates"])
+
+    def test_array_with_bool_but_unused(self):
+        path = self.path("test_array_with_bool_but_unused")
+
+        dom = tiledb.Domain(
+            tiledb.Dim(name="d", domain=(1, 3), tile=1, dtype=np.uint32)
+        )
+        attrs = [
+            tiledb.Attr(name="myint", dtype=int),
+            tiledb.Attr(name="mystr", dtype=str),
+            tiledb.Attr(name="mybool", dtype=bool),
+        ]
+
+        schema = tiledb.ArraySchema(domain=dom, attrs=attrs, sparse=True)
+        tiledb.Array.create(path, schema)
+
+        data = {
+            "myint": np.asarray([10, 20, 30]),
+            "mystr": np.asarray(["apple", "ball", "cat"]),
+            "mybool": np.asarray([True, False, True]),
+        }
+
+        with tiledb.open(path, "w") as A:
+            A[np.arange(1, 4)] = data
+
+        with tiledb.open(path) as A:
+            qc = tiledb.QueryCondition("myint > 10")
+            result = A.query(attr_cond=qc, attrs=["myint"])[:]
+            assert all(result["myint"] > 10)