From e09ac7275b5fe71435850af00e7dc9b2da1ea5ab Mon Sep 17 00:00:00 2001 From: Sai Asish Y Date: Mon, 5 Oct 2026 12:06:05 -0700 Subject: [PATCH] fix: return selector union columns in schema order Signed-off-by: Sai Asish Y --- src/narwhals/_compliant/selectors.py | 16 +++++----------- tests/selectors_test.py | 3 ++- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/src/narwhals/_compliant/selectors.py b/src/narwhals/_compliant/selectors.py index 1a3ffcd24e..3d68ac21a9 100644 --- a/src/narwhals/_compliant/selectors.py +++ b/src/narwhals/_compliant/selectors.py @@ -279,20 +279,14 @@ def __or__( def series(df: FrameT) -> Sequence[SeriesOrExprT]: lhs_names, rhs_names = _eval_lhs_rhs(df, self, other) - rhs_set = frozenset(rhs_names) - return [ - *( - x - for x, name in zip(self(df), lhs_names, strict=True) - if name not in rhs_set - ), - *other(df), - ] + by_name = dict(zip(lhs_names, self(df), strict=True)) + by_name.update(zip(rhs_names, other(df), strict=True)) + return [by_name[name] for name in df.columns if name in by_name] def names(df: FrameT) -> Sequence[str]: lhs_names, rhs_names = _eval_lhs_rhs(df, self, other) - rhs_set = frozenset(rhs_names) - return [*(x for x in lhs_names if x not in rhs_set), *rhs_names] + selected = {*lhs_names, *rhs_names} + return [name for name in df.columns if name in selected] return self.selectors._selector.from_callables(series, names, context=self) return self._to_expr() | other diff --git a/tests/selectors_test.py b/tests/selectors_test.py index 72fb3835a8..7b596c94a3 100644 --- a/tests/selectors_test.py +++ b/tests/selectors_test.py @@ -271,6 +271,7 @@ def test_datetime_no_tz(constructor: Constructor) -> None: ("selector", "expected"), [ (ncs.numeric() | ncs.boolean(), ["a", "c", "d"]), + (ncs.boolean() | ncs.numeric(), ["a", "c", "d"]), (ncs.numeric() & ncs.boolean(), []), (ncs.numeric() & ncs.by_dtype(nw.Int64), ["a"]), (ncs.numeric() | ncs.by_dtype(nw.Int64), ["a", "c"]), @@ -295,7 +296,7 @@ def test_set_ops( request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) result = df.select(selector).collect_schema().names() - assert sorted(result) == expected + assert result == expected def test_subtract_expr(constructor: Constructor) -> None: