From 9aa91a7176b96a143b592dbd0bc363c9ee5fb3d8 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Sun, 4 Aug 2024 11:28:13 +0200 Subject: [PATCH] fix(python): Fix drop selector (#18034) --- .../src/plans/conversion/expr_expansion.rs | 11 +++++++++-- py-polars/tests/unit/test_selectors.py | 3 +++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/crates/polars-plan/src/plans/conversion/expr_expansion.rs b/crates/polars-plan/src/plans/conversion/expr_expansion.rs index 9b5b532479f6..d069fc38ae85 100644 --- a/crates/polars-plan/src/plans/conversion/expr_expansion.rs +++ b/crates/polars-plan/src/plans/conversion/expr_expansion.rs @@ -854,11 +854,18 @@ pub(super) fn expand_selectors( ) -> PolarsResult> { let mut columns = vec![]; + // Skip the column fast paths. + fn skip(name: &str) -> bool { + is_regex_projection(name) || name == "*" + } + for s in s { match s { Selector::Root(e) => match *e { - Expr::Column(name) => columns.push(name), - Expr::Columns(names) => columns.extend_from_slice(names.as_ref()), + Expr::Column(name) if !skip(name.as_ref()) => columns.push(name), + Expr::Columns(names) if names.iter().all(|n| !skip(n.as_ref())) => { + columns.extend_from_slice(names.as_ref()) + }, Expr::Selector(s) => { let names = expand_selector(s, schema, keys)?; columns.extend_from_slice(names.as_ref()); diff --git a/py-polars/tests/unit/test_selectors.py b/py-polars/tests/unit/test_selectors.py index 2d2bc3239dcf..e292e709a970 100644 --- a/py-polars/tests/unit/test_selectors.py +++ b/py-polars/tests/unit/test_selectors.py @@ -347,6 +347,9 @@ def test_selector_drop(df: pl.DataFrame) -> None: dfd = df.drop(cs.numeric(), cs.temporal()) assert dfd.columns == ["eee", "fgg", "qqR"] + df = pl.DataFrame([["x"], [1]], schema={"foo": pl.String, "foo_right": pl.Int8}) + assert df.drop(cs.ends_with("_right")).schema == {"foo": pl.String()} + def test_selector_duration(df: pl.DataFrame) -> None: assert df.select(cs.duration("ms")).columns == []