Repository: lincc-frameworks/nested-pandas

Issue:
Cross-nest queries in NestedFrame.query should be prohibited or split
**Bug report**


**Before submitting**
Please check the following:

- [X] I have described the situation in which the bug arose, including what code was executed, information about my environment, and any applicable data others will need to reproduce the problem.
- [X] I have included available evidence of the unexpected behavior (including error messages, screenshots, and/or plots) as well as a description of what I expected instead.
- [X] If I have a solution in mind, I have provided an explanation and/or pseudocode and/or task list.

Processing queries across base and nested columns is not possible in `NestedFrame.query` today, so expressions of this kind are detected and prohibited.  However, it does not presently prohibit queries that go across two different nests:

```python
nf.query("n1.c < 0.5 and n2.d > 1e-7")
```

The above will "work", but only selects the nested rows in `n1`, simply because `n1` is earlier in the string.

The solution is to treat such cross-nest queries exactly the same as queries that cross base and nested columns: error out today, and as split queries, in the future.

Hints from issue comments:

Solution patch:
diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py
index 08907321..8736c196 100644
--- a/src/nested_pandas/nestedframe/core.py
+++ b/src/nested_pandas/nestedframe/core.py
@@ -16,7 +16,7 @@
 from nested_pandas.series.dtype import NestedDtype
 
 from ..series.packer import pack_sorted_df_into_struct
-from .utils import NestingType, check_expr_nesting
+from .utils import extract_nest_names
 
 
 class NestedPandasExprVisitor(PandasExprVisitor):
@@ -506,11 +506,11 @@ def query(self, expr: str, *, inplace: bool = False, **kwargs) -> NestedFrame |
             raise ValueError(msg)
         kwargs["level"] = kwargs.pop("level", 0) + 1
         kwargs["target"] = None
-        # At present, the query expression must be either entirely within the
-        # nested namespace or the base namespace.  Mixed structures are not
+        # At present, the query expression must be either entirely within a
+        # single nest, or have nothing but base columns.  Mixed structures are not
         # supported, so preflight the expression.
-        nesting_types = check_expr_nesting(expr)
-        if NestingType.NESTED in nesting_types and NestingType.BASE in nesting_types:
+        nest_names = extract_nest_names(expr)
+        if len(nest_names) > 1:
             raise ValueError("Queries cannot target multiple structs/layers, write a separate query for each")
         result = self.eval(expr, **kwargs)
         # If the result is a _SeriesFromNest, then the evaluation has caused unpacking,
diff --git a/src/nested_pandas/nestedframe/utils.py b/src/nested_pandas/nestedframe/utils.py
index b219cea1..4e56b5b4 100644
--- a/src/nested_pandas/nestedframe/utils.py
+++ b/src/nested_pandas/nestedframe/utils.py
@@ -2,39 +2,60 @@
 from __future__ import annotations
 
 import ast
-from enum import Enum
 
 
-class NestingType(Enum):
-    """Types of sub-expressions possible in a NestedFrame string expression."""
-
-    BASE = "base"
-    NESTED = "nested"
-
-
-def _expr_nesting_type(node: ast.expr | None) -> set[NestingType]:
+def _subexprs_by_nest(parents: list[ast.expr], node: ast.expr | None) -> dict[str, list]:
+    """
+    Given an expression which contains references to both base and nested
+    columns, return a dictionary of the sub-expressions that should be
+    evaluated independently, keyed by nesting context.
+
+    The key of the dictionary is the name of the nested column, and will
+    be a blank string in the case of base columns.  The value is a list
+    of the parent nodes that lead to sub-expressions that can be evaluated
+    successfully.
+
+    While this is not in use today for automatically splitting expressions,
+    it can be used to detect whether an expression is suitably structured
+    for evaluation: the returned dictionary should have a single key.
+    """
     if not isinstance(node, ast.expr):
-        return set()
+        return {}
     if isinstance(node, ast.Name):
-        return {NestingType.BASE}
-    if isinstance(node, ast.Attribute):
-        return {NestingType.NESTED}
+        return {"": parents}
+    if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
+        return {node.value.id: parents}
     sources = (
         [getattr(node, "left", None), getattr(node, "right", None)]
         + getattr(node, "values", [])
         + getattr(node, "comparators", [])
     )
-    result: set[NestingType] = set()
-    for s in sources:
-        result.update(_expr_nesting_type(s))
+    result: dict[str, list] = {}
+    for source in sources:
+        child = _subexprs_by_nest(parents, source)
+        for k, v in child.items():
+            result.setdefault(k, []).append(v)
+    # After a complete traversal across sources, check for any necessary splits.
+    # If it's homogenous, move the split-node up the tree.
+    if len(result) == 1:
+        # Let the record of each parent node drift up the tree,
+        # and merge the subtrees into a single node, since by definition,
+        # this node is homogeneous over all of its children, and can
+        # be evaluated in a single step.
+        result = {k: [node] for k in result}
+    # If the result is either empty or has more than one key, leave the result
+    # alone.  Each key represents a different nest (with a blank string for the base),
+    # and the value is the highest point in the expression tree where the expression
+    # was still within a single nest.
     return result
 
 
-def check_expr_nesting(expr: str) -> set[NestingType]:
+def extract_nest_names(expr: str) -> set[str]:
     """
     Given a string expression, parse it and visit the resulting AST, surfacing
     the nesting types.  The purpose is to identify expressions that attempt
-    to mix base and nested columns, which will need to be handled specially.
+    to mix base and nested columns, or columns from two different nests.
     """
     expr_tree = ast.parse(expr, mode="eval").body
-    return _expr_nesting_type(expr_tree)
+    separable = _subexprs_by_nest([], expr_tree)
+    return set(separable.keys())
diff --git a/src/nested_pandas/series/ext_array.py b/src/nested_pandas/series/ext_array.py
index 2c5a839f..6f98816a 100644
--- a/src/nested_pandas/series/ext_array.py
+++ b/src/nested_pandas/series/ext_array.py
@@ -204,7 +204,7 @@ def __setitem__(self, key, value) -> None:
             # Copy will happen later in replace_with_mask() anyway
             value = self._box_pa_array(value, pa_type=self._pyarrow_dtype)
         else:
-            # Our replace_with_mask implementation doesm't work with scalars
+            # Our replace_with_mask implementation doesn't work with scalars
             value = pa.array([scalar] * pa.compute.sum(pa_mask).as_py())
 
         if argsort is not None:
@@ -387,7 +387,7 @@ def copy(self) -> Self:  # type: ignore[name-defined] # noqa: F821
         This implementation returns a shallow copy of the extension array,
         because the underlying PyArrow array is immutable.
         """
-        return type(self)(self._chunked_array)
+        return type(self)(self._chunked_array, validate=False)
 
     def _formatter(self, boxed: bool = False) -> Callable[[Any], str | None]:
         # TODO: make formatted strings more pretty


Test patch:
diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py
index fa5da349..151b5f09 100644
--- a/tests/nested_pandas/nestedframe/test_nestedframe.py
+++ b/tests/nested_pandas/nestedframe/test_nestedframe.py
@@ -447,6 +447,11 @@ def test_query():
     # Check for the multi-layer error
     with pytest.raises(ValueError):
         base.query("a > 2 & nested.c > 1")
+    # Create another nest in order to further test the multi-layer error
+    base_2 = base.eval("nest2.c = nested.c + 1")
+    assert len(base_2.nested_columns) == 2
+    with pytest.raises(ValueError):
+        base_2.query("nested.c > 1 & nest2.c > 2")
 
     # Test nested queries
     nest_queried = base.query("nested.c > 1")
diff --git a/tests/nested_pandas/utils/test_utils.py b/tests/nested_pandas/utils/test_utils.py
index d76d5928..e2ec5266 100644
--- a/tests/nested_pandas/utils/test_utils.py
+++ b/tests/nested_pandas/utils/test_utils.py
@@ -2,7 +2,7 @@
 import pandas as pd
 import pytest
 from nested_pandas import NestedFrame
-from nested_pandas.nestedframe.utils import NestingType, check_expr_nesting
+from nested_pandas.nestedframe.utils import extract_nest_names
 from nested_pandas.utils import count_nested
 
 
@@ -52,21 +52,16 @@ def test_check_expr_nesting():
     used to ensure that an expression-based query does not try to combine base and nested
     sub-expressions.
     """
-    assert check_expr_nesting("a > 2 & nested.c > 1") == {NestingType.NESTED, NestingType.BASE}
-    assert check_expr_nesting("(nested.c > 1) and (nested.d>2)") == {NestingType.NESTED}
-    assert check_expr_nesting("-1.52e-5 < abc < 35.2e2") == {NestingType.BASE}
-    assert check_expr_nesting("(n.a > 1) and ((b + c) > (d - 1e-8)) or n.q > c") == {
-        NestingType.NESTED,
-        NestingType.BASE,
-    }
+    assert extract_nest_names("a > 2 & nested.c > 1") == {"", "nested"}
+    assert extract_nest_names("(nested.c > 1) and (nested.d>2)") == {"nested"}
+    assert extract_nest_names("-1.52e-5 < abc < 35.2e2") == {""}
+    assert extract_nest_names("(n.a > 1) and ((b + c) > (d - 1e-8)) or n.q > c") == {"n", ""}
 
-    # NOTE: this correctly captures the desired behavior here, but suggests that the two nests
-    # are interoperable, which is too strong a claim.
-    assert check_expr_nesting("a.b > 2 & c.d < 5") == {NestingType.NESTED}
+    assert extract_nest_names("a.b > 2 & c.d < 5") == {"a", "c"}
 
-    assert check_expr_nesting("a>3") == {NestingType.BASE}
-    assert check_expr_nesting("a > 3") == {NestingType.BASE}
-    assert check_expr_nesting("test.a>5&b==2") == {NestingType.NESTED, NestingType.BASE}
-    assert check_expr_nesting("test.a > 5 & b == 2") == {NestingType.NESTED, NestingType.BASE}
-    assert check_expr_nesting("(a.b > 3)&(a.c == 'f')") == {NestingType.NESTED}
-    assert check_expr_nesting("(a.b > 3) & (a.c == 'f')") == {NestingType.NESTED}
+    assert extract_nest_names("a>3") == {""}
+    assert extract_nest_names("a > 3") == {""}
+    assert extract_nest_names("test.a>5&b==2") == {"test", ""}
+    assert extract_nest_names("test.a > 5 & b == 2") == {"test", ""}
+    assert extract_nest_names("(a.b > 3)&(a.c == 'f')") == {"a"}
+    assert extract_nest_names("(a.b > 3) & (a.c == 'f')") == {"a"}