diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs b/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs index 61fe7b611260..8733215f563e 100644 --- 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, schema: &Schema, exclude: &PlHashSet>, + 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 index d40201c0bf8f..aa218fac5872 100644 --- 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"]