Repository: lincc-frameworks/nested-pandas

Issue:
Fix NestedFrame.query/eval to handle column names including a space
We recently introduced an AST parser for evaluation of expressions against Nested-Pandas columns in `eval` and `query`.

However this seems to fail against column names which include a space (which is allowed in base Pandas).

This was discovered when updating the `nested-dask` version for `lsdb` since one of their unit tests exercises a query against such a column. Here is the example test failure with a full stack trace when our current query/eval parser is used by LSDB: https://github.com/astronomy-commons/lsdb/actions/runs/11713280862/job/32625751711?pr=493

A pretty trivial reproduction:

```
from nested_pandas import NestedFrame
import pandas as pd

nf = NestedFrame(data={"dog": [1, 2, 3], "good dog": [2, 4, 6]}, index=[0, 1, 2])

nested = pd.DataFrame(
    data={"a": [0, 2, 4, 1, 4, 3, 1, 4, 1], "b": [5, 4, 7, 5, 3, 1, 9, 3, 4]},
    index=[0, 0, 0, 1, 1, 1, 2, 2, 2],
)

nf = nf.add_nested(nested, "nested")
nf.query("`good dog` > 3")
```

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/pull/175?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.40%. Comparing base [(`402ab66`)](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/commit/402ab66e60db290b0c869387fcf68dd818b4425f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks) to head [(`3a5852f`)](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/commit/3a5852fbd14d0b3fd585e4b17c63166e498ab135?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #175   +/-   ##
=======================================
  Coverage   99.40%   99.40%           
=======================================
  Files          14       13    -1     
  Lines        1004     1016   +12     
=======================================
+ Hits          998     1010   +12     
  Misses          6        6           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/pull/175?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks).

github-actions[bot]: | Before [402ab66e] <v0.3.1>   | After [bd9afc5e]    | Ratio   | Benchmark (Parameter)                               |
|------------------------------|---------------------|---------|-----------------------------------------------------|
| 38.6±9ms                     | 48.5±10ms           | ~1.26   | benchmarks.ReassignHalfOfNestedSeries.time_run      |
| 7.91±0.1ms                   | 15.4±0.2ms          | 1.95    | benchmarks.NestedFrameQuery.time_run                |
| 22.1±0.1ms                   | 23.3±0.4ms          | 1.05    | benchmarks.AssignSingleDfToNestedSeries.time_run    |
| 9.73±0.2ms                   | 9.80±0.2ms          | 1.01    | benchmarks.NestedFrameAddNested.time_run            |
| 266M                         | 266M                | 1.00    | benchmarks.AssignSingleDfToNestedSeries.peakmem_run |
| 88.1M                        | 88M                 | 1.00    | benchmarks.NestedFrameAddNested.peakmem_run         |
| 93.2M                        | 93.2M               | 1.00    | benchmarks.NestedFrameQuery.peakmem_run             |
| 91.7M                        | 91.6M               | 1.00    | benchmarks.NestedFrameReduce.peakmem_run            |
| 1.59±0.02ms                  | 1.59±0.02ms         | 1.00    | benchmarks.NestedFrameReduce.time_run               |
| 285M                         | 285M                | 1.00    | benchmarks.ReassignHalfOfNestedSeries.peakmem_run   |


Click [here](https://github.com/lincc-frameworks/nested-pandas/actions/runs/11828211871/job/32957877028#step:10:1) to view all benchmarks.

Patch:
diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py
index 192a421..26e0467 100644
--- a/src/nested_pandas/nestedframe/core.py
+++ b/src/nested_pandas/nestedframe/core.py
@@ -10,9 +10,11 @@ import pyarrow as pa
 from pandas._libs import lib
 from pandas._typing import Any, AnyAll, Axis, IndexLabel
 from pandas.api.extensions import no_default
+from pandas.core.computation import ops
+from pandas.core.computation.eval import Expr, ensure_scope
 from pandas.core.computation.expr import PARSERS, PandasExprVisitor
+from pandas.core.computation.parsing import clean_column_name
 
-from nested_pandas.nestedframe.utils import extract_nest_names
 from nested_pandas.series.dtype import NestedDtype
 from nested_pandas.series.packer import pack, pack_lists, pack_sorted_df_into_struct
 
@@ -79,6 +81,22 @@ class _NestResolver(dict):
     def __init__(self, outer: NestedFrame):
         self._outer = outer
         super().__init__()
+        # Pre-load the field resolvers for all columns which are known at present.
+        for column in outer.nested_columns:
+            self._initialize_field_resolver(column, outer)
+
+    def _initialize_field_resolver(self, column: str, outer: NestedFrame):
+        """
+        Initialize a resolver for the given nested column, and also an alias
+        for it, in the case of column names that have spaces or are otherwise
+        not identifier-like.
+        """
+        super().__setitem__(column, _NestedFieldResolver(column, outer))
+        clean_id = clean_column_name(column)
+        # And once more for the cleaned name, if it's different.
+        # This allows us to capture references to it from the Pandas evaluator.
+        if clean_id != column:
+            super().__setitem__(clean_id, _NestedFieldResolver(column, outer))
 
     def __contains__(self, item):
         top_nest = item if "." not in item else item.split(".")[0].strip()
@@ -89,7 +107,7 @@ class _NestResolver(dict):
         if not super().__contains__(top_nest):
             if top_nest not in self._outer.nested_columns:
                 raise KeyError(f"Unknown nest {top_nest}")
-            super().__setitem__(top_nest, _NestedFieldResolver(top_nest, self._outer))
+            self._initialize_field_resolver(top_nest, self._outer)
         return super().__getitem__(top_nest)
 
     def __setitem__(self, item, _):
@@ -133,6 +151,48 @@ class _NestedFieldResolver:
         raise AttributeError(f"No attribute {item_name}")
 
 
+def _subexprs_by_nest(parents: list, node) -> 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 isinstance(node, ops.Term) and not isinstance(node, ops.Constant):
+        if isinstance(node.value, _SeriesFromNest):
+            return {node.value.nest_name: parents}
+        return {getattr(node, "upper_name", ""): parents}
+    if not isinstance(node, ops.Op):
+        return {}
+    sources = [getattr(node, "lhs", None), getattr(node, "rhs", None)]
+    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
+
+
 class NestedFrame(pd.DataFrame):
     """A Pandas Dataframe extension with support for nested structure.
 
@@ -457,6 +517,39 @@ class NestedFrame(pd.DataFrame):
         kwargs["parser"] = "nested-pandas"
         return super().eval(expr, **kwargs)
 
+    def extract_nest_names(
+        self,
+        expr: str,
+        local_dict=None,
+        global_dict=None,
+        resolvers=(),
+        level: int = 0,
+        target=None,
+        **kwargs,
+    ) -> set[str]:
+        """
+        Given a string expression, parse it and visit the resulting expression tree,
+        surfacing the nesting types.  The purpose is to identify expressions that attempt
+        to mix base and nested columns, or columns from two different nests.
+        """
+        index_resolvers = self._get_index_resolvers()
+        column_resolvers = self._get_cleaned_column_resolvers()
+        resolvers = resolvers + (_NestResolver(self), column_resolvers, index_resolvers)
+        # Parser needs to be the "nested-pandas" parser.
+        # We also need the same variable context that eval() will have, so that
+        # backtick-quoted names are substituted as expected.
+        env = ensure_scope(
+            level + 1,
+            global_dict=global_dict,
+            local_dict=local_dict,
+            resolvers=resolvers,
+            target=target,
+        )
+        parsed_expr = Expr(expr, parser="nested-pandas", env=env)
+        expr_tree = parsed_expr.terms
+        separable = _subexprs_by_nest([], expr_tree)
+        return set(separable.keys())
+
     def query(self, expr: str, *, inplace: bool = False, **kwargs) -> NestedFrame | None:
         """
         Query the columns of a NestedFrame with a boolean expression. Specified
@@ -514,7 +607,7 @@ class NestedFrame(pd.DataFrame):
         # 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.
-        nest_names = extract_nest_names(expr)
+        nest_names = self.extract_nest_names(expr, **kwargs)
         if len(nest_names) > 1:
             raise ValueError("Queries cannot target multiple structs/layers, write a separate query for each")
         result = self.eval(expr, **kwargs)
diff --git a/src/nested_pandas/nestedframe/utils.py b/src/nested_pandas/nestedframe/utils.py
deleted file mode 100644
index 4e56b5b..0000000
--- a/src/nested_pandas/nestedframe/utils.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# typing.Self and "|" union syntax don't exist in Python 3.9
-from __future__ import annotations
-
-import ast
-
-
-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 {}
-    if isinstance(node, ast.Name):
-        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: 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 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, or columns from two different nests.
-    """
-    expr_tree = ast.parse(expr, mode="eval").body
-    separable = _subexprs_by_nest([], expr_tree)
-    return set(separable.keys())


Test patch:
diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py
index 938476f..252386e 100644
--- a/tests/nested_pandas/nestedframe/test_nestedframe.py
+++ b/tests/nested_pandas/nestedframe/test_nestedframe.py
@@ -594,6 +594,25 @@ def test_query():
     assert base["nested.d"].shape == (2,)
 
 
+def test_query_on_non_identifier_columns():
+    """
+    Column names very often follow the same rules as Python identifiers, but
+    they are not required to.  Test that query() can handle such names.
+    """
+    # Taken from GH#174
+    nf = NestedFrame(data={"dog": [1, 2, 3], "good dog": [2, 4, 6]}, index=[0, 1, 2])
+    nested = pd.DataFrame(
+        data={"a": [0, 2, 4, 1, 4, 3, 1, 4, 1], "b": [5, 4, 7, 5, 3, 1, 9, 3, 4]},
+        index=[0, 0, 0, 1, 1, 1, 2, 2, 2],
+    )
+    nf = nf.add_nested(nested, "bad dog")
+    nf2 = nf.query("`good dog` > 3")
+    assert nf.shape == (3, 3)
+    assert nf2.shape == (2, 3)
+    nf3 = nf.query("`bad dog`.a > 2")
+    assert nf3["bad dog"].nest["a"].size == 4
+
+
 def test_dropna():
     """Test that dropna works on all layers"""
 
diff --git a/tests/nested_pandas/utils/test_utils.py b/tests/nested_pandas/utils/test_utils.py
index e2ec526..5c1abc9 100644
--- a/tests/nested_pandas/utils/test_utils.py
+++ b/tests/nested_pandas/utils/test_utils.py
@@ -2,7 +2,6 @@ import numpy as np
 import pandas as pd
 import pytest
 from nested_pandas import NestedFrame
-from nested_pandas.nestedframe.utils import extract_nest_names
 from nested_pandas.utils import count_nested
 
 
@@ -52,16 +51,41 @@ def test_check_expr_nesting():
     used to ensure that an expression-based query does not try to combine base and nested
     sub-expressions.
     """
-    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", ""}
+    base = NestedFrame(data={"a": [1, 2, 3], "b": [2, np.nan, 6]}, index=[0, 1, 2])
+    nested = pd.DataFrame(
+        data={
+            "c": [0, 2, 4, 1, np.nan, 3, 1, 4, 1],
+            "d": [5, 4, 7, 5, 3, 1, 9, 3, 4],
+            "label": ["b", "a", "b", "b", "a", "a", "b", "a", "b"],
+        },
+        index=[0, 0, 0, 1, 1, 1, 2, 2, 2],
+    )
+    b1 = base.add_nested(nested, "nested")
+    assert b1.extract_nest_names("a > 2 & nested.c > 1") == {"", "nested"}
+    assert b1.extract_nest_names("(nested.c > 1) and (nested.d>2)") == {"nested"}
+    assert b1.extract_nest_names("-1.52e-5 < b < 35.2e2") == {""}
+
+    b2 = base.add_nested(nested.copy(), "n")
+    assert b2.extract_nest_names("(n.c > 1) and ((b + a) > (b - 1e-8)) or n.d > a") == {"n", ""}
+
+    abc = pd.DataFrame(
+        data={
+            "c": [3, 1, 4, 1, 5, 9, 2, 6, 5],
+            "d": [1, 4, 1, 2, 1, 3, 5, 6, 2],
+            "g": ["a", "b", "c", "d", "e", "f", "g", "h", "i"],
+        },
+        index=[0, 0, 0, 1, 1, 1, 2, 2, 2],
+    )
+    b3 = base.add_nested(abc, "abc").add_nested(abc, "c")
+    assert b3.extract_nest_names("abc.c > 2 & c.d < 5") == {"abc", "c"}
+
+    assert b3.extract_nest_names("(abc.d > 3) & (abc.c == [2, 5])") == {"abc"}
+    assert b3.extract_nest_names("(abc.d > 3)&(abc.g == 'f')") == {"abc"}
+    assert b3.extract_nest_names("(abc.d > 3) & (abc.g == 'f')") == {"abc"}
 
-    assert extract_nest_names("a.b > 2 & c.d < 5") == {"a", "c"}
+    assert b1.extract_nest_names("a>3") == {""}
+    assert b1.extract_nest_names("a > 3") == {""}
 
-    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"}
+    b4 = base.add_nested(nested, "test")
+    assert b4.extract_nest_names("test.c>5&b==2") == {"test", ""}
+    assert b4.extract_nest_names("test.c > 5 & b == 2") == {"test", ""}