feat(expr-ir): Add {DataFrame,Series}.explode(empty_as_null, keep_nulls)

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

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:

narwhals/_plan/arrow/functions.py

@@ -4,8 +4,9 @@
 
 import math
 import typing as t
-from collections.abc import Callable, Sequence
-from typing import TYPE_CHECKING, Any, Final, Literal, overload
+from collections.abc import Callable, Collection, Iterator, Sequence
+from itertools import chain
+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
@@ -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._plan.options import ExplodeOptions
 from narwhals._utils import Implementation, Version
+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
@@ -46,6 +49,7 @@
         BoolType,
         ChunkedArray,
         ChunkedArrayAny,
+        ChunkedI64,
         ChunkedList,
         ChunkedOrArray,
         ChunkedOrArrayAny,
@@ -64,6 +68,7 @@
         ListArray,
         ListScalar,
         NativeScalar,
+        NonListType,
         NumericScalar,
         Predicate,
         SameArrowT,
@@ -386,6 +391,148 @@ 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
+
+    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`.
+        """
+        if self.options.any():
+            safe = self._replace_mask(native, self._predicate(list_len(native)))
+        else:
+            safe = 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)
+        if self.options.any():
+            safe = self._replace_mask(ca, self._predicate(list_len(ca)))
+        else:
+            safe = 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`."""
+        arrays = native.select(list(subset)).columns
+        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 = (
+                _list_explode(self._replace_mask(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)
+        # 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(
+        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]:
+        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]
+
+
+@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])
+
+
+@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

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)

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]: ...
@@ -195,6 +207,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"

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: ...

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(