From 8ae6f9471a74be0e53a1ace1762141ee6a27c103 Mon Sep 17 00:00:00 2001 From: Pavel Zwerschke Date: Tue, 8 Oct 2024 12:00:11 +0200 Subject: [PATCH 1/3] fix: Invalid selectors not being recognized --- .../src/plans/conversion/expr_expansion.rs | 24 ++++++++++--------- py-polars/tests/unit/test_selectors.py | 8 +++++++ 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/crates/polars-plan/src/plans/conversion/expr_expansion.rs b/crates/polars-plan/src/plans/conversion/expr_expansion.rs index 4d1aa76caff5..72a04263268e 100644 --- a/crates/polars-plan/src/plans/conversion/expr_expansion.rs +++ b/crates/polars-plan/src/plans/conversion/expr_expansion.rs @@ -921,22 +921,24 @@ pub(super) fn expand_selector( let mut members = PlIndexSet::new(); replace_selector_inner(s, &mut members, &mut vec![], schema, keys)?; - if members.len() <= 1 { - members - .into_iter() - .map(|e| { - let Expr::Column(name) = e else { - polars_bail!(InvalidOperation: "invalid selector expression: {}", e) - }; - Ok(name) - }) - .collect() + let column_names = members + .into_iter() + .map(|e| { + let Expr::Column(name) = e else { + polars_bail!(InvalidOperation: "invalid selector expression: {}", e) + }; + Ok(name) + }) + .collect::>>()?; + + if column_names.len() <= 1 { + Ok(column_names) } else { // Ensure that multiple columns returned from combined/nested selectors remain in schema order let selected = schema .iter_fields() .map(|field| field.name().clone()) - .filter(|field_name| members.contains(&Expr::Column(field_name.clone()))) + .filter(|field_name| column_names.contains(&field_name)) .collect(); Ok(selected) diff --git a/py-polars/tests/unit/test_selectors.py b/py-polars/tests/unit/test_selectors.py index bf44ff87bac5..72d3e91cc081 100644 --- a/py-polars/tests/unit/test_selectors.py +++ b/py-polars/tests/unit/test_selectors.py @@ -814,3 +814,11 @@ def test_selector_list_of_lists_18499() -> None: with pytest.raises(InvalidOperationError, match="invalid selector expression"): lf.unique(subset=[["bar", "ham"]]) # type: ignore[list-item] + + +def test_invalid_selector() -> None: + df = pl.DataFrame(data={"x": [1, 2], "z": ["a", "b"]}) + with pytest.raises(InvalidOperationError, match="invalid selector expression"): + df.drop(pl.col("x", "z") + 2) + with pytest.raises(InvalidOperationError, match="invalid selector expression"): + df.drop(pl.col("x") + 2) From 027f115dc2afaf5cabdacbd93bfa4aec69df5c05 Mon Sep 17 00:00:00 2001 From: Pavel Zwerschke Date: Tue, 8 Oct 2024 12:05:43 +0200 Subject: [PATCH 2/3] fix --- crates/polars-plan/src/plans/conversion/expr_expansion.rs | 2 +- py-polars/tests/unit/test_selectors.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/polars-plan/src/plans/conversion/expr_expansion.rs b/crates/polars-plan/src/plans/conversion/expr_expansion.rs index 72a04263268e..d04adcd932e8 100644 --- a/crates/polars-plan/src/plans/conversion/expr_expansion.rs +++ b/crates/polars-plan/src/plans/conversion/expr_expansion.rs @@ -938,7 +938,7 @@ pub(super) fn expand_selector( let selected = schema .iter_fields() .map(|field| field.name().clone()) - .filter(|field_name| column_names.contains(&field_name)) + .filter(|field_name| column_names.contains(field_name)) .collect(); Ok(selected) diff --git a/py-polars/tests/unit/test_selectors.py b/py-polars/tests/unit/test_selectors.py index 72d3e91cc081..03fe83f0e3b0 100644 --- a/py-polars/tests/unit/test_selectors.py +++ b/py-polars/tests/unit/test_selectors.py @@ -819,6 +819,6 @@ def test_selector_list_of_lists_18499() -> None: def test_invalid_selector() -> None: df = pl.DataFrame(data={"x": [1, 2], "z": ["a", "b"]}) with pytest.raises(InvalidOperationError, match="invalid selector expression"): - df.drop(pl.col("x", "z") + 2) + df.drop(pl.col("x", "z") + 2) # type: ignore[arg-type] with pytest.raises(InvalidOperationError, match="invalid selector expression"): - df.drop(pl.col("x") + 2) + df.drop(pl.col("x") + 2) # type: ignore[arg-type] From 3947452bc1061030e34cc59a81de830213a162d1 Mon Sep 17 00:00:00 2001 From: Pavel Zwerschke Date: Mon, 14 Oct 2024 15:09:43 +0200 Subject: [PATCH 3/3] update test --- py-polars/tests/unit/test_selectors.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/py-polars/tests/unit/test_selectors.py b/py-polars/tests/unit/test_selectors.py index 03fe83f0e3b0..9cae617ae087 100644 --- a/py-polars/tests/unit/test_selectors.py +++ b/py-polars/tests/unit/test_selectors.py @@ -818,7 +818,7 @@ def test_selector_list_of_lists_18499() -> None: def test_invalid_selector() -> None: df = pl.DataFrame(data={"x": [1, 2], "z": ["a", "b"]}) + sel = pl.selectors.all() + pl.col("x") + assert isinstance(sel, pl.selectors._selector_proxy_) with pytest.raises(InvalidOperationError, match="invalid selector expression"): - df.drop(pl.col("x", "z") + 2) # type: ignore[arg-type] - with pytest.raises(InvalidOperationError, match="invalid selector expression"): - df.drop(pl.col("x") + 2) # type: ignore[arg-type] + df.drop(sel)