fix(rust, python): fix regression in regex expansion (pola-rs#9952)

alexander-beedie · Jul 18, 2023 · 760a4f0 · 760a4f0
1 parent 45c0d0e
commit 760a4f0
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs b/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs
@@ -131,6 +131,7 @@ fn replace_regex(
     result: &mut Vec<Expr>,
     schema: &Schema,
     exclude: &PlHashSet<Arc<str>>,
+    has_exclude: bool,
 ) -> PolarsResult<()> {
     let roots = expr_to_leaf_column_names(expr);
     let mut regex = None;
@@ -139,18 +140,19 @@ fn replace_regex(
             match regex {
                 None => {
                     regex = Some(name);
-                    if exclude.is_empty() {
-                        expand_regex(expr, result, schema, name, exclude)?
-                    } else {
+                    if has_exclude {
                         // iterate until we find the Exclude node
                         // we remove that node from the expression
+                        // the `exclude` set is already filled
+                        // by prepare exclude
                         for e in expr.into_iter() {
                             if let Expr::Exclude(e, _) = e {
                                 expand_regex(e, result, schema, name, exclude)?;
-                                break;
+                                return Ok(());
                             }
                         }
                     }
+                    expand_regex(expr, result, schema, name, exclude)?
                 }
                 Some(r) => {
                     polars_ensure!(
@@ -274,7 +276,13 @@ fn prepare_excluded(
                         match to_exclude_single {
                             Excluded::Name(name) => {
                                 let e = Expr::Column(name.clone());
-                                replace_regex(&e, &mut buf, schema, &Default::default())?;
+                                replace_regex(
+                                    &e,
+                                    &mut buf,
+                                    schema,
+                                    &Default::default(),
+                                    has_exclude,
+                                )?;
                                 // we cannot loop because of bchck
                                 while let Some(col) = buf.pop() {
                                     if let Expr::Column(name) = col {
@@ -512,7 +520,7 @@ fn replace_and_add_to_results(
         {
             // keep track of column excluded from the dtypes
             let exclude = prepare_excluded(&expr, schema, keys, flags.has_exclude)?;
-            replace_regex(&expr, result, schema, &exclude)?;
+            replace_regex(&expr, result, schema, &exclude, flags.has_exclude)?;
         }
         #[cfg(not(feature = "regex"))]
         {

diff --git a/py-polars/tests/unit/test_selectors.py b/py-polars/tests/unit/test_selectors.py
@@ -427,3 +427,8 @@ def test_selector_expr_dispatch() -> None:
                 pl.when(nan_or_inf).then(0.0).otherwise(cs.float()).keep_name()
             ).fill_null(0),
         )
+
+
+def test_regex_expansion_groupby_9947() -> None:
+    df = pl.DataFrame({"g": [3], "abc": [1], "abcd": [3]})
+    assert df.groupby("g").agg(pl.col("^ab.*$")).columns == ["g", "abc", "abcd"]