fix(rust, python): Don't reuse member for Selector::Add (pola-rs#11026)

cmdlineluser · Sep 11, 2023 · 8d5de95 · 8d5de95
1 parent 2433c58
commit 8d5de95
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/crates/polars-plan/src/logical_plan/projection.rs b/crates/polars-plan/src/logical_plan/projection.rs
@@ -580,7 +580,9 @@ fn replace_selector_inner(
         },
         Selector::Add(lhs, rhs) => {
             replace_selector_inner(*lhs, members, scratch, schema, keys)?;
-            replace_selector_inner(*rhs, members, scratch, schema, keys)?;
+            let mut rhs_members: PlIndexSet<Expr> = Default::default();
+            replace_selector_inner(*rhs, &mut rhs_members, scratch, schema, keys)?;
+            members.extend(rhs_members)
         },
         Selector::Sub(lhs, rhs) => {
             // fill lhs

diff --git a/py-polars/tests/unit/test_selectors.py b/py-polars/tests/unit/test_selectors.py
@@ -483,3 +483,16 @@ def test_regex_expansion_exclude_10002() -> None:
         ).to_dict(as_series=False)
         == expected
     )
+
+
+def test_selector_or() -> None:
+    df = pl.DataFrame(
+        {
+            "int": [1, 2, 3],
+            "float": [1.0, 2.0, 3.0],
+            "str": ["x", "y", "z"],
+        }
+    ).with_row_count("rn")
+
+    out = df.select(cs.by_name("rn") | ~cs.numeric())
+    assert out.to_dict(False) == {"rn": [0, 1, 2], "str": ["x", "y", "z"]}