Repository: TileDB-Inc/TileDB-Py
Pull request #2171 (2025-03-05 23:18:43)

Fix for query conditions that contain attributes with dot(s) when using the `in` operator
Similarly to https://github.com/single-cell-data/TileDB-SOMA/pull/3746, let's fix the issue where a query condition contains attributes with dot(s) when using the `in` operator.

---

[sc-63843]


Patch:
diff --git a/tiledb/query_condition.py b/tiledb/query_condition.py
index 0902906c..8606460b 100644
--- a/tiledb/query_condition.py
+++ b/tiledb/query_condition.py
@@ -212,7 +212,33 @@ class QueryConditionTree(ast.NodeVisitor):
                     "`in` operator syntax must be written as `variable in ['l', 'i', 's', 't']`"
                 )
 
-            variable = node.left.id
+            # For 'my_string in ["red", "yellow"]': node.left is ast.Name
+            # For 'attr(my.string) in ["red", "yellow"]': node.left is ast.Call
+            if isinstance(node.left, ast.Call):
+                if node.left.func.id != "attr":
+                    raise TileDBError(
+                        f"query condition left-hand side function call must be 'attr'; got '{node.left.func.id}'"
+                    )
+
+                if len(node.left.args) != 1:
+                    raise TileDBError(
+                        f"query condition left-hand side 'attr' function call must have one argument; got '{len(node.left.args)}'"
+                    )
+
+                arg = node.left.args[0]
+                if not isinstance(arg, ast.Constant):
+                    raise TileDBError(
+                        "query condition left-hand side 'attr' argument must be a constant"
+                    )
+                variable = arg.value
+
+            elif isinstance(node.left, ast.Name):
+                variable = node.left.id
+            else:
+                raise TileDBError(
+                    f"cannot handle query condition left-hand side of type '{type(node.left)}'"
+                )
+
             values = [self.get_value_from_node(val) for val in self.visit(rhs)]
             if len(values) == 0:
                 raise TileDBError(
@@ -230,7 +256,7 @@ class QueryConditionTree(ast.NodeVisitor):
 
             dtype = "string" if dt.kind in "SUa" else dt.name
             op = qc.TILEDB_IN if isinstance(operator, ast.In) else qc.TILEDB_NOT_IN
-            result = self.create_pyqc(dtype)(self.ctx, node.left.id, values, op)
+            result = self.create_pyqc(dtype)(self.ctx, variable, values, op)
 
         else:
             raise TileDBError(f"unrecognized operator in <<{ast.dump(node)}>>")


Test patch:
diff --git a/tiledb/tests/test_query_condition.py b/tiledb/tests/test_query_condition.py
index e41d319b..1f69539d 100644
--- a/tiledb/tests/test_query_condition.py
+++ b/tiledb/tests/test_query_condition.py
@@ -765,11 +765,16 @@ class QueryConditionTest(DiskTestCase):
         attrs = [
             tiledb.Attr(name="attr.one", dtype=np.uint32),
             tiledb.Attr(name="attr.two", dtype=np.uint32),
+            tiledb.Attr(name="at.tr.thr.ee", dtype=np.uint32),
         ]
         schema = tiledb.ArraySchema(domain=dom, attrs=attrs, sparse=True)
         tiledb.Array.create(path, schema)
         with tiledb.open(path, "w") as A:
-            A[np.arange(11)] = {"attr.one": np.arange(11), "attr.two": np.arange(11)}
+            A[np.arange(11)] = {
+                "attr.one": np.arange(11),
+                "attr.two": np.arange(11),
+                "at.tr.thr.ee": np.arange(11),
+            }
         with tiledb.open(path, "r") as A:
             with pytest.raises(tiledb.TileDBError) as exc_info:
                 A.query(cond="attr.one < 6")[:]
@@ -783,12 +788,26 @@ class QueryConditionTest(DiskTestCase):
                 "TileDBError: Unhandled dot operator in Attribute(value=Name(id='attr', ctx=Load()), attr='two', ctx=Load()) -- if your attribute name has a dot in it, e.g. `orig.ident`, please wrap it with `attr(\"...\")`, e.g. `attr(\"orig.ident\")`"
                 in str(exc_info.value)
             )
+            with pytest.raises(tiledb.TileDBError) as exc_info:
+                A.query(cond="attr.one in [1, 2, 3]")[:]
+            assert (
+                "TileDBError: cannot handle query condition left-hand side of type"
+                in str(exc_info.value)
+            )
 
             # now test with the correct syntax
             result = A.query(cond='attr("attr.one") < 6')[:]
             assert_array_equal(result["attr.one"], A[:6]["attr.one"])
+            result = A.query(cond='attr("at.tr.thr.ee") >= 6')[:]
+            assert_array_equal(result["at.tr.thr.ee"], A[6:]["at.tr.thr.ee"])
             result = A.query(cond='attr("attr.two") >= 6')[:]
             assert_array_equal(result["attr.two"], A[6:]["attr.two"])
+            result = A.query(cond='attr("at.tr.thr.ee") >= 6')[:]
+            assert_array_equal(result["at.tr.thr.ee"], A[6:]["at.tr.thr.ee"])
+            result = A.query(cond='attr("attr.one") in [1, 2, 3]')[:]
+            assert_array_equal(result["attr.one"], A[1:4]["attr.one"])
+            result = A.query(cond='attr("at.tr.thr.ee") in [1, 2, 3]')[:]
+            assert_array_equal(result["at.tr.thr.ee"], A[1:4]["at.tr.thr.ee"])
 
     @pytest.mark.skipif(not has_pandas(), reason="pandas>=1.0,<3.0 not installed")
     def test_do_not_return_attrs(self):