From 58c75953d44f2004515adcb4025acd6d3e79ab42 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 14:34:55 +0000 Subject: [PATCH 01/21] feat: Add `DataFrame.explode` --- narwhals/_plan/arrow/dataframe.py | 4 ++++ narwhals/_plan/compliant/dataframe.py | 1 + narwhals/_plan/dataframe.py | 9 +++++++++ 3 files changed, 14 insertions(+) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index 1c2f3e1d6c..465afdbf39 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -162,6 +162,10 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self: native = self.native.filter(~to_drop) return self._with_native(native) + def explode(self, subset: Sequence[str]) -> Self: + msg = "ArrowDataFrame.explode" + raise NotImplementedError(msg) + def rename(self, mapping: Mapping[str, str]) -> Self: names: dict[str, str] | list[str] if fn.BACKEND_VERSION >= (17,): diff --git a/narwhals/_plan/compliant/dataframe.py b/narwhals/_plan/compliant/dataframe.py index b502848534..b40019d788 100644 --- a/narwhals/_plan/compliant/dataframe.py +++ b/narwhals/_plan/compliant/dataframe.py @@ -59,6 +59,7 @@ def to_narwhals(self) -> BaseFrame[NativeFrameT_co]: ... def columns(self) -> list[str]: ... def drop(self, columns: Sequence[str]) -> Self: ... def drop_nulls(self, subset: Sequence[str] | None) -> Self: ... + def explode(self, subset: Sequence[str]) -> Self: ... # Shouldn't *need* to be `NamedIR`, but current impl depends on a name being passed around def filter(self, predicate: NamedIR, /) -> Self: ... def rename(self, mapping: Mapping[str, str]) -> Self: ... diff --git a/narwhals/_plan/dataframe.py b/narwhals/_plan/dataframe.py index 7455ad8cc5..1bdff8cf88 100644 --- a/narwhals/_plan/dataframe.py +++ b/narwhals/_plan/dataframe.py @@ -158,6 +158,15 @@ def with_row_index( by_names = expand_selector_irs_names(by_selectors, schema=self, require_any=True) return self._with_compliant(self._compliant.with_row_index_by(name, by_names)) + def explode( + self, + columns: OneOrIterable[ColumnNameOrSelector], + *more_columns: ColumnNameOrSelector, + ) -> Self: + s_ir = _parse.parse_into_combined_selector_ir(columns, *more_columns) + subset = expand_selector_irs_names((s_ir,), schema=self, require_any=True) + return self._with_compliant(self._compliant.explode(subset)) + def _dataframe_from_dict( data: Mapping[str, Any], From 2f4ffed9687e6645fa04a1a51f89dea82cb24a79 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 14:35:13 +0000 Subject: [PATCH 02/21] test: Port tests from `main` --- tests/plan/explode_test.py | 115 +++++++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 tests/plan/explode_test.py diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py new file mode 100644 index 0000000000..45bd737b02 --- /dev/null +++ b/tests/plan/explode_test.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +import narwhals as nw +import narwhals._plan as nwp +from narwhals.exceptions import InvalidOperationError, ShapeError +from tests.plan.utils import assert_equal_data, dataframe + +if TYPE_CHECKING: + from collections.abc import Sequence + + from tests.conftest import Data + + +@pytest.fixture(scope="module") +def data() -> Data: + # For context, polars allows to explode multiple columns only if the columns + # have matching element counts, therefore, l1 and l2 but not l1 and l3 together. + return { + "a": ["x", "y", "z", "w"], + "l1": [[1, 2], None, [None], []], + "l2": [[3, None], None, [42], []], + "l3": [[1, 2], [3], [None], [1]], + "l4": [[1, 2], [3], [123], [456]], + } + + +@pytest.mark.xfail( + reason="TODO:` DataFrame.explode` (single column)", raises=NotImplementedError +) +@pytest.mark.parametrize( + ("column", "expected_values"), + [("l2", [None, 3, None, None, 42]), ("l3", [1, 1, 2, 3, None])], +) +def test_explode_single_col( + column: str, expected_values: list[int | None], data: Data +) -> None: # pragma: no cover + result = ( + dataframe(data) + .with_columns(nwp.col(column).cast(nw.List(nw.Int32()))) + .explode(column) + .select("a", column) + .sort("a", column, nulls_last=True) + ) + expected = {"a": ["w", "x", "x", "y", "z"], column: expected_values} + assert_equal_data(result, expected) + + +@pytest.mark.xfail( + reason="TODO:` DataFrame.explode` (multi column)", raises=NotImplementedError +) +@pytest.mark.parametrize( + ("column", "more_columns", "expected"), + [ + ( + "l1", + ["l2"], + { + "a": ["w", "x", "x", "y", "z"], + "l1": [None, 1, 2, None, None], + "l2": [None, 3, None, None, 42], + }, + ), + ( + "l3", + ["l4"], + { + "a": ["w", "x", "x", "y", "z"], + "l3": [1, 1, 2, 3, None], + "l4": [456, 1, 2, 3, 123], + }, + ), + ], +) +def test_explode_multiple_cols( + column: str, + more_columns: Sequence[str], + expected: dict[str, list[str | int | None]], + data: Data, +) -> None: # pragma: no cover + result = ( + dataframe(data) + .with_columns(nwp.col(column, *more_columns).cast(nw.List(nw.Int32()))) + .explode(column, *more_columns) + .select("a", column, *more_columns) + .sort("a", column, nulls_last=True) + ) + assert_equal_data(result, expected) + + +@pytest.mark.xfail( + reason="TODO:` DataFrame.explode` (validate shape)", + raises=(AssertionError, NotImplementedError), +) +def test_explode_shape_error(data: Data) -> None: # pragma: no cover + with pytest.raises( + ShapeError, match=r".*exploded columns (must )?have matching element counts" + ): + dataframe(data).with_columns( + nwp.col("l1", "l2", "l3").cast(nw.List(nw.Int32())) + ).explode("l1", "l3") + + +@pytest.mark.xfail( + reason="TODO:` DataFrame.explode` (validate dtype)", + raises=(AssertionError, NotImplementedError), +) +def test_explode_invalid_operation_error(data: Data) -> None: # pragma: no cover + with pytest.raises( + InvalidOperationError, match="`explode` operation not supported for dtype" + ): + dataframe(data).explode("a") From 74125f376e14c47fab00fce79245220d28c5e108 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 14:49:13 +0000 Subject: [PATCH 03/21] feat: Validate dtypes at narwhals-level One `collect_schema()` is now shared for selector expansion & type checking --- narwhals/_plan/dataframe.py | 12 ++++++++++-- tests/plan/explode_test.py | 11 ++++------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/narwhals/_plan/dataframe.py b/narwhals/_plan/dataframe.py index 1bdff8cf88..c14343be09 100644 --- a/narwhals/_plan/dataframe.py +++ b/narwhals/_plan/dataframe.py @@ -25,7 +25,7 @@ ) from narwhals._utils import Implementation, Version, generate_repr from narwhals.dependencies import is_pyarrow_table -from narwhals.exceptions import ShapeError +from narwhals.exceptions import InvalidOperationError, ShapeError from narwhals.schema import Schema from narwhals.typing import EagerAllowed, IntoBackend, IntoDType, IntoSchema, JoinStrategy @@ -164,7 +164,15 @@ def explode( *more_columns: ColumnNameOrSelector, ) -> Self: s_ir = _parse.parse_into_combined_selector_ir(columns, *more_columns) - subset = expand_selector_irs_names((s_ir,), schema=self, require_any=True) + schema = self.collect_schema() + subset = expand_selector_irs_names((s_ir,), schema=schema, require_any=True) + dtypes = self.version.dtypes + tp_list = dtypes.List + for col_to_explode in subset: + dtype = schema[col_to_explode] + if dtype != tp_list: + msg = f"`explode` operation is not supported for dtype `{dtype}`, expected List type" + raise InvalidOperationError(msg) return self._with_compliant(self._compliant.explode(subset)) diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index 45bd737b02..3a5867299a 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -7,7 +7,7 @@ import narwhals as nw import narwhals._plan as nwp from narwhals.exceptions import InvalidOperationError, ShapeError -from tests.plan.utils import assert_equal_data, dataframe +from tests.plan.utils import assert_equal_data, dataframe, re_compile if TYPE_CHECKING: from collections.abc import Sequence @@ -104,12 +104,9 @@ def test_explode_shape_error(data: Data) -> None: # pragma: no cover ).explode("l1", "l3") -@pytest.mark.xfail( - reason="TODO:` DataFrame.explode` (validate dtype)", - raises=(AssertionError, NotImplementedError), -) -def test_explode_invalid_operation_error(data: Data) -> None: # pragma: no cover +def test_explode_invalid_operation_error(data: Data) -> None: with pytest.raises( - InvalidOperationError, match="`explode` operation not supported for dtype" + InvalidOperationError, + match=re_compile(r"explode.+not supported for.+string.+expected.+list"), ): dataframe(data).explode("a") From 3afdaba62ec1525754bbdb7f50d22517f9f97d21 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 17:23:44 +0000 Subject: [PATCH 04/21] feat: Support `ArrowDataFrame.explode([column])` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mentioned in (https://github.com/narwhals-dev/narwhals/pull/1644#issuecomment-3622051763) Multi-column coming up next 😄 --- narwhals/_plan/arrow/dataframe.py | 35 ++++++++++++++++- narwhals/_plan/arrow/expr.py | 6 +-- narwhals/_plan/arrow/functions.py | 63 +++++++++++++++++++++++++++++++ narwhals/_plan/arrow/typing.py | 1 + narwhals/_plan/exceptions.py | 5 +++ tests/plan/explode_test.py | 5 +-- 6 files changed, 106 insertions(+), 9 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index 465afdbf39..aa4d097acb 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -16,6 +16,7 @@ from narwhals._plan.arrow.series import ArrowSeries as Series from narwhals._plan.compliant.dataframe import EagerDataFrame from narwhals._plan.compliant.typing import namespace +from narwhals._plan.exceptions import shape_error from narwhals._plan.expressions import NamedIR from narwhals._utils import Version, generate_repr from narwhals.schema import Schema @@ -26,7 +27,7 @@ import polars as pl from typing_extensions import Self, TypeAlias - from narwhals._plan.arrow.typing import ChunkedArrayAny + from narwhals._plan.arrow.typing import ChunkedArrayAny, ChunkedOrArrayAny from narwhals._plan.compliant.group_by import GroupByResolver from narwhals._plan.expressions import ExprIR, NamedIR from narwhals._plan.options import SortMultipleOptions @@ -163,7 +164,14 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self: return self._with_native(native) def explode(self, subset: Sequence[str]) -> Self: - msg = "ArrowDataFrame.explode" + if len(subset) == 1: + name = subset[0] + exploded, indices = fn.table_explode_1(self.native.column(name)) + # TODO @dangotbanned: Might be more efficient to null-out the column, before `gather`? + df = self.gather(indices) if len(indices) != len(self) else self + ser = Series.from_native(exploded, name, version=self.version) + return df.with_series(ser) + msg = "TODO: `ArrowDataFrame.explode((..., ...))`" raise NotImplementedError(msg) def rename(self, mapping: Mapping[str, str]) -> Self: @@ -174,6 +182,22 @@ def rename(self, mapping: Mapping[str, str]) -> Self: names = [mapping.get(c, c) for c in self.columns] return self._with_native(self.native.rename_columns(names)) + def with_series(self, series: Series) -> Self: + """Add a new column or replace an existing one. + + Uses similar semantics as `with_columns`, but: + - for a single named `Series` + - no broadcasting (use `Scalar.broadcast` instead) + - no length checking (use `with_series_checked` instead) + """ + return self._with_native(with_array(self.native, series.name, series.native)) + + def with_series_checked(self, series: Series) -> Self: + expected, actual = len(self), len(series) + if len(series) != len(self): + raise shape_error(expected, actual) + return self.with_series(series) + # NOTE: Use instead of `with_columns` for trivial cases def _with_columns(self, exprs: Iterable[Expr | Scalar], /) -> Self: native = self.native @@ -230,3 +254,10 @@ def partition_by(self, by: Sequence[str], *, include_key: bool = True) -> list[S from_native = self._with_native partitions = partition_by(self.native, by, include_key=include_key) return [from_native(df) for df in partitions] + + +def with_array(table: pa.Table, name: str, column: ChunkedOrArrayAny) -> pa.Table: + columns = table.column_names + if name in columns: + return table.set_column(columns.index(name), name, column) + return table.append_column(name, column) diff --git a/narwhals/_plan/arrow/expr.py b/narwhals/_plan/arrow/expr.py index 14e3ac893e..5213b72442 100644 --- a/narwhals/_plan/arrow/expr.py +++ b/narwhals/_plan/arrow/expr.py @@ -28,6 +28,7 @@ from narwhals._plan.compliant.expr import EagerExpr from narwhals._plan.compliant.scalar import EagerScalar from narwhals._plan.compliant.typing import namespace +from narwhals._plan.exceptions import shape_error from narwhals._plan.expressions import FunctionExpr as FExpr, functions as F from narwhals._plan.expressions.boolean import ( IsDuplicated, @@ -48,7 +49,7 @@ not_implemented, qualified_type_name, ) -from narwhals.exceptions import InvalidOperationError, ShapeError +from narwhals.exceptions import InvalidOperationError if TYPE_CHECKING: from collections.abc import Callable, Mapping, Sequence @@ -405,8 +406,7 @@ def to_series(self) -> Series: def broadcast(self, length: int, /) -> Series: if (actual_len := len(self)) != length: - msg = f"Expected object of length {length}, got {actual_len}." - raise ShapeError(msg) + raise shape_error(length, actual_len) return self._evaluated def __len__(self) -> int: diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index d69787110d..d9b501fc89 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -46,6 +46,7 @@ BoolType, ChunkedArray, ChunkedArrayAny, + ChunkedI64, ChunkedList, ChunkedOrArray, ChunkedOrArrayAny, @@ -386,6 +387,68 @@ def get_categories(native: ArrowAny) -> ChunkedArrayAny: return chunked_array(da.dictionary) +@t.overload +def list_explode( + native: ChunkedList[DataTypeT] | ListScalar[DataTypeT], +) -> ChunkedArray[Scalar[DataTypeT]]: ... +@t.overload +def list_explode(native: ListArray[DataTypeT]) -> Array[Scalar[DataTypeT]]: ... +@t.overload +def list_explode( + native: Arrow[ListScalar[DataTypeT]], +) -> ChunkedOrArray[Scalar[DataTypeT]]: ... +def list_explode( + native: ArrowAny, + *, + empty_as_null: bool = True, # noqa: ARG001 + keep_nulls: bool = True, # noqa: ARG001 +) -> ChunkedOrArray[Scalar[DataTypeT]]: + """Explode list elements, expanding one-level into a new array. + + Equivalent to `polars.Expr.(list.)explode`. + + Unused arguments are related to ([#1644 (comment)](https://github.com/narwhals-dev/narwhals/pull/1644#issuecomment-3622051763)) + """ + lengths = list_len(native) + needs_replacing = or_(is_null(lengths), eq(lengths, lit(0))) + to_explode = when_then(needs_replacing, lit([None], native.type), native) + # NOTE: Maybe reconsider scalar re-wrap for multiple levels of nesting? + # For the single case, it matches polars + if isinstance(native, pa.Scalar): + return chunked_array(_list_explode_unchecked(to_explode)) + result: ChunkedOrArray[Scalar[DataTypeT]] = _list_explode_unchecked(to_explode) + return result + + +def table_explode_1( + native: ChunkedList[DataTypeT], +) -> tuple[ChunkedArray[Scalar[DataTypeT]], ChunkedI64]: + """Variant of `list_explode`, providing indices for `DataFrame.explode([column])`.""" + lengths = list_len(native) + needs_replacing = or_(is_null(lengths), eq(lengths, lit(0))) + to_explode = when_then(needs_replacing, lit([None], native.type), native) + result: ChunkedArray[Scalar[DataTypeT]] = _list_explode_unchecked(to_explode) + return result, _list_parent_indices(to_explode) + + +def _list_explode_unchecked(native: Incomplete) -> Incomplete: + return pc.call_function("list_flatten", [native]) + + +@t.overload +def _list_parent_indices(native: ChunkedList) -> ChunkedI64: ... +@t.overload +def _list_parent_indices(native: ListArray) -> pa.Int64Array: ... +def _list_parent_indices( + native: ChunkedOrArray[ListScalar], +) -> ChunkedOrArray[pa.Int64Scalar]: + """Don't use this withut handling nulls!""" + result: ChunkedOrArray[pa.Int64Scalar] = pc.call_function( + "list_parent_indices", [native] + ) + return result + + @t.overload def list_len(native: ChunkedList) -> ChunkedArray[pa.UInt32Scalar]: ... @t.overload diff --git a/narwhals/_plan/arrow/typing.py b/narwhals/_plan/arrow/typing.py index 63ee251997..047e436eb4 100644 --- a/narwhals/_plan/arrow/typing.py +++ b/narwhals/_plan/arrow/typing.py @@ -195,6 +195,7 @@ class BinaryLogical(BinaryFunction["BooleanScalar", "BooleanScalar"], Protocol): StructArray: TypeAlias = "pa.StructArray | Array[pa.StructScalar]" ChunkedList: TypeAlias = "ChunkedArray[ListScalar[DataTypeT_co]]" ListArray: TypeAlias = "Array[ListScalar[DataTypeT_co]]" +ChunkedI64: TypeAlias = "ChunkedArray[pa.Int64Scalar]" Arrow: TypeAlias = "ChunkedOrScalar[ScalarT_co] | Array[ScalarT_co]" ArrowAny: TypeAlias = "ChunkedOrScalarAny | ArrayAny" diff --git a/narwhals/_plan/exceptions.py b/narwhals/_plan/exceptions.py index cac4616b9b..53aaa49f67 100644 --- a/narwhals/_plan/exceptions.py +++ b/narwhals/_plan/exceptions.py @@ -66,6 +66,11 @@ def hist_bins_monotonic_error(bins: Seq[float]) -> ComputeError: # noqa: ARG001 return ComputeError(msg) +def shape_error(expected_length: int, actual_length: int) -> ShapeError: + msg = f"Expected object of length {expected_length}, got {actual_length}." + return ShapeError(msg) + + def _binary_underline( left: ir.ExprIR, operator: Operator, diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index 3a5867299a..c05fb72fd5 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -28,16 +28,13 @@ def data() -> Data: } -@pytest.mark.xfail( - reason="TODO:` DataFrame.explode` (single column)", raises=NotImplementedError -) @pytest.mark.parametrize( ("column", "expected_values"), [("l2", [None, 3, None, None, 42]), ("l3", [1, 1, 2, 3, None])], ) def test_explode_single_col( column: str, expected_values: list[int | None], data: Data -) -> None: # pragma: no cover +) -> None: result = ( dataframe(data) .with_columns(nwp.col(column).cast(nw.List(nw.Int32()))) From 3fd946bebf96b7688bf06bfb4b950fc0384bbcdb Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 19:49:01 +0000 Subject: [PATCH 05/21] feat(expr-ir): Support multi* column explode Its enough to pass the test suite, but next up is refactoring `_table_explode_2` to handle everything >=2 --- narwhals/_plan/arrow/dataframe.py | 43 +++++++++++++++++-------------- narwhals/_plan/arrow/functions.py | 32 +++++++++++++++++++++++ tests/plan/explode_test.py | 11 ++------ 3 files changed, 58 insertions(+), 28 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index aa4d097acb..d0e7136203 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -18,7 +18,7 @@ from narwhals._plan.compliant.typing import namespace from narwhals._plan.exceptions import shape_error from narwhals._plan.expressions import NamedIR -from narwhals._utils import Version, generate_repr +from narwhals._utils import Version, generate_repr, zip_strict from narwhals.schema import Schema if TYPE_CHECKING: @@ -164,15 +164,18 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self: return self._with_native(native) def explode(self, subset: Sequence[str]) -> Self: + native = self.native if len(subset) == 1: name = subset[0] - exploded, indices = fn.table_explode_1(self.native.column(name)) + exploded, indices = fn.table_explode_1(native.column(name)) # TODO @dangotbanned: Might be more efficient to null-out the column, before `gather`? df = self.gather(indices) if len(indices) != len(self) else self ser = Series.from_native(exploded, name, version=self.version) return df.with_series(ser) - msg = "TODO: `ArrowDataFrame.explode((..., ...))`" - raise NotImplementedError(msg) + explodeds, indices = fn.table_explode_multi(*native.select(list(subset)).columns) + df = self.gather(indices) if len(indices) != len(self) else self + names_and_columns = zip_strict(subset, explodeds) + return self._with_native(with_arrays(df.native, names_and_columns)) def rename(self, mapping: Mapping[str, str]) -> Self: names: dict[str, str] | list[str] @@ -198,20 +201,10 @@ def with_series_checked(self, series: Series) -> Self: raise shape_error(expected, actual) return self.with_series(series) - # NOTE: Use instead of `with_columns` for trivial cases def _with_columns(self, exprs: Iterable[Expr | Scalar], /) -> Self: - native = self.native - columns = self.columns height = len(self) - for into_series in exprs: - name = into_series.name - chunked = into_series.broadcast(height).native - if name in columns: - i = columns.index(name) - native = native.set_column(i, name, chunked) - else: - native = native.append_column(name, chunked) - return self._with_native(native) + names_and_columns = ((e.name, e.broadcast(height).native) for e in exprs) + return self._with_native(with_arrays(self.native, names_and_columns)) def select_names(self, *column_names: str) -> Self: return self._with_native(self.native.select(list(column_names))) @@ -257,7 +250,19 @@ def partition_by(self, by: Sequence[str], *, include_key: bool = True) -> list[S def with_array(table: pa.Table, name: str, column: ChunkedOrArrayAny) -> pa.Table: - columns = table.column_names - if name in columns: - return table.set_column(columns.index(name), name, column) + column_names = table.column_names + if name in column_names: + return table.set_column(column_names.index(name), name, column) return table.append_column(name, column) + + +def with_arrays( + table: pa.Table, names_and_columns: Iterable[tuple[str, ChunkedOrArrayAny]], / +) -> pa.Table: + column_names = table.column_names + for name, column in names_and_columns: + if name in column_names: + table = table.set_column(column_names.index(name), name, column) + else: + table = table.append_column(name, column) + return table diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index d9b501fc89..d1f5ae9575 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -21,6 +21,7 @@ from narwhals._plan.arrow import options as pa_options from narwhals._plan.expressions import functions as F, operators as ops from narwhals._utils import Implementation, Version +from narwhals.exceptions import ShapeError if TYPE_CHECKING: import datetime as dt @@ -431,6 +432,37 @@ def table_explode_1( return result, _list_parent_indices(to_explode) +# TODO @dangotbanned: generalize to `*arrays` +# temp, while I figure out a clean loop/reduce version +def _table_explode_2( + left: ChunkedList, right: ChunkedList +) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: + left_len, right_len = list_len(left), list_len(right) + if not left_len.equals(right_len): + msg = "exploded columns must have matching element counts" + raise ShapeError(msg) + zero = lit(0) + needs_replacing = or_(is_null(left_len), eq(left_len, zero)) + to_explode_l = when_then(needs_replacing, lit([None], left.type), left) + to_explode_r = when_then(needs_replacing, lit([None], right.type), right) + result_l, result_r = ( + _list_explode_unchecked(to_explode_l), + _list_explode_unchecked(to_explode_r), + ) + indices = _list_parent_indices(to_explode_l) + results: Sequence[ChunkedArrayAny] = result_l, result_r + return results, indices + + +def table_explode_multi( + *arrays: ChunkedList, +) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: + if len(arrays) == 2: + return _table_explode_2(arrays[0], arrays[1]) + msg = "TODO: `ArrowDataFrame.explode((arr,arr, *arrays))`" + raise NotImplementedError(msg) + + def _list_explode_unchecked(native: Incomplete) -> Incomplete: return pc.call_function("list_flatten", [native]) diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index c05fb72fd5..c10dc1ba6e 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -46,9 +46,6 @@ def test_explode_single_col( assert_equal_data(result, expected) -@pytest.mark.xfail( - reason="TODO:` DataFrame.explode` (multi column)", raises=NotImplementedError -) @pytest.mark.parametrize( ("column", "more_columns", "expected"), [ @@ -77,7 +74,7 @@ def test_explode_multiple_cols( more_columns: Sequence[str], expected: dict[str, list[str | int | None]], data: Data, -) -> None: # pragma: no cover +) -> None: result = ( dataframe(data) .with_columns(nwp.col(column, *more_columns).cast(nw.List(nw.Int32()))) @@ -88,11 +85,7 @@ def test_explode_multiple_cols( assert_equal_data(result, expected) -@pytest.mark.xfail( - reason="TODO:` DataFrame.explode` (validate shape)", - raises=(AssertionError, NotImplementedError), -) -def test_explode_shape_error(data: Data) -> None: # pragma: no cover +def test_explode_shape_error(data: Data) -> None: with pytest.raises( ShapeError, match=r".*exploded columns (must )?have matching element counts" ): From fcb58b402dfa28764cc127e58e6074910787a22a Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 21:43:18 +0000 Subject: [PATCH 06/21] feat(expr-ir): Fully support multi column explode --- narwhals/_plan/arrow/dataframe.py | 4 +-- narwhals/_plan/arrow/functions.py | 44 +++++++++++-------------------- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index d0e7136203..db91a65718 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -167,12 +167,12 @@ def explode(self, subset: Sequence[str]) -> Self: native = self.native if len(subset) == 1: name = subset[0] - exploded, indices = fn.table_explode_1(native.column(name)) + exploded, indices = fn.table_explode_array(native.column(name)) # TODO @dangotbanned: Might be more efficient to null-out the column, before `gather`? df = self.gather(indices) if len(indices) != len(self) else self ser = Series.from_native(exploded, name, version=self.version) return df.with_series(ser) - explodeds, indices = fn.table_explode_multi(*native.select(list(subset)).columns) + explodeds, indices = fn.table_explode_arrays(*native.select(list(subset)).columns) df = self.gather(indices) if len(indices) != len(self) else self names_and_columns = zip_strict(subset, explodeds) return self._with_native(with_arrays(df.native, names_and_columns)) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index d1f5ae9575..3348f0bab4 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -4,6 +4,7 @@ import math import typing as t +from collections import deque from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, overload @@ -421,7 +422,7 @@ def list_explode( return result -def table_explode_1( +def table_explode_array( native: ChunkedList[DataTypeT], ) -> tuple[ChunkedArray[Scalar[DataTypeT]], ChunkedI64]: """Variant of `list_explode`, providing indices for `DataFrame.explode([column])`.""" @@ -432,35 +433,22 @@ def table_explode_1( return result, _list_parent_indices(to_explode) -# TODO @dangotbanned: generalize to `*arrays` -# temp, while I figure out a clean loop/reduce version -def _table_explode_2( - left: ChunkedList, right: ChunkedList -) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: - left_len, right_len = list_len(left), list_len(right) - if not left_len.equals(right_len): - msg = "exploded columns must have matching element counts" - raise ShapeError(msg) - zero = lit(0) - needs_replacing = or_(is_null(left_len), eq(left_len, zero)) - to_explode_l = when_then(needs_replacing, lit([None], left.type), left) - to_explode_r = when_then(needs_replacing, lit([None], right.type), right) - result_l, result_r = ( - _list_explode_unchecked(to_explode_l), - _list_explode_unchecked(to_explode_r), - ) - indices = _list_parent_indices(to_explode_l) - results: Sequence[ChunkedArrayAny] = result_l, result_r - return results, indices - - -def table_explode_multi( +def table_explode_arrays( *arrays: ChunkedList, ) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: - if len(arrays) == 2: - return _table_explode_2(arrays[0], arrays[1]) - msg = "TODO: `ArrowDataFrame.explode((arr,arr, *arrays))`" - raise NotImplementedError(msg) + """Variant of `table_explode_array`, with shape checking against the first array.""" + explode = _list_explode_unchecked + first = arrays[0] + first_len = list_len(first) + needs_replacing = or_(is_null(first_len), eq(first_len, lit(0))) + first_to_explode = when_then(needs_replacing, lit([None], first.type), first) + results = deque["ChunkedArrayAny"]([explode(first_to_explode)]) + for arr in arrays[1:]: + if not first_len.equals(list_len(arr)): + msg = "exploded columns must have matching element counts" + raise ShapeError(msg) + results.append(explode(when_then(needs_replacing, lit([None], arr.type), arr))) + return results, _list_parent_indices(first_to_explode) def _list_explode_unchecked(native: Incomplete) -> Incomplete: From 414d0f6b30429c28f712a0071ce6bb57d8fde789 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 21:43:48 +0000 Subject: [PATCH 07/21] test: Cover selectors, >2 columns --- tests/plan/explode_test.py | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index c10dc1ba6e..b01ecb3590 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -6,6 +6,7 @@ import narwhals as nw import narwhals._plan as nwp +import narwhals._plan.selectors as ncs from narwhals.exceptions import InvalidOperationError, ShapeError from tests.plan.utils import assert_equal_data, dataframe, re_compile @@ -25,6 +26,7 @@ def data() -> Data: "l2": [[3, None], None, [42], []], "l3": [[1, 2], [3], [None], [1]], "l4": [[1, 2], [3], [123], [456]], + "l5": [[None, None], [None], [99], [83]], } @@ -85,13 +87,46 @@ def test_explode_multiple_cols( assert_equal_data(result, expected) +@pytest.mark.parametrize( + ("expr", "expected"), + [ + ( + ncs.by_index(-1, -2, -3), + { + "a": ["w", "x", "x", "y", "z"], + "l5": [83, None, None, None, 99], + "l4": [456, 1, 2, 3, 123], + "l3": [1, 1, 2, 3, None], + }, + ), + ( + ncs.matches(r"l[3|5]"), + { + "a": ["w", "x", "x", "y", "z"], + "l3": [1, 1, 2, 3, None], + "l5": [83, None, None, None, 99], + }, + ), + ], +) +def test_explode_selectors(expr: nwp.Selector, expected: Data, data: Data) -> None: + result = ( + dataframe(data) + .with_columns(expr.cast(nw.List(nw.Int32()))) + .explode(expr) + .select("a", expr) + .sort("a", expr, nulls_last=True) + ) + assert_equal_data(result, expected) + + def test_explode_shape_error(data: Data) -> None: with pytest.raises( ShapeError, match=r".*exploded columns (must )?have matching element counts" ): dataframe(data).with_columns( nwp.col("l1", "l2", "l3").cast(nw.List(nw.Int32())) - ).explode("l1", "l3") + ).explode(ncs.list()) def test_explode_invalid_operation_error(data: Data) -> None: From 2174ae34a159c564b0e08b70f1a4068c963afc18 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Sun, 7 Dec 2025 23:19:58 +0000 Subject: [PATCH 08/21] feat(expr-ir): Add `Series.explode` Pretty sure the `AttributeError` is unrelated to this, but nice to have found --- narwhals/_plan/arrow/series.py | 6 +++ narwhals/_plan/compliant/series.py | 1 + narwhals/_plan/series.py | 5 ++ tests/plan/explode_test.py | 85 +++++++++++++++++++++++++++--- 4 files changed, 90 insertions(+), 7 deletions(-) diff --git a/narwhals/_plan/arrow/series.py b/narwhals/_plan/arrow/series.py index 72e01a4fb1..e625cee46c 100644 --- a/narwhals/_plan/arrow/series.py +++ b/narwhals/_plan/arrow/series.py @@ -301,6 +301,12 @@ def drop_nans(self) -> Self: self.native.filter(predicate, null_selection_behavior="emit_null") ) + def explode(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> Self: + if empty_as_null and keep_nulls: + return self._with_native(fn.list_explode(self.native)) + msg = f"TODO: `ArrowSeries.explode({empty_as_null=}, {keep_nulls=})" + raise NotImplementedError(msg) + @property def struct(self) -> SeriesStructNamespace: return SeriesStructNamespace(self) diff --git a/narwhals/_plan/compliant/series.py b/narwhals/_plan/compliant/series.py index 7a51d7802e..6416adbc06 100644 --- a/narwhals/_plan/compliant/series.py +++ b/narwhals/_plan/compliant/series.py @@ -133,6 +133,7 @@ def cum_sum(self, *, reverse: bool = False) -> Self: ... def diff(self, n: int = 1) -> Self: ... def drop_nulls(self) -> Self: ... def drop_nans(self) -> Self: ... + def explode(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> Self: ... def fill_nan(self, value: float | Self | None) -> Self: ... def fill_null(self, value: NonNestedLiteral | Self) -> Self: ... def fill_null_with_strategy( diff --git a/narwhals/_plan/series.py b/narwhals/_plan/series.py index 3b6407a0da..db799ff37f 100644 --- a/narwhals/_plan/series.py +++ b/narwhals/_plan/series.py @@ -298,6 +298,11 @@ def hist( return result return result.to_series().struct.unnest() + def explode(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> Self: + return type(self)( + self._compliant.explode(empty_as_null=empty_as_null, keep_nulls=keep_nulls) + ) + @property def struct(self) -> SeriesStructNamespace[Self]: return SeriesStructNamespace(self) diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index b01ecb3590..6aa17bac93 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Final import pytest @@ -8,7 +8,13 @@ import narwhals._plan as nwp import narwhals._plan.selectors as ncs from narwhals.exceptions import InvalidOperationError, ShapeError -from tests.plan.utils import assert_equal_data, dataframe, re_compile +from tests.plan.utils import ( + assert_equal_data, + assert_equal_series, + dataframe, + re_compile, + series, +) if TYPE_CHECKING: from collections.abc import Sequence @@ -34,7 +40,7 @@ def data() -> Data: ("column", "expected_values"), [("l2", [None, 3, None, None, 42]), ("l3", [1, 1, 2, 3, None])], ) -def test_explode_single_col( +def test_explode_frame_single_col( column: str, expected_values: list[int | None], data: Data ) -> None: result = ( @@ -71,7 +77,7 @@ def test_explode_single_col( ), ], ) -def test_explode_multiple_cols( +def test_explode_frame_multiple_cols( column: str, more_columns: Sequence[str], expected: dict[str, list[str | int | None]], @@ -109,7 +115,7 @@ def test_explode_multiple_cols( ), ], ) -def test_explode_selectors(expr: nwp.Selector, expected: Data, data: Data) -> None: +def test_explode_frame_selectors(expr: nwp.Selector, expected: Data, data: Data) -> None: result = ( dataframe(data) .with_columns(expr.cast(nw.List(nw.Int32()))) @@ -120,7 +126,7 @@ def test_explode_selectors(expr: nwp.Selector, expected: Data, data: Data) -> No assert_equal_data(result, expected) -def test_explode_shape_error(data: Data) -> None: +def test_explode_frame_shape_error(data: Data) -> None: with pytest.raises( ShapeError, match=r".*exploded columns (must )?have matching element counts" ): @@ -129,9 +135,74 @@ def test_explode_shape_error(data: Data) -> None: ).explode(ncs.list()) -def test_explode_invalid_operation_error(data: Data) -> None: +def test_explode_frame_invalid_operation_error(data: Data) -> None: with pytest.raises( InvalidOperationError, match=re_compile(r"explode.+not supported for.+string.+expected.+list"), ): dataframe(data).explode("a") + + +@pytest.mark.parametrize( + ("values", "expected"), + [ + ([[1, 2, 3]], [1, 2, 3]), + ([[1, 2, 3], None], [1, 2, 3, None]), + ([[1, 2, 3], []], [1, 2, 3, None]), + ], +) +def test_explode_series_default(values: list[Any], expected: list[Any]) -> None: + # Based on `test_explode_basic` in https://github.com/pola-rs/polars/issues/25289 + # https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L465-L505 + result = series(values).explode() + assert_equal_series(result, expected, "") + + +@pytest.mark.xfail( + reason="TODO: 'ArrowExpr' object has no attribute '_evaluated'", raises=AttributeError +) +@pytest.mark.parametrize( + ("values", "expected"), + [ + ([[1, 2, 3], [1, 2], [1, 2]], [1, 2, 3, None, 1, 2]), + ([[1, 2, 3], [], [1, 2]], [1, 2, 3, None, 1, 2]), + ], +) +def test_explode_series_default_masked( + values: list[Any], expected: list[Any] +) -> None: # pragma: no cover + result = ( + series(values) + .to_frame() + .select(nwp.when(series([True, False, True])).then(nwp.col(""))) + .to_series() + .explode() + ) + assert_equal_series(result, expected, "") + + +DROP_EMPTY: Final = {"empty_as_null": False} +DROP_NULLS: Final = {"keep_nulls": False} +DROP_BOTH: Final = {"empty_as_null": False, "keep_nulls": False} + + +@pytest.mark.xfail( + reason="TODO: Implement non-default `Series.explode(...)", raises=NotImplementedError +) +@pytest.mark.parametrize( + ("values", "kwds", "expected"), + [ + ([[1, 2, 3]], DROP_BOTH, [1, 2, 3]), + ([[1, 2, 3], None], DROP_NULLS, [1, 2, 3]), + ([[1, 2, 3], [None]], DROP_NULLS, [1, 2, 3, None]), + ([[1, 2, 3], []], DROP_EMPTY, [1, 2, 3]), + ([[1, 2, 3], [None]], DROP_EMPTY, [1, 2, 3, None]), + ], +) +def test_explode_series_options( + values: list[Any], kwds: dict[str, Any], expected: list[Any] +) -> None: # pragma: no cover + # Based on `test_explode_basic` in https://github.com/pola-rs/polars/issues/25289 + # https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L465-L505 + result = series(values).explode(**kwds) + assert_equal_series(result, expected, "") From c93620589e473edf88affd138e0010d2522376c9 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 12:16:13 +0000 Subject: [PATCH 09/21] chore: Pass through `explode` options The 3 passing tests are a coincidence, since they're supposed to be equivalent to the default --- narwhals/_plan/arrow/functions.py | 10 +++++++++- narwhals/_plan/arrow/series.py | 8 ++++---- tests/plan/explode_test.py | 13 +++++++------ 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index 3348f0bab4..99728c4e10 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -392,12 +392,20 @@ def get_categories(native: ArrowAny) -> ChunkedArrayAny: @t.overload def list_explode( native: ChunkedList[DataTypeT] | ListScalar[DataTypeT], + *, + empty_as_null: bool = ..., + keep_nulls: bool = ..., ) -> ChunkedArray[Scalar[DataTypeT]]: ... @t.overload -def list_explode(native: ListArray[DataTypeT]) -> Array[Scalar[DataTypeT]]: ... +def list_explode( + native: ListArray[DataTypeT], *, empty_as_null: bool = ..., keep_nulls: bool = ... +) -> Array[Scalar[DataTypeT]]: ... @t.overload def list_explode( native: Arrow[ListScalar[DataTypeT]], + *, + empty_as_null: bool = ..., + keep_nulls: bool = ..., ) -> ChunkedOrArray[Scalar[DataTypeT]]: ... def list_explode( native: ArrowAny, diff --git a/narwhals/_plan/arrow/series.py b/narwhals/_plan/arrow/series.py index e625cee46c..dc606d0927 100644 --- a/narwhals/_plan/arrow/series.py +++ b/narwhals/_plan/arrow/series.py @@ -302,10 +302,10 @@ def drop_nans(self) -> Self: ) def explode(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> Self: - if empty_as_null and keep_nulls: - return self._with_native(fn.list_explode(self.native)) - msg = f"TODO: `ArrowSeries.explode({empty_as_null=}, {keep_nulls=})" - raise NotImplementedError(msg) + result = fn.list_explode( + self.native, empty_as_null=empty_as_null, keep_nulls=keep_nulls + ) + return self._with_native(result) @property def struct(self) -> SeriesStructNamespace: diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index 6aa17bac93..af914ec1e1 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -185,23 +185,24 @@ def test_explode_series_default_masked( DROP_NULLS: Final = {"keep_nulls": False} DROP_BOTH: Final = {"empty_as_null": False, "keep_nulls": False} - -@pytest.mark.xfail( - reason="TODO: Implement non-default `Series.explode(...)", raises=NotImplementedError +XFAIL_EXPLODE_KWDS = pytest.mark.xfail( + reason="TODO: Implement non-default `Series.explode(...)" ) + + @pytest.mark.parametrize( ("values", "kwds", "expected"), [ ([[1, 2, 3]], DROP_BOTH, [1, 2, 3]), - ([[1, 2, 3], None], DROP_NULLS, [1, 2, 3]), + pytest.param([[1, 2, 3], None], DROP_NULLS, [1, 2, 3], marks=XFAIL_EXPLODE_KWDS), ([[1, 2, 3], [None]], DROP_NULLS, [1, 2, 3, None]), - ([[1, 2, 3], []], DROP_EMPTY, [1, 2, 3]), + pytest.param([[1, 2, 3], []], DROP_EMPTY, [1, 2, 3], marks=XFAIL_EXPLODE_KWDS), ([[1, 2, 3], [None]], DROP_EMPTY, [1, 2, 3, None]), ], ) def test_explode_series_options( values: list[Any], kwds: dict[str, Any], expected: list[Any] -) -> None: # pragma: no cover +) -> None: # Based on `test_explode_basic` in https://github.com/pola-rs/polars/issues/25289 # https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L465-L505 result = series(values).explode(**kwds) From 314f3599e99360a0de2026ef49933685258ca4ba Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 12:26:24 +0000 Subject: [PATCH 10/21] feat(expr-ir): Support `Series.explode(empty_as_null=..., keep_nulls=...)` --- narwhals/_plan/arrow/functions.py | 23 +++++++++++++---------- narwhals/_plan/arrow/series.py | 5 ++--- tests/plan/explode_test.py | 8 ++------ 3 files changed, 17 insertions(+), 19 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index 99728c4e10..157257a2e2 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -408,20 +408,23 @@ def list_explode( keep_nulls: bool = ..., ) -> ChunkedOrArray[Scalar[DataTypeT]]: ... def list_explode( - native: ArrowAny, - *, - empty_as_null: bool = True, # noqa: ARG001 - keep_nulls: bool = True, # noqa: ARG001 + native: ArrowAny, *, empty_as_null: bool = True, keep_nulls: bool = True ) -> ChunkedOrArray[Scalar[DataTypeT]]: """Explode list elements, expanding one-level into a new array. - Equivalent to `polars.Expr.(list.)explode`. - - Unused arguments are related to ([#1644 (comment)](https://github.com/narwhals-dev/narwhals/pull/1644#issuecomment-3622051763)) + Equivalent to `polars.{Expr,Series}.explode`. """ - lengths = list_len(native) - needs_replacing = or_(is_null(lengths), eq(lengths, lit(0))) - to_explode = when_then(needs_replacing, lit([None], native.type), native) + if empty_as_null or keep_nulls: + lengths = list_len(native) + if empty_as_null and keep_nulls: + needs_replacing = or_(is_null(lengths), eq(lengths, lit(0))) + elif empty_as_null: + needs_replacing = eq(lengths, lit(0)) + else: + needs_replacing = is_null(lengths) + to_explode = when_then(needs_replacing, lit([None], native.type), native) + else: + to_explode = native # NOTE: Maybe reconsider scalar re-wrap for multiple levels of nesting? # For the single case, it matches polars if isinstance(native, pa.Scalar): diff --git a/narwhals/_plan/arrow/series.py b/narwhals/_plan/arrow/series.py index dc606d0927..3adebde9c1 100644 --- a/narwhals/_plan/arrow/series.py +++ b/narwhals/_plan/arrow/series.py @@ -302,9 +302,8 @@ def drop_nans(self) -> Self: ) def explode(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> Self: - result = fn.list_explode( - self.native, empty_as_null=empty_as_null, keep_nulls=keep_nulls - ) + ca = self.native + result = fn.list_explode(ca, empty_as_null=empty_as_null, keep_nulls=keep_nulls) return self._with_native(result) @property diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index af914ec1e1..affa4936e9 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -185,18 +185,14 @@ def test_explode_series_default_masked( DROP_NULLS: Final = {"keep_nulls": False} DROP_BOTH: Final = {"empty_as_null": False, "keep_nulls": False} -XFAIL_EXPLODE_KWDS = pytest.mark.xfail( - reason="TODO: Implement non-default `Series.explode(...)" -) - @pytest.mark.parametrize( ("values", "kwds", "expected"), [ ([[1, 2, 3]], DROP_BOTH, [1, 2, 3]), - pytest.param([[1, 2, 3], None], DROP_NULLS, [1, 2, 3], marks=XFAIL_EXPLODE_KWDS), + ([[1, 2, 3], None], DROP_NULLS, [1, 2, 3]), ([[1, 2, 3], [None]], DROP_NULLS, [1, 2, 3, None]), - pytest.param([[1, 2, 3], []], DROP_EMPTY, [1, 2, 3], marks=XFAIL_EXPLODE_KWDS), + ([[1, 2, 3], []], DROP_EMPTY, [1, 2, 3]), ([[1, 2, 3], [None]], DROP_EMPTY, [1, 2, 3, None]), ], ) From 734792fb64a3291d27c0c586864bf8c53a78efc5 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 12:50:18 +0000 Subject: [PATCH 11/21] fix: Don't fall back to `ArrowExpr.name` --- narwhals/_plan/arrow/expr.py | 2 +- tests/plan/explode_test.py | 14 ++++---------- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/narwhals/_plan/arrow/expr.py b/narwhals/_plan/arrow/expr.py index 5213b72442..15ee7e5f91 100644 --- a/narwhals/_plan/arrow/expr.py +++ b/narwhals/_plan/arrow/expr.py @@ -373,7 +373,7 @@ def _with_native(self, result: ChunkedOrScalarAny, name: str, /) -> Scalar | Sel def _with_native(self, result: ChunkedOrScalarAny, name: str, /) -> Scalar | Self: if isinstance(result, pa.Scalar): return ArrowScalar.from_native(result, name, version=self.version) - return self.from_native(result, name or self.name, self.version) + return self.from_native(result, name, self.version) # NOTE: I'm not sure what I meant by # > "isn't natively supported on `ChunkedArray`" diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index affa4936e9..fe395db711 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -152,15 +152,11 @@ def test_explode_frame_invalid_operation_error(data: Data) -> None: ], ) def test_explode_series_default(values: list[Any], expected: list[Any]) -> None: - # Based on `test_explode_basic` in https://github.com/pola-rs/polars/issues/25289 - # https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L465-L505 + # Based on https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L465-L470 result = series(values).explode() assert_equal_series(result, expected, "") -@pytest.mark.xfail( - reason="TODO: 'ArrowExpr' object has no attribute '_evaluated'", raises=AttributeError -) @pytest.mark.parametrize( ("values", "expected"), [ @@ -168,9 +164,8 @@ def test_explode_series_default(values: list[Any], expected: list[Any]) -> None: ([[1, 2, 3], [], [1, 2]], [1, 2, 3, None, 1, 2]), ], ) -def test_explode_series_default_masked( - values: list[Any], expected: list[Any] -) -> None: # pragma: no cover +def test_explode_series_default_masked(values: list[Any], expected: list[Any]) -> None: + # Based on https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L471-484 result = ( series(values) .to_frame() @@ -199,7 +194,6 @@ def test_explode_series_default_masked( def test_explode_series_options( values: list[Any], kwds: dict[str, Any], expected: list[Any] ) -> None: - # Based on `test_explode_basic` in https://github.com/pola-rs/polars/issues/25289 - # https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L465-L505 + # Based on https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L486-L505 result = series(values).explode(**kwds) assert_equal_series(result, expected, "") From 4987ebb39b2d2e92931f893a400d4ba594c42a20 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 16:15:17 +0000 Subject: [PATCH 12/21] prep for `DataFrame.explode(..., **kwds)` --- narwhals/_plan/arrow/dataframe.py | 4 ++-- narwhals/_plan/compliant/dataframe.py | 4 ++-- narwhals/_plan/dataframe.py | 7 ++++-- narwhals/_plan/options.py | 12 +++++++++++ tests/plan/explode_test.py | 31 +++++++++++++++++++++++++++ 5 files changed, 52 insertions(+), 6 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index db91a65718..3cbf4d7bff 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -30,7 +30,7 @@ from narwhals._plan.arrow.typing import ChunkedArrayAny, ChunkedOrArrayAny from narwhals._plan.compliant.group_by import GroupByResolver from narwhals._plan.expressions import ExprIR, NamedIR - from narwhals._plan.options import SortMultipleOptions + from narwhals._plan.options import ExplodeOptions, SortMultipleOptions from narwhals._plan.typing import NonCrossJoinStrategy from narwhals.dtypes import DType from narwhals.typing import IntoSchema @@ -163,7 +163,7 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self: native = self.native.filter(~to_drop) return self._with_native(native) - def explode(self, subset: Sequence[str]) -> Self: + def explode(self, subset: Sequence[str], options: ExplodeOptions) -> Self: native = self.native if len(subset) == 1: name = subset[0] diff --git a/narwhals/_plan/compliant/dataframe.py b/narwhals/_plan/compliant/dataframe.py index b40019d788..2787f449e2 100644 --- a/narwhals/_plan/compliant/dataframe.py +++ b/narwhals/_plan/compliant/dataframe.py @@ -30,7 +30,7 @@ from narwhals._plan.compliant.namespace import EagerNamespace from narwhals._plan.dataframe import BaseFrame, DataFrame from narwhals._plan.expressions import NamedIR - from narwhals._plan.options import SortMultipleOptions + from narwhals._plan.options import ExplodeOptions, SortMultipleOptions from narwhals._plan.typing import Seq from narwhals._typing import _EagerAllowedImpl from narwhals._utils import Implementation, Version @@ -59,7 +59,7 @@ def to_narwhals(self) -> BaseFrame[NativeFrameT_co]: ... def columns(self) -> list[str]: ... def drop(self, columns: Sequence[str]) -> Self: ... def drop_nulls(self, subset: Sequence[str] | None) -> Self: ... - def explode(self, subset: Sequence[str]) -> Self: ... + def explode(self, subset: Sequence[str], options: ExplodeOptions) -> Self: ... # Shouldn't *need* to be `NamedIR`, but current impl depends on a name being passed around def filter(self, predicate: NamedIR, /) -> Self: ... def rename(self, mapping: Mapping[str, str]) -> Self: ... diff --git a/narwhals/_plan/dataframe.py b/narwhals/_plan/dataframe.py index c14343be09..2941c30a30 100644 --- a/narwhals/_plan/dataframe.py +++ b/narwhals/_plan/dataframe.py @@ -7,7 +7,7 @@ from narwhals._plan._guards import is_series from narwhals._plan.common import ensure_seq_str, temp from narwhals._plan.group_by import GroupBy, Grouped -from narwhals._plan.options import SortMultipleOptions +from narwhals._plan.options import ExplodeOptions, SortMultipleOptions from narwhals._plan.series import Series from narwhals._plan.typing import ( ColumnNameOrSelector, @@ -162,6 +162,8 @@ def explode( self, columns: OneOrIterable[ColumnNameOrSelector], *more_columns: ColumnNameOrSelector, + empty_as_null: bool = True, + keep_nulls: bool = True, ) -> Self: s_ir = _parse.parse_into_combined_selector_ir(columns, *more_columns) schema = self.collect_schema() @@ -173,7 +175,8 @@ def explode( if dtype != tp_list: msg = f"`explode` operation is not supported for dtype `{dtype}`, expected List type" raise InvalidOperationError(msg) - return self._with_compliant(self._compliant.explode(subset)) + options = ExplodeOptions(empty_as_null=empty_as_null, keep_nulls=keep_nulls) + return self._with_compliant(self._compliant.explode(subset, options)) def _dataframe_from_dict( diff --git a/narwhals/_plan/options.py b/narwhals/_plan/options.py index 4914b8b59b..ab43aa23ed 100644 --- a/narwhals/_plan/options.py +++ b/narwhals/_plan/options.py @@ -328,3 +328,15 @@ def default(cls) -> Self: FEOptions = FunctionExprOptions + + +class ExplodeOptions(Immutable): + __slots__ = ("empty_as_null", "keep_nulls") + empty_as_null: bool + """Explode an empty list into a `null`.""" + keep_nulls: bool + """Explode a `null` into a `null`.""" + + def any(self) -> bool: + """Return True if we need to handle empty lists and/or nulls.""" + return self.empty_as_null or self.keep_nulls diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index fe395db711..614ab65d74 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -197,3 +197,34 @@ def test_explode_series_options( # Based on https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L486-L505 result = series(values).explode(**kwds) assert_equal_series(result, expected, "") + + +XFAIL_NOT_IMPL = pytest.mark.xfail(reason="TODO: `DataFrame.explode(..., **kwds)`") + + +@pytest.mark.parametrize( + ("kwds", "expected"), + [ + ({}, {"a": [1, 2, 3, None, 4, 5, 6, None], "b": [1, 1, 1, 2, 3, 3, 3, 4]}), + pytest.param( + DROP_EMPTY, + {"a": [1, 2, 3, None, 4, 5, 6], "b": [1, 1, 1, 2, 3, 3, 3]}, + marks=XFAIL_NOT_IMPL, + ), + pytest.param( + DROP_NULLS, + {"a": [1, 2, 3, 4, 5, 6, None], "b": [1, 1, 1, 3, 3, 3, 4]}, + marks=XFAIL_NOT_IMPL, + ), + pytest.param( + DROP_BOTH, + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 1, 3, 3, 3]}, + marks=XFAIL_NOT_IMPL, + ), + ], +) +def test_explode_frame_options(kwds: dict[str, Any], expected: Data) -> None: + data = {"a": [[1, 2, 3], None, [4, 5, 6], []], "b": [1, 2, 3, 4]} + # Based on https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L596-L616 + result = dataframe(data).explode("a", **kwds) + assert_equal_data(result, expected) From 5a799ce3012d729919382be5b44b028878e8a0f8 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 17:20:52 +0000 Subject: [PATCH 13/21] feat(expr-ir): Support `DataFrame.explode(*columns, **kwds)` --- narwhals/_plan/arrow/dataframe.py | 8 +- narwhals/_plan/arrow/functions.py | 138 ++++++++++++++++++++---------- tests/plan/explode_test.py | 69 +++++++++++---- 3 files changed, 152 insertions(+), 63 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index 3cbf4d7bff..9a6c8c7214 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -163,16 +163,20 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self: native = self.native.filter(~to_drop) return self._with_native(native) + # TODO @dangotbanned: Move move of this into `ExplodeBuilder` def explode(self, subset: Sequence[str], options: ExplodeOptions) -> Self: native = self.native + builder = fn.ExplodeBuilder.from_options(options) if len(subset) == 1: name = subset[0] - exploded, indices = fn.table_explode_array(native.column(name)) + exploded, indices = builder.explode_into(native.column(name)) # TODO @dangotbanned: Might be more efficient to null-out the column, before `gather`? df = self.gather(indices) if len(indices) != len(self) else self ser = Series.from_native(exploded, name, version=self.version) return df.with_series(ser) - explodeds, indices = fn.table_explode_arrays(*native.select(list(subset)).columns) + explodeds, indices = builder.explode_arrays_into( + *native.select(list(subset)).columns + ) df = self.gather(indices) if len(indices) != len(self) else self names_and_columns = zip_strict(subset, explodeds) return self._with_native(with_arrays(df.native, names_and_columns)) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index 157257a2e2..db94be3de9 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -6,7 +6,7 @@ import typing as t from collections import deque from collections.abc import Callable, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, overload +from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload import pyarrow as pa # ignore-banned-import import pyarrow.compute as pc # ignore-banned-import @@ -21,6 +21,7 @@ from narwhals._plan._guards import is_non_nested_literal from narwhals._plan.arrow import options as pa_options from narwhals._plan.expressions import functions as F, operators as ops +from narwhals._plan.options import ExplodeOptions from narwhals._utils import Implementation, Version from narwhals.exceptions import ShapeError @@ -28,7 +29,7 @@ import datetime as dt from collections.abc import Iterable, Mapping - from typing_extensions import TypeAlias, TypeIs, TypeVarTuple, Unpack + from typing_extensions import Self, TypeAlias, TypeIs, TypeVarTuple, Unpack from narwhals._arrow.typing import Incomplete, PromoteOptions from narwhals._plan.arrow.acero import Field @@ -414,52 +415,97 @@ def list_explode( Equivalent to `polars.{Expr,Series}.explode`. """ - if empty_as_null or keep_nulls: - lengths = list_len(native) - if empty_as_null and keep_nulls: - needs_replacing = or_(is_null(lengths), eq(lengths, lit(0))) - elif empty_as_null: - needs_replacing = eq(lengths, lit(0)) - else: - needs_replacing = is_null(lengths) - to_explode = when_then(needs_replacing, lit([None], native.type), native) - else: - to_explode = native - # NOTE: Maybe reconsider scalar re-wrap for multiple levels of nesting? - # For the single case, it matches polars - if isinstance(native, pa.Scalar): - return chunked_array(_list_explode_unchecked(to_explode)) - result: ChunkedOrArray[Scalar[DataTypeT]] = _list_explode_unchecked(to_explode) - return result + return ExplodeBuilder(empty_as_null=empty_as_null, keep_nulls=keep_nulls).explode( + native + ) + + +_ArrowListT = TypeVar("_ArrowListT", bound="Arrow[ListScalar[Any]]") + +class ExplodeBuilder: + options: ExplodeOptions -def table_explode_array( - native: ChunkedList[DataTypeT], -) -> tuple[ChunkedArray[Scalar[DataTypeT]], ChunkedI64]: - """Variant of `list_explode`, providing indices for `DataFrame.explode([column])`.""" - lengths = list_len(native) - needs_replacing = or_(is_null(lengths), eq(lengths, lit(0))) - to_explode = when_then(needs_replacing, lit([None], native.type), native) - result: ChunkedArray[Scalar[DataTypeT]] = _list_explode_unchecked(to_explode) - return result, _list_parent_indices(to_explode) - - -def table_explode_arrays( - *arrays: ChunkedList, -) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: - """Variant of `table_explode_array`, with shape checking against the first array.""" - explode = _list_explode_unchecked - first = arrays[0] - first_len = list_len(first) - needs_replacing = or_(is_null(first_len), eq(first_len, lit(0))) - first_to_explode = when_then(needs_replacing, lit([None], first.type), first) - results = deque["ChunkedArrayAny"]([explode(first_to_explode)]) - for arr in arrays[1:]: - if not first_len.equals(list_len(arr)): - msg = "exploded columns must have matching element counts" - raise ShapeError(msg) - results.append(explode(when_then(needs_replacing, lit([None], arr.type), arr))) - return results, _list_parent_indices(first_to_explode) + def __init__(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> None: + self.options = ExplodeOptions(empty_as_null=empty_as_null, keep_nulls=keep_nulls) + + @classmethod + def from_options(cls, options: ExplodeOptions, /) -> Self: + obj = cls.__new__(cls) + obj.options = options + return obj + + def explode( + self, native: Arrow[ListScalar[DataTypeT]] + ) -> ChunkedOrArray[Scalar[DataTypeT]]: + explode = _list_explode_unchecked + if self.options.any(): + lengths = list_len(native) + needs_replacing = self._predicate(lengths) + safe = self._replace_mask(native, needs_replacing) + else: + safe = native + result: ChunkedOrArray[Scalar[DataTypeT]] = ( + chunked_array(explode(safe)) if isinstance(safe, pa.Scalar) else explode(safe) + ) + return result + + # TODO @dangotbanned: likely want to be passing in the native table *somewhere* + # - Rather than returning a tuple that represents a half-done job + def explode_into( + self, native: ChunkedList[DataTypeT] + ) -> tuple[ChunkedArray[Scalar[DataTypeT]], ChunkedI64]: + """Variant of `explode`, providing indices for `DataFrame.explode([column])`.""" + result: ChunkedArray[Scalar[DataTypeT]] + if self.options.any(): + lengths = list_len(native) + needs_replacing = self._predicate(lengths) + safe = self._replace_mask(native, needs_replacing) + else: + safe = native + result = _list_explode_unchecked(safe) + return result, _list_parent_indices(safe) + + # TODO @dangotbanned: De-duplicate + def explode_arrays_into( + self, *arrays: ChunkedList + ) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: + """Variant of `explode_into`, with shape checking against the first array.""" + explode = _list_explode_unchecked + first = arrays[0] + first_len = list_len(first) + results = deque["ChunkedArrayAny"]() + if self.options.any(): + needs_replacing = self._predicate(first_len) + first_safe = self._replace_mask(first, needs_replacing) + results.append(explode(first_safe)) + for arr in arrays[1:]: + if not first_len.equals(list_len(arr)): + msg = "exploded columns must have matching element counts" + raise ShapeError(msg) + results.append(explode(self._replace_mask(arr, needs_replacing))) + else: + first_safe = first + results.append(explode(first_safe)) + for arr in arrays[1:]: + if not first_len.equals(list_len(arr)): + msg = "exploded columns must have matching element counts" + raise ShapeError(msg) + results.append(explode(arr)) + return results, _list_parent_indices(first_safe) + + def _predicate(self, lengths: ArrowAny) -> Arrow[BooleanScalar]: + empty_as_null, keep_nulls = self.options.empty_as_null, self.options.keep_nulls + if empty_as_null and keep_nulls: + return or_(is_null(lengths), eq(lengths, lit(0))) + if empty_as_null: + return eq(lengths, lit(0)) + return is_null(lengths) + + def _replace_mask( + self, native: _ArrowListT, mask: Arrow[BooleanScalar] + ) -> _ArrowListT: + return when_then(mask, lit([None], native.type), native) # type: ignore[no-any-return] def _list_explode_unchecked(native: Incomplete) -> Incomplete: diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index 614ab65d74..ecffea2818 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: from collections.abc import Sequence + from narwhals._plan.typing import ColumnNameOrSelector from tests.conftest import Data @@ -199,32 +200,70 @@ def test_explode_series_options( assert_equal_series(result, expected, "") -XFAIL_NOT_IMPL = pytest.mark.xfail(reason="TODO: `DataFrame.explode(..., **kwds)`") +A = ("a",) +BA = "b", "a" @pytest.mark.parametrize( - ("kwds", "expected"), + ("columns", "kwds", "expected"), [ - ({}, {"a": [1, 2, 3, None, 4, 5, 6, None], "b": [1, 1, 1, 2, 3, 3, 3, 4]}), - pytest.param( + (A, {}, {"a": [1, 2, 3, None, 4, 5, 6, None], "i": [1, 1, 1, 2, 3, 3, 3, 4]}), + (A, DROP_EMPTY, {"a": [1, 2, 3, None, 4, 5, 6], "i": [1, 1, 1, 2, 3, 3, 3]}), + (A, DROP_NULLS, {"a": [1, 2, 3, 4, 5, 6, None], "i": [1, 1, 1, 3, 3, 3, 4]}), + (A, DROP_BOTH, {"a": [1, 2, 3, 4, 5, 6], "i": [1, 1, 1, 3, 3, 3]}), + ( + BA, + {}, + { + "b": [None, "dog", "cat", None, "narwhal", None, "orca", None], + "a": [1, 2, 3, None, 4, 5, 6, None], + "i": [1, 1, 1, 2, 3, 3, 3, 4], + }, + ), + ( + BA, DROP_EMPTY, - {"a": [1, 2, 3, None, 4, 5, 6], "b": [1, 1, 1, 2, 3, 3, 3]}, - marks=XFAIL_NOT_IMPL, + { + "b": [None, "dog", "cat", None, "narwhal", None, "orca"], + "a": [1, 2, 3, None, 4, 5, 6], + "i": [1, 1, 1, 2, 3, 3, 3], + }, ), - pytest.param( + ( + BA, DROP_NULLS, - {"a": [1, 2, 3, 4, 5, 6, None], "b": [1, 1, 1, 3, 3, 3, 4]}, - marks=XFAIL_NOT_IMPL, + { + "b": [None, "dog", "cat", "narwhal", None, "orca", None], + "a": [1, 2, 3, 4, 5, 6, None], + "i": [1, 1, 1, 3, 3, 3, 4], + }, ), - pytest.param( + ( + BA, DROP_BOTH, - {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 1, 3, 3, 3]}, - marks=XFAIL_NOT_IMPL, + { + "b": [None, "dog", "cat", "narwhal", None, "orca"], + "a": [1, 2, 3, 4, 5, 6], + "i": [1, 1, 1, 3, 3, 3], + }, ), ], ) -def test_explode_frame_options(kwds: dict[str, Any], expected: Data) -> None: - data = {"a": [[1, 2, 3], None, [4, 5, 6], []], "b": [1, 2, 3, 4]} +def test_explode_frame_options( + columns: Sequence[ColumnNameOrSelector], kwds: dict[str, Any], expected: Data +) -> None: # Based on https://github.com/pola-rs/polars/blob/1684cc09dfaa46656dfecc45ab866d01aa69bc78/py-polars/tests/unit/operations/test_explode.py#L596-L616 - result = dataframe(data).explode("a", **kwds) + data = { + "a": [[1, 2, 3], None, [4, 5, 6], []], + "b": [[None, "dog", "cat"], None, ["narwhal", None, "orca"], []], + "i": [1, 2, 3, 4], + } + result = ( + dataframe(data) + .with_columns( + nwp.col("a").cast(nw.List(nw.Int32())), nwp.col("b").cast(nw.List(nw.String)) + ) + .select(*columns, "i") + .explode(columns, **kwds) + ) assert_equal_data(result, expected) From 32997bc37daf1d77f728e73b18afa52a649655e2 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 17:59:38 +0000 Subject: [PATCH 14/21] refactor: Looking cleaner --- narwhals/_plan/arrow/dataframe.py | 7 +------ narwhals/_plan/arrow/functions.py | 31 ++++++++++++++++++------------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index 9a6c8c7214..0fcf4282e9 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -168,12 +168,7 @@ def explode(self, subset: Sequence[str], options: ExplodeOptions) -> Self: native = self.native builder = fn.ExplodeBuilder.from_options(options) if len(subset) == 1: - name = subset[0] - exploded, indices = builder.explode_into(native.column(name)) - # TODO @dangotbanned: Might be more efficient to null-out the column, before `gather`? - df = self.gather(indices) if len(indices) != len(self) else self - ser = Series.from_native(exploded, name, version=self.version) - return df.with_series(ser) + return self._with_native(builder.explode_column(native, subset[0])) explodeds, indices = builder.explode_arrays_into( *native.select(list(subset)).columns ) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index db94be3de9..7291c95169 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -438,6 +438,7 @@ def from_options(cls, options: ExplodeOptions, /) -> Self: def explode( self, native: Arrow[ListScalar[DataTypeT]] ) -> ChunkedOrArray[Scalar[DataTypeT]]: + """Explode list elements, expanding one-level into a new array.""" explode = _list_explode_unchecked if self.options.any(): lengths = list_len(native) @@ -450,23 +451,27 @@ def explode( ) return result - # TODO @dangotbanned: likely want to be passing in the native table *somewhere* - # - Rather than returning a tuple that represents a half-done job - def explode_into( - self, native: ChunkedList[DataTypeT] - ) -> tuple[ChunkedArray[Scalar[DataTypeT]], ChunkedI64]: - """Variant of `explode`, providing indices for `DataFrame.explode([column])`.""" - result: ChunkedArray[Scalar[DataTypeT]] + def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: + """Explode a list-typed column in the context of `native`.""" + ca = native.column(column_name) if self.options.any(): - lengths = list_len(native) - needs_replacing = self._predicate(lengths) - safe = self._replace_mask(native, needs_replacing) + safe = self._replace_mask(ca, self._predicate(list_len(ca))) else: - safe = native - result = _list_explode_unchecked(safe) - return result, _list_parent_indices(safe) + safe = ca + exploded = _list_explode_unchecked(safe) + indices = _list_parent_indices(safe) + col_idx = native.schema.get_field_index(column_name) + if len(indices) == len(native): + return native.set_column(col_idx, column_name, exploded) + return ( + native.remove_column(col_idx) + .take(indices) + .add_column(col_idx, column_name, exploded) + ) # TODO @dangotbanned: De-duplicate + # TODO @dangotbanned: likely want to be passing in the native table *somewhere* + # - Rather than returning a tuple that represents a half-done job def explode_arrays_into( self, *arrays: ChunkedList ) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: From 0126f923af4fa51f3a27f3b01dfa7ee4b83d7f26 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 22:07:20 +0000 Subject: [PATCH 15/21] cover non-take, refactor to `explode_columns` --- narwhals/_plan/arrow/dataframe.py | 13 ++---- narwhals/_plan/arrow/functions.py | 69 +++++++++++++++++-------------- tests/plan/explode_test.py | 13 ++++++ 3 files changed, 54 insertions(+), 41 deletions(-) diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index 0fcf4282e9..ecef91b772 100644 --- a/narwhals/_plan/arrow/dataframe.py +++ b/narwhals/_plan/arrow/dataframe.py @@ -18,7 +18,7 @@ from narwhals._plan.compliant.typing import namespace from narwhals._plan.exceptions import shape_error from narwhals._plan.expressions import NamedIR -from narwhals._utils import Version, generate_repr, zip_strict +from narwhals._utils import Version, generate_repr from narwhals.schema import Schema if TYPE_CHECKING: @@ -163,18 +163,11 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self: native = self.native.filter(~to_drop) return self._with_native(native) - # TODO @dangotbanned: Move move of this into `ExplodeBuilder` def explode(self, subset: Sequence[str], options: ExplodeOptions) -> Self: - native = self.native builder = fn.ExplodeBuilder.from_options(options) if len(subset) == 1: - return self._with_native(builder.explode_column(native, subset[0])) - explodeds, indices = builder.explode_arrays_into( - *native.select(list(subset)).columns - ) - df = self.gather(indices) if len(indices) != len(self) else self - names_and_columns = zip_strict(subset, explodeds) - return self._with_native(with_arrays(df.native, names_and_columns)) + return self._with_native(builder.explode_column(self.native, subset[0])) + return self._with_native(builder.explode_columns(self.native, subset)) def rename(self, mapping: Mapping[str, str]) -> Self: names: dict[str, str] | list[str] diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index 7291c95169..1234a5ab66 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -4,8 +4,8 @@ import math import typing as t -from collections import deque -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterator, Sequence +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload import pyarrow as pa # ignore-banned-import @@ -441,9 +441,7 @@ def explode( """Explode list elements, expanding one-level into a new array.""" explode = _list_explode_unchecked if self.options.any(): - lengths = list_len(native) - needs_replacing = self._predicate(lengths) - safe = self._replace_mask(native, needs_replacing) + safe = self._replace_mask(native, self._predicate(list_len(native))) else: safe = native result: ChunkedOrArray[Scalar[DataTypeT]] = ( @@ -459,45 +457,54 @@ def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: else: safe = ca exploded = _list_explode_unchecked(safe) - indices = _list_parent_indices(safe) col_idx = native.schema.get_field_index(column_name) - if len(indices) == len(native): + if len(exploded) == len(native): return native.set_column(col_idx, column_name, exploded) return ( native.remove_column(col_idx) - .take(indices) + .take(_list_parent_indices(safe)) .add_column(col_idx, column_name, exploded) ) - # TODO @dangotbanned: De-duplicate - # TODO @dangotbanned: likely want to be passing in the native table *somewhere* - # - Rather than returning a tuple that represents a half-done job - def explode_arrays_into( - self, *arrays: ChunkedList - ) -> tuple[Sequence[ChunkedArrayAny], ChunkedI64]: - """Variant of `explode_into`, with shape checking against the first array.""" + def explode_columns(self, native: pa.Table, subset: Sequence[str], /) -> pa.Table: + """Explode multiple list-typed columns in the context of `native`.""" + arrays = native.select(list(subset)).columns explode = _list_explode_unchecked first = arrays[0] first_len = list_len(first) - results = deque["ChunkedArrayAny"]() if self.options.any(): - needs_replacing = self._predicate(first_len) - first_safe = self._replace_mask(first, needs_replacing) - results.append(explode(first_safe)) - for arr in arrays[1:]: - if not first_len.equals(list_len(arr)): - msg = "exploded columns must have matching element counts" - raise ShapeError(msg) - results.append(explode(self._replace_mask(arr, needs_replacing))) + mask = self._predicate(first_len) + first_safe = self._replace_mask(first, mask) + it = ( + explode(self._replace_mask(arr, mask)) + for arr in self._iter_ensure_shape(first_len, arrays[1:]) + ) else: first_safe = first - results.append(explode(first_safe)) - for arr in arrays[1:]: - if not first_len.equals(list_len(arr)): - msg = "exploded columns must have matching element counts" - raise ShapeError(msg) - results.append(explode(arr)) - return results, _list_parent_indices(first_safe) + it = (explode(arr) for arr in self._iter_ensure_shape(first_len, arrays[1:])) + first_result = explode(first_safe) + + # NOTE: Not great ... + from narwhals._plan.arrow.dataframe import with_arrays + + if len(first_result) != len(native): + # TODO @dangotbanned: Try to avoid repeating the lists in take + # Complicated by needing the columns in the same position + native = native.take(_list_parent_indices(first_safe)) + + return with_arrays(native, zip(subset, chain([first_result], it))) + + def _iter_ensure_shape( + self, + first_len: ChunkedArray[pa.UInt32Scalar], + other_arrays: Iterable[ChunkedArrayAny], + /, + ) -> Iterator[ChunkedArrayAny]: + for arr in other_arrays: + if not first_len.equals(list_len(arr)): + msg = "exploded columns must have matching element counts" + raise ShapeError(msg) + yield arr def _predicate(self, lengths: ArrowAny) -> Arrow[BooleanScalar]: empty_as_null, keep_nulls = self.options.empty_as_null, self.options.keep_nulls diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index ecffea2818..18bf3c00a0 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -267,3 +267,16 @@ def test_explode_frame_options( .explode(columns, **kwds) ) assert_equal_data(result, expected) + + +def test_explode_frame_single_elements() -> None: + data = {"a": [[1], [2], [3]], "b": [[4], [5], [6]], "i": [0, 10, 20]} + df = dataframe(data).with_columns(nwp.col("a", "b").cast(nw.List(nw.Int32()))) + + result = df.explode("a") + expected = {"a": [1, 2, 3], "b": [[4], [5], [6]], "i": [0, 10, 20]} + assert_equal_data(result, expected) + + result = df.explode("b", "a") + expected = {"a": [1, 2, 3], "b": [4, 5, 6], "i": [0, 10, 20]} + assert_equal_data(result, expected) From 8c420d0d6a5023a7914c767ea90f31df393ef3a8 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 22:17:37 +0000 Subject: [PATCH 16/21] "fix" typing --- narwhals/_plan/arrow/functions.py | 46 +++++++++++++++++++------------ narwhals/_plan/arrow/typing.py | 12 ++++++++ 2 files changed, 41 insertions(+), 17 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index 1234a5ab66..af1fdadc34 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -4,7 +4,7 @@ import math import typing as t -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Collection, Iterator, Sequence from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload @@ -68,6 +68,7 @@ ListArray, ListScalar, NativeScalar, + NonListType, NumericScalar, Predicate, SameArrowT, @@ -421,6 +422,8 @@ def list_explode( _ArrowListT = TypeVar("_ArrowListT", bound="Arrow[ListScalar[Any]]") +_NonListT = TypeVar("_NonListT", bound="NonListType") +_ListT = TypeVar("_ListT", bound="pa.ListType[Any]") class ExplodeBuilder: @@ -439,15 +442,13 @@ def explode( self, native: Arrow[ListScalar[DataTypeT]] ) -> ChunkedOrArray[Scalar[DataTypeT]]: """Explode list elements, expanding one-level into a new array.""" - explode = _list_explode_unchecked if self.options.any(): safe = self._replace_mask(native, self._predicate(list_len(native))) else: safe = native - result: ChunkedOrArray[Scalar[DataTypeT]] = ( - chunked_array(explode(safe)) if isinstance(safe, pa.Scalar) else explode(safe) - ) - return result + if not isinstance(safe, pa.Scalar): + return _list_explode(safe) + return chunked_array(_list_explode(safe)) def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: """Explode a list-typed column in the context of `native`.""" @@ -456,7 +457,7 @@ def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: safe = self._replace_mask(ca, self._predicate(list_len(ca))) else: safe = ca - exploded = _list_explode_unchecked(safe) + exploded = _list_explode(safe) col_idx = native.schema.get_field_index(column_name) if len(exploded) == len(native): return native.set_column(col_idx, column_name, exploded) @@ -466,25 +467,26 @@ def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: .add_column(col_idx, column_name, exploded) ) - def explode_columns(self, native: pa.Table, subset: Sequence[str], /) -> pa.Table: + def explode_columns(self, native: pa.Table, subset: Collection[str], /) -> pa.Table: """Explode multiple list-typed columns in the context of `native`.""" arrays = native.select(list(subset)).columns - explode = _list_explode_unchecked first = arrays[0] first_len = list_len(first) if self.options.any(): mask = self._predicate(first_len) first_safe = self._replace_mask(first, mask) it = ( - explode(self._replace_mask(arr, mask)) + _list_explode(self._replace_mask(arr, mask)) for arr in self._iter_ensure_shape(first_len, arrays[1:]) ) else: first_safe = first - it = (explode(arr) for arr in self._iter_ensure_shape(first_len, arrays[1:])) - first_result = explode(first_safe) - - # NOTE: Not great ... + it = ( + _list_explode(arr) + for arr in self._iter_ensure_shape(first_len, arrays[1:]) + ) + first_result = _list_explode(first_safe) + # NOTE: Not too happy about this import from narwhals._plan.arrow.dataframe import with_arrays if len(first_result) != len(native): @@ -497,10 +499,10 @@ def explode_columns(self, native: pa.Table, subset: Sequence[str], /) -> pa.Tabl def _iter_ensure_shape( self, first_len: ChunkedArray[pa.UInt32Scalar], - other_arrays: Iterable[ChunkedArrayAny], + arrays: Iterable[ChunkedArrayAny], /, ) -> Iterator[ChunkedArrayAny]: - for arr in other_arrays: + for arr in arrays: if not first_len.equals(list_len(arr)): msg = "exploded columns must have matching element counts" raise ShapeError(msg) @@ -520,7 +522,17 @@ def _replace_mask( return when_then(mask, lit([None], native.type), native) # type: ignore[no-any-return] -def _list_explode_unchecked(native: Incomplete) -> Incomplete: +@t.overload +def _list_explode(native: ChunkedList[DataTypeT]) -> ChunkedArray[Scalar[DataTypeT]]: ... +@t.overload +def _list_explode( + native: ListArray[_NonListT] | ListScalar[_NonListT], +) -> Array[Scalar[_NonListT]]: ... +@t.overload +def _list_explode(native: ListArray[DataTypeT]) -> Array[Scalar[DataTypeT]]: ... +@t.overload +def _list_explode(native: ListScalar[_ListT]) -> ListArray[_ListT]: ... +def _list_explode(native: Incomplete) -> Incomplete: return pc.call_function("list_flatten", [native]) diff --git a/narwhals/_plan/arrow/typing.py b/narwhals/_plan/arrow/typing.py index 047e436eb4..f16b6365e4 100644 --- a/narwhals/_plan/arrow/typing.py +++ b/narwhals/_plan/arrow/typing.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: import pyarrow as pa import pyarrow.compute as pc + from pyarrow import lib, types from pyarrow.lib import ( BoolType as BoolType, Date32Type, @@ -37,6 +38,17 @@ BooleanScalar: TypeAlias = "Scalar[BoolType]" NumericScalar: TypeAlias = "pc.NumericScalar" + PrimitiveNumericType: TypeAlias = "types._Integer | types._Floating" + NumericType: TypeAlias = "PrimitiveNumericType | types._Decimal" + NumericOrTemporalType: TypeAlias = "NumericType | types._Temporal" + StringOrBinaryType: TypeAlias = "StringType | LargeStringType | lib.StringViewType | lib.BinaryType | lib.LargeBinaryType | lib.BinaryViewType" + BasicType: TypeAlias = ( + "NumericOrTemporalType | StringOrBinaryType | BoolType | lib.NullType" + ) + NonListNestedType: TypeAlias = "pa.StructType | pa.DictionaryType[Any, Any] | pa.MapType[Any, Any] | pa.UnionType" + NonListType: TypeAlias = "BasicType | NonListNestedType" + NestedType: TypeAlias = "NonListNestedType | pa.ListType[Any]" + class NativeArrowSeries(NativeSeries, Protocol): @property def chunks(self) -> list[Any]: ... From 2510d571d538331e209f676417692ee1fa160d67 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Mon, 8 Dec 2025 22:51:01 +0000 Subject: [PATCH 17/21] refactor: Just one version, thanks --- narwhals/_plan/arrow/functions.py | 45 ++++++++++--------------------- narwhals/_plan/arrow/series.py | 5 ++-- 2 files changed, 16 insertions(+), 34 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index af1fdadc34..a5df10111c 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -391,36 +391,6 @@ def get_categories(native: ArrowAny) -> ChunkedArrayAny: return chunked_array(da.dictionary) -@t.overload -def list_explode( - native: ChunkedList[DataTypeT] | ListScalar[DataTypeT], - *, - empty_as_null: bool = ..., - keep_nulls: bool = ..., -) -> ChunkedArray[Scalar[DataTypeT]]: ... -@t.overload -def list_explode( - native: ListArray[DataTypeT], *, empty_as_null: bool = ..., keep_nulls: bool = ... -) -> Array[Scalar[DataTypeT]]: ... -@t.overload -def list_explode( - native: Arrow[ListScalar[DataTypeT]], - *, - empty_as_null: bool = ..., - keep_nulls: bool = ..., -) -> ChunkedOrArray[Scalar[DataTypeT]]: ... -def list_explode( - native: ArrowAny, *, empty_as_null: bool = True, keep_nulls: bool = True -) -> ChunkedOrArray[Scalar[DataTypeT]]: - """Explode list elements, expanding one-level into a new array. - - Equivalent to `polars.{Expr,Series}.explode`. - """ - return ExplodeBuilder(empty_as_null=empty_as_null, keep_nulls=keep_nulls).explode( - native - ) - - _ArrowListT = TypeVar("_ArrowListT", bound="Arrow[ListScalar[Any]]") _NonListT = TypeVar("_NonListT", bound="NonListType") _ListT = TypeVar("_ListT", bound="pa.ListType[Any]") @@ -438,10 +408,23 @@ def from_options(cls, options: ExplodeOptions, /) -> Self: obj.options = options return obj + @t.overload + def explode( + self, native: ChunkedList[DataTypeT] | ListScalar[DataTypeT] + ) -> ChunkedArray[Scalar[DataTypeT]]: ... + @t.overload + def explode(self, native: ListArray[DataTypeT]) -> Array[Scalar[DataTypeT]]: ... + @t.overload + def explode( + self, native: Arrow[ListScalar[DataTypeT]] + ) -> ChunkedOrArray[Scalar[DataTypeT]]: ... def explode( self, native: Arrow[ListScalar[DataTypeT]] ) -> ChunkedOrArray[Scalar[DataTypeT]]: - """Explode list elements, expanding one-level into a new array.""" + """Explode list elements, expanding one-level into a new array. + + Equivalent to `polars.{Expr,Series}.explode`. + """ if self.options.any(): safe = self._replace_mask(native, self._predicate(list_len(native))) else: diff --git a/narwhals/_plan/arrow/series.py b/narwhals/_plan/arrow/series.py index 3adebde9c1..8a59ce42cc 100644 --- a/narwhals/_plan/arrow/series.py +++ b/narwhals/_plan/arrow/series.py @@ -302,9 +302,8 @@ def drop_nans(self) -> Self: ) def explode(self, *, empty_as_null: bool = True, keep_nulls: bool = True) -> Self: - ca = self.native - result = fn.list_explode(ca, empty_as_null=empty_as_null, keep_nulls=keep_nulls) - return self._with_native(result) + exploder = fn.ExplodeBuilder(empty_as_null=empty_as_null, keep_nulls=keep_nulls) + return self._with_native(exploder.explode(self.native)) @property def struct(self) -> SeriesStructNamespace: From 59d4f3892e7a1899176409942289f25b5450777c Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Tue, 9 Dec 2025 12:25:01 +0000 Subject: [PATCH 18/21] perf: Drop list columns before `take` https://github.com/narwhals-dev/narwhals/pull/3347#discussion_r2602378565 --- narwhals/_plan/arrow/functions.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index a5df10111c..c9524a6605 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -452,7 +452,8 @@ def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: def explode_columns(self, native: pa.Table, subset: Collection[str], /) -> pa.Table: """Explode multiple list-typed columns in the context of `native`.""" - arrays = native.select(list(subset)).columns + subset = list(subset) + arrays = native.select(subset).columns first = arrays[0] first_len = list_len(first) if self.options.any(): @@ -469,14 +470,14 @@ def explode_columns(self, native: pa.Table, subset: Collection[str], /) -> pa.Ta for arr in self._iter_ensure_shape(first_len, arrays[1:]) ) first_result = _list_explode(first_safe) + if len(first_result) != len(native): + gathered = native.drop_columns(subset).take(_list_parent_indices(first_safe)) + for name, arr in zip(subset, chain([first_result], it)): + gathered = gathered.append_column(name, arr) + return gathered.select(native.column_names) # NOTE: Not too happy about this import from narwhals._plan.arrow.dataframe import with_arrays - if len(first_result) != len(native): - # TODO @dangotbanned: Try to avoid repeating the lists in take - # Complicated by needing the columns in the same position - native = native.take(_list_parent_indices(first_safe)) - return with_arrays(native, zip(subset, chain([first_result], it))) def _iter_ensure_shape( From 164c8c878c3c35120c39900feed20101bfa4eb41 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Tue, 9 Dec 2025 13:18:49 +0000 Subject: [PATCH 19/21] fill/fix typing holes --- narwhals/_plan/arrow/functions.py | 32 ++++++++++++++----------------- narwhals/_plan/arrow/typing.py | 4 +++- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index c9524a6605..1e241f7c69 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -6,7 +6,7 @@ import typing as t from collections.abc import Callable, Collection, Iterator, Sequence from itertools import chain -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload +from typing import TYPE_CHECKING, Any, Final, Literal, overload import pyarrow as pa # ignore-banned-import import pyarrow.compute as pc # ignore-banned-import @@ -38,6 +38,7 @@ ArrayAny, Arrow, ArrowAny, + ArrowListT, ArrowT, BinaryComp, BinaryFunction, @@ -49,7 +50,6 @@ BoolType, ChunkedArray, ChunkedArrayAny, - ChunkedI64, ChunkedList, ChunkedOrArray, ChunkedOrArrayAny, @@ -67,8 +67,9 @@ LargeStringType, ListArray, ListScalar, + ListTypeT, NativeScalar, - NonListType, + NonListTypeT, NumericScalar, Predicate, SameArrowT, @@ -391,11 +392,6 @@ def get_categories(native: ArrowAny) -> ChunkedArrayAny: return chunked_array(da.dictionary) -_ArrowListT = TypeVar("_ArrowListT", bound="Arrow[ListScalar[Any]]") -_NonListT = TypeVar("_NonListT", bound="NonListType") -_ListT = TypeVar("_ListT", bound="pa.ListType[Any]") - - class ExplodeBuilder: options: ExplodeOptions @@ -500,28 +496,28 @@ def _predicate(self, lengths: ArrowAny) -> Arrow[BooleanScalar]: return eq(lengths, lit(0)) return is_null(lengths) - def _replace_mask( - self, native: _ArrowListT, mask: Arrow[BooleanScalar] - ) -> _ArrowListT: - return when_then(mask, lit([None], native.type), native) # type: ignore[no-any-return] + def _replace_mask(self, native: ArrowListT, mask: Arrow[BooleanScalar]) -> ArrowListT: + result: ArrowListT = when_then(mask, lit([None], native.type), native) + return result @t.overload def _list_explode(native: ChunkedList[DataTypeT]) -> ChunkedArray[Scalar[DataTypeT]]: ... @t.overload def _list_explode( - native: ListArray[_NonListT] | ListScalar[_NonListT], -) -> Array[Scalar[_NonListT]]: ... + native: ListArray[NonListTypeT] | ListScalar[NonListTypeT], +) -> Array[Scalar[NonListTypeT]]: ... @t.overload def _list_explode(native: ListArray[DataTypeT]) -> Array[Scalar[DataTypeT]]: ... @t.overload -def _list_explode(native: ListScalar[_ListT]) -> ListArray[_ListT]: ... -def _list_explode(native: Incomplete) -> Incomplete: - return pc.call_function("list_flatten", [native]) +def _list_explode(native: ListScalar[ListTypeT]) -> ListArray[ListTypeT]: ... +def _list_explode(native: Arrow[ListScalar]) -> ChunkedOrArrayAny: + result: ChunkedOrArrayAny = pc.call_function("list_flatten", [native]) + return result @t.overload -def _list_parent_indices(native: ChunkedList) -> ChunkedI64: ... +def _list_parent_indices(native: ChunkedList) -> ChunkedArray[pa.Int64Scalar]: ... @t.overload def _list_parent_indices(native: ListArray) -> pa.Int64Array: ... def _list_parent_indices( diff --git a/narwhals/_plan/arrow/typing.py b/narwhals/_plan/arrow/typing.py index f16b6365e4..766053f76a 100644 --- a/narwhals/_plan/arrow/typing.py +++ b/narwhals/_plan/arrow/typing.py @@ -48,6 +48,8 @@ NonListNestedType: TypeAlias = "pa.StructType | pa.DictionaryType[Any, Any] | pa.MapType[Any, Any] | pa.UnionType" NonListType: TypeAlias = "BasicType | NonListNestedType" NestedType: TypeAlias = "NonListNestedType | pa.ListType[Any]" + NonListTypeT = TypeVar("NonListTypeT", bound="NonListType") + ListTypeT = TypeVar("ListTypeT", bound="pa.ListType[Any]") class NativeArrowSeries(NativeSeries, Protocol): @property @@ -207,12 +209,12 @@ class BinaryLogical(BinaryFunction["BooleanScalar", "BooleanScalar"], Protocol): StructArray: TypeAlias = "pa.StructArray | Array[pa.StructScalar]" ChunkedList: TypeAlias = "ChunkedArray[ListScalar[DataTypeT_co]]" ListArray: TypeAlias = "Array[ListScalar[DataTypeT_co]]" -ChunkedI64: TypeAlias = "ChunkedArray[pa.Int64Scalar]" Arrow: TypeAlias = "ChunkedOrScalar[ScalarT_co] | Array[ScalarT_co]" ArrowAny: TypeAlias = "ChunkedOrScalarAny | ArrayAny" SameArrowT = TypeVar("SameArrowT", ChunkedArrayAny, ArrayAny, ScalarAny) ArrowT = TypeVar("ArrowT", bound=ArrowAny) +ArrowListT = TypeVar("ArrowListT", bound="Arrow[ListScalar[Any]]") Predicate: TypeAlias = "Arrow[BooleanScalar]" """Any `pyarrow` container that wraps boolean.""" From 81b114a9aa70e3d01112c210023abb63a94a6269 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Tue, 9 Dec 2025 13:31:14 +0000 Subject: [PATCH 20/21] test: Shrink cases for `test_explode_frame_options` Also makes it clearer which cases are related in an IDE --- tests/plan/explode_test.py | 62 ++++++++++++++------------------------ 1 file changed, 22 insertions(+), 40 deletions(-) diff --git a/tests/plan/explode_test.py b/tests/plan/explode_test.py index 18bf3c00a0..c7cc478d67 100644 --- a/tests/plan/explode_test.py +++ b/tests/plan/explode_test.py @@ -180,6 +180,7 @@ def test_explode_series_default_masked(values: list[Any], expected: list[Any]) - DROP_EMPTY: Final = {"empty_as_null": False} DROP_NULLS: Final = {"keep_nulls": False} DROP_BOTH: Final = {"empty_as_null": False, "keep_nulls": False} +DEFAULT: Final[Data] = {} @pytest.mark.parametrize( @@ -203,50 +204,31 @@ def test_explode_series_options( A = ("a",) BA = "b", "a" +DEFAULT_A: Final = [1, 2, 3, None, 4, 5, 6, None] +DEFAULT_I: Final = [1, 1, 1, 2, 3, 3, 3, 4] +DEFAULT_B: Final = [None, "dog", "cat", None, "narwhal", None, "orca", None] +EMPTY_A: Final = [1, 2, 3, None, 4, 5, 6] +EMPTY_I: Final = [1, 1, 1, 2, 3, 3, 3] +EMPTY_B: Final = [None, "dog", "cat", None, "narwhal", None, "orca"] +NULLS_A: Final = [1, 2, 3, 4, 5, 6, None] +NULLS_I: Final = [1, 1, 1, 3, 3, 3, 4] +NULLS_B: Final = [None, "dog", "cat", "narwhal", None, "orca", None] +BOTH_A: Final = [1, 2, 3, 4, 5, 6] +BOTH_I: Final = [1, 1, 1, 3, 3, 3] +BOTH_B: Final = [None, "dog", "cat", "narwhal", None, "orca"] + @pytest.mark.parametrize( ("columns", "kwds", "expected"), [ - (A, {}, {"a": [1, 2, 3, None, 4, 5, 6, None], "i": [1, 1, 1, 2, 3, 3, 3, 4]}), - (A, DROP_EMPTY, {"a": [1, 2, 3, None, 4, 5, 6], "i": [1, 1, 1, 2, 3, 3, 3]}), - (A, DROP_NULLS, {"a": [1, 2, 3, 4, 5, 6, None], "i": [1, 1, 1, 3, 3, 3, 4]}), - (A, DROP_BOTH, {"a": [1, 2, 3, 4, 5, 6], "i": [1, 1, 1, 3, 3, 3]}), - ( - BA, - {}, - { - "b": [None, "dog", "cat", None, "narwhal", None, "orca", None], - "a": [1, 2, 3, None, 4, 5, 6, None], - "i": [1, 1, 1, 2, 3, 3, 3, 4], - }, - ), - ( - BA, - DROP_EMPTY, - { - "b": [None, "dog", "cat", None, "narwhal", None, "orca"], - "a": [1, 2, 3, None, 4, 5, 6], - "i": [1, 1, 1, 2, 3, 3, 3], - }, - ), - ( - BA, - DROP_NULLS, - { - "b": [None, "dog", "cat", "narwhal", None, "orca", None], - "a": [1, 2, 3, 4, 5, 6, None], - "i": [1, 1, 1, 3, 3, 3, 4], - }, - ), - ( - BA, - DROP_BOTH, - { - "b": [None, "dog", "cat", "narwhal", None, "orca"], - "a": [1, 2, 3, 4, 5, 6], - "i": [1, 1, 1, 3, 3, 3], - }, - ), + (A, DEFAULT, {"a": DEFAULT_A, "i": DEFAULT_I}), + (A, DROP_EMPTY, {"a": EMPTY_A, "i": EMPTY_I}), + (A, DROP_NULLS, {"a": NULLS_A, "i": NULLS_I}), + (A, DROP_BOTH, {"a": BOTH_A, "i": BOTH_I}), + (BA, DEFAULT, {"b": DEFAULT_B, "a": DEFAULT_A, "i": DEFAULT_I}), + (BA, DROP_EMPTY, {"b": EMPTY_B, "a": EMPTY_A, "i": EMPTY_I}), + (BA, DROP_NULLS, {"b": NULLS_B, "a": NULLS_A, "i": NULLS_I}), + (BA, DROP_BOTH, {"b": BOTH_B, "a": BOTH_A, "i": BOTH_I}), ], ) def test_explode_frame_options( From e86c66dd5976fdbf24b68eff58bd4878b909d227 Mon Sep 17 00:00:00 2001 From: dangotbanned <125183946+dangotbanned@users.noreply.github.com> Date: Tue, 9 Dec 2025 14:05:35 +0000 Subject: [PATCH 21/21] refactor: Clean up and docs --- narwhals/_plan/arrow/functions.py | 37 ++++++++++++++++++------------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/narwhals/_plan/arrow/functions.py b/narwhals/_plan/arrow/functions.py index 1e241f7c69..1969cd33d3 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -22,7 +22,7 @@ from narwhals._plan.arrow import options as pa_options from narwhals._plan.expressions import functions as F, operators as ops from narwhals._plan.options import ExplodeOptions -from narwhals._utils import Implementation, Version +from narwhals._utils import Implementation, Version, no_default from narwhals.exceptions import ShapeError if TYPE_CHECKING: @@ -86,6 +86,7 @@ from narwhals._plan.compliant.typing import SeriesT from narwhals._plan.options import RankOptions, SortMultipleOptions, SortOptions from narwhals._plan.typing import Seq + from narwhals._typing import NoDefault from narwhals.typing import ( ClosedInterval, FillNullStrategy, @@ -421,10 +422,7 @@ def explode( Equivalent to `polars.{Expr,Series}.explode`. """ - if self.options.any(): - safe = self._replace_mask(native, self._predicate(list_len(native))) - else: - safe = native + safe = self._fill_with_null(native) if self.options.any() else native if not isinstance(safe, pa.Scalar): return _list_explode(safe) return chunked_array(_list_explode(safe)) @@ -432,10 +430,7 @@ def explode( def explode_column(self, native: pa.Table, column_name: str, /) -> pa.Table: """Explode a list-typed column in the context of `native`.""" ca = native.column(column_name) - if self.options.any(): - safe = self._replace_mask(ca, self._predicate(list_len(ca))) - else: - safe = ca + safe = self._fill_with_null(ca) if self.options.any() else ca exploded = _list_explode(safe) col_idx = native.schema.get_field_index(column_name) if len(exploded) == len(native): @@ -454,9 +449,9 @@ def explode_columns(self, native: pa.Table, subset: Collection[str], /) -> pa.Ta first_len = list_len(first) if self.options.any(): mask = self._predicate(first_len) - first_safe = self._replace_mask(first, mask) + first_safe = self._fill_with_null(first, mask) it = ( - _list_explode(self._replace_mask(arr, mask)) + _list_explode(self._fill_with_null(arr, mask)) for arr in self._iter_ensure_shape(first_len, arrays[1:]) ) else: @@ -488,7 +483,8 @@ def _iter_ensure_shape( raise ShapeError(msg) yield arr - def _predicate(self, lengths: ArrowAny) -> Arrow[BooleanScalar]: + def _predicate(self, lengths: ArrowAny, /) -> Arrow[BooleanScalar]: + """Return True for each sublist length that indicates the original sublist should be replaced with `[None]`.""" empty_as_null, keep_nulls = self.options.empty_as_null, self.options.keep_nulls if empty_as_null and keep_nulls: return or_(is_null(lengths), eq(lengths, lit(0))) @@ -496,8 +492,17 @@ def _predicate(self, lengths: ArrowAny) -> Arrow[BooleanScalar]: return eq(lengths, lit(0)) return is_null(lengths) - def _replace_mask(self, native: ArrowListT, mask: Arrow[BooleanScalar]) -> ArrowListT: - result: ArrowListT = when_then(mask, lit([None], native.type), native) + def _fill_with_null( + self, native: ArrowListT, mask: Arrow[BooleanScalar] | NoDefault = no_default + ) -> ArrowListT: + """Replace each sublist in `native` with `[None]`, according to `self.options`. + + Arguments: + native: List-typed arrow data. + mask: An optional, pre-computed replacement mask. By default, this is generated from `native`. + """ + predicate = self._predicate(list_len(native)) if mask is no_default else mask + result: ArrowListT = when_then(predicate, lit([None], native.type), native) return result @@ -537,9 +542,9 @@ def list_len(native: ListArray) -> pa.UInt32Array: ... @t.overload def list_len(native: ListScalar) -> pa.UInt32Scalar: ... @t.overload -def list_len(native: SameArrowT) -> SameArrowT: ... -@t.overload def list_len(native: ChunkedOrScalar[ListScalar]) -> ChunkedOrScalar[pa.UInt32Scalar]: ... +@t.overload +def list_len(native: Arrow[ListScalar[Any]]) -> Arrow[pa.UInt32Scalar]: ... def list_len(native: ArrowAny) -> ArrowAny: length: Incomplete = pc.list_value_length result: ArrowAny = length(native).cast(pa.uint32())