diff --git a/narwhals/_plan/arrow/dataframe.py b/narwhals/_plan/arrow/dataframe.py index 1c2f3e1d6c..ecef91b772 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,10 +27,10 @@ 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 + from narwhals._plan.options import ExplodeOptions, SortMultipleOptions from narwhals._plan.typing import NonCrossJoinStrategy from narwhals.dtypes import DType from narwhals.typing import IntoSchema @@ -162,6 +163,12 @@ 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], options: ExplodeOptions) -> Self: + builder = fn.ExplodeBuilder.from_options(options) + if len(subset) == 1: + 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] if fn.BACKEND_VERSION >= (17,): @@ -170,20 +177,26 @@ 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)) - # NOTE: Use instead of `with_columns` for trivial cases + 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) + 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))) @@ -226,3 +239,22 @@ 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: + 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/expr.py b/narwhals/_plan/arrow/expr.py index 14e3ac893e..15ee7e5f91 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 @@ -372,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`" @@ -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..1969cd33d3 100644 --- a/narwhals/_plan/arrow/functions.py +++ b/narwhals/_plan/arrow/functions.py @@ -4,7 +4,8 @@ import math import typing as t -from collections.abc import Callable, Sequence +from collections.abc import Callable, Collection, Iterator, Sequence +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, overload import pyarrow as pa # ignore-banned-import @@ -20,13 +21,15 @@ 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._utils import Implementation, Version +from narwhals._plan.options import ExplodeOptions +from narwhals._utils import Implementation, Version, no_default +from narwhals.exceptions import ShapeError if TYPE_CHECKING: 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 @@ -35,6 +38,7 @@ ArrayAny, Arrow, ArrowAny, + ArrowListT, ArrowT, BinaryComp, BinaryFunction, @@ -63,7 +67,9 @@ LargeStringType, ListArray, ListScalar, + ListTypeT, NativeScalar, + NonListTypeT, NumericScalar, Predicate, SameArrowT, @@ -80,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, @@ -386,6 +393,148 @@ def get_categories(native: ArrowAny) -> ChunkedArrayAny: return chunked_array(da.dictionary) +class ExplodeBuilder: + options: ExplodeOptions + + 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 + + @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. + + Equivalent to `polars.{Expr,Series}.explode`. + """ + 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)) + + 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) + 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): + return native.set_column(col_idx, column_name, exploded) + return ( + native.remove_column(col_idx) + .take(_list_parent_indices(safe)) + .add_column(col_idx, column_name, exploded) + ) + + def explode_columns(self, native: pa.Table, subset: Collection[str], /) -> pa.Table: + """Explode multiple list-typed columns in the context of `native`.""" + subset = list(subset) + arrays = native.select(subset).columns + first = arrays[0] + first_len = list_len(first) + if self.options.any(): + mask = self._predicate(first_len) + first_safe = self._fill_with_null(first, mask) + it = ( + _list_explode(self._fill_with_null(arr, mask)) + for arr in self._iter_ensure_shape(first_len, arrays[1:]) + ) + else: + first_safe = first + it = ( + _list_explode(arr) + 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 + + return with_arrays(native, zip(subset, chain([first_result], it))) + + def _iter_ensure_shape( + self, + first_len: ChunkedArray[pa.UInt32Scalar], + arrays: Iterable[ChunkedArrayAny], + /, + ) -> Iterator[ChunkedArrayAny]: + for arr in 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]: + """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))) + if empty_as_null: + return eq(lengths, lit(0)) + return is_null(lengths) + + 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 + + +@t.overload +def _list_explode(native: ChunkedList[DataTypeT]) -> ChunkedArray[Scalar[DataTypeT]]: ... +@t.overload +def _list_explode( + 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[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) -> ChunkedArray[pa.Int64Scalar]: ... +@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 @@ -393,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()) diff --git a/narwhals/_plan/arrow/series.py b/narwhals/_plan/arrow/series.py index 72e01a4fb1..8a59ce42cc 100644 --- a/narwhals/_plan/arrow/series.py +++ b/narwhals/_plan/arrow/series.py @@ -301,6 +301,10 @@ 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: + 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: return SeriesStructNamespace(self) diff --git a/narwhals/_plan/arrow/typing.py b/narwhals/_plan/arrow/typing.py index 63ee251997..766053f76a 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,19 @@ 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]" + NonListTypeT = TypeVar("NonListTypeT", bound="NonListType") + ListTypeT = TypeVar("ListTypeT", bound="pa.ListType[Any]") + class NativeArrowSeries(NativeSeries, Protocol): @property def chunks(self) -> list[Any]: ... @@ -200,6 +214,7 @@ class BinaryLogical(BinaryFunction["BooleanScalar", "BooleanScalar"], Protocol): 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.""" diff --git a/narwhals/_plan/compliant/dataframe.py b/narwhals/_plan/compliant/dataframe.py index b502848534..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,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], 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/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/dataframe.py b/narwhals/_plan/dataframe.py index 7455ad8cc5..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, @@ -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 @@ -158,6 +158,26 @@ 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, + 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() + 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) + 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( data: Mapping[str, Any], 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/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/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 new file mode 100644 index 0000000000..c7cc478d67 --- /dev/null +++ b/tests/plan/explode_test.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Final + +import pytest + +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, + assert_equal_series, + dataframe, + re_compile, + series, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from narwhals._plan.typing import ColumnNameOrSelector + 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]], + "l5": [[None, None], [None], [99], [83]], + } + + +@pytest.mark.parametrize( + ("column", "expected_values"), + [("l2", [None, 3, None, None, 42]), ("l3", [1, 1, 2, 3, None])], +) +def test_explode_frame_single_col( + column: str, expected_values: list[int | None], data: Data +) -> None: + 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.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_frame_multiple_cols( + column: str, + more_columns: Sequence[str], + expected: dict[str, list[str | int | None]], + data: Data, +) -> None: + 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.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_frame_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_frame_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(ncs.list()) + + +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 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.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: + # 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() + .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} +DEFAULT: Final[Data] = {} + + +@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: + # 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, "") + + +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, 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( + 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 + 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) + + +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)