feat: support list.contains for PyArrow, pandas and Dask

src/narwhals/_arrow/series_list.py

@@ -5,18 +5,22 @@
 import pyarrow as pa
 import pyarrow.compute as pc
 
-from narwhals._arrow.utils import ArrowSeriesNamespace, list_agg, list_sort
+from narwhals._arrow.utils import ArrowSeriesNamespace, list_agg, list_contains, list_sort
 from narwhals._compliant.any_namespace import ListNamespace
 from narwhals._utils import not_implemented
 
 if TYPE_CHECKING:
     from narwhals._arrow.series import ArrowSeries
+    from narwhals.typing import NonNestedLiteral
 
 
 class ArrowSeriesListNamespace(ArrowSeriesNamespace, ListNamespace["ArrowSeries"]):
     def len(self) -> ArrowSeries:
         return self.with_native(pc.list_value_length(self.native).cast(pa.uint32()))
 
+    def contains(self, item: NonNestedLiteral) -> ArrowSeries:
+        return self.with_native(list_contains(self.native, item))
+
     def get(self, index: int) -> ArrowSeries:
         return self.with_native(pc.list_element(self.native, index))
 
@@ -41,4 +45,3 @@ def sort(self, *, descending: bool, nulls_last: bool) -> ArrowSeries:
         )
 
     unique = not_implemented()
-    contains = not_implemented()

src/narwhals/_arrow/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any, cast
 
@@ -33,11 +34,12 @@
     )
     from narwhals._duration import IntervalUnit
     from narwhals.dtypes import DType
-    from narwhals.typing import IntoDType, PythonLiteral
+    from narwhals.typing import IntoDType, NonNestedLiteral, PythonLiteral
 
     # NOTE: stubs don't allow for `ChunkedArray[StructArray]`
     # Intended to represent the `.chunks` property storing `list[pa.StructArray]`
     ChunkedArrayStructArray: TypeAlias = ChunkedArrayAny
+    ListArrayAny: TypeAlias = "pa.ListArray[Any] | pa.LargeListArray[Any]"
 
     def is_timestamp(t: Any) -> TypeIs[pa.TimestampType[Any, Any]]: ...
     def is_duration(t: Any) -> TypeIs[pa.DurationType[Any]]: ...
@@ -563,6 +565,82 @@ def list_agg(
     )
 
 
+_LIST_BLOCK_VALUES = 1 << 21
+_LIST_MIN_BLOCK_VALUES = _LIST_BLOCK_VALUES >> 4
+
+
+def list_contains(array: ChunkedArrayAny, item: NonNestedLiteral) -> ChunkedArrayAny:
+    """Whether each list holds `item`, or a null element if `item` is None.
+
+    A running count of matches over the flattened values changes within a list iff the
+    list holds a match. That keeps this linear, where a group-by per list would sort.
+    """
+    blocks = [_list_block_contains(block, item) for block in _list_blocks(array)]
+    return pa.chunked_array(blocks, pa.bool_())
+
+
+def _list_blocks(array: ChunkedArrayAny) -> Iterator[ListArrayAny]:
+    """Split `array` into runs of whole lists, of up to about `_LIST_BLOCK_VALUES` values.
+
+    That keeps the intermediates in cache. Large chunks are sliced without copying, and
+    chunks under `_LIST_MIN_BLOCK_VALUES` values are combined until a run reaches it, to
+    amortise the ~35us of pyarrow calls per block. Combining every chunk could overflow
+    32-bit list offsets.
+    """
+    run: list[ListArrayAny] = []
+    run_values = 0
+    for chunk in cast("list[ListArrayAny]", array.chunks):
+        if not len(chunk):
+            continue
+        n_values = chunk.offsets[-1].as_py() - chunk.offsets[0].as_py()
+        is_small = n_values < _LIST_MIN_BLOCK_VALUES
+        if run and (not is_small or run_values >= _LIST_MIN_BLOCK_VALUES):
+            yield pa.concat_arrays(run) if len(run) > 1 else run[0]
+            run, run_values = [], 0
+        if is_small:
+            run.append(chunk)
+            run_values += n_values
+        else:
+            step = max(1, len(chunk) * _LIST_BLOCK_VALUES // n_values)
+            yield from (chunk.slice(start, step) for start in range(0, len(chunk), step))
+    if run:
+        yield pa.concat_arrays(run) if len(run) > 1 else run[0]
+
+
+def _list_block_contains(block: ListArrayAny, item: NonNestedLiteral) -> pa.BooleanArray:
+    offsets = block.offsets
+    first = offsets[0].as_py()
+    values = block.values.slice(first, offsets[-1].as_py() - first)
+    # A plain `int` would widen int32 offsets to int64, which is slower to `take`.
+    offsets = pc.subtract(offsets, lit(first, offsets.type))  # type: ignore[arg-type]
+    ends, starts = offsets.slice(1), offsets.slice(0, len(block))
+
+    if item is None:
+        matches = pc.is_null(values)
+    elif isinstance(item, float) and math.isnan(item):
+        matches = pc.is_nan(values)  # NaN matches NaN, as in Polars.
+    else:
+        matches = pc.equal(values, lit(item))
+
+    # Counting modulo 2^k stays exact within lists shorter than 2^k, so the narrowest
+    # type that fits the longest list is enough.
+    longest = pc.max(pc.subtract(ends, starts)).as_py()
+    count_type = next(
+        t
+        for t in (pa.uint8(), pa.uint16(), pa.uint32(), pa.uint64())
+        if longest < 1 << t.bit_width
+    )
+    # `matches_before[i]` counts the matches before flattened position `i`.
+    matches = pa.concat_arrays([pa.array([False]), matches.fill_null(False)])
+    matches_before = pc.cumulative_sum(matches.cast(count_type))
+    contains = pc.not_equal(matches_before.take(ends), matches_before.take(starts))
+    return (
+        contains
+        if block.null_count == 0
+        else pc.if_else(block.is_valid(), contains, None)
+    )
+
+
 def sortable(array: ChunkedArrayAny, /) -> ChunkedArrayAny:
     """Return an array which orders like `array`, but with a type `pyarrow` can sort.
 

src/narwhals/_dask/expr.py

@@ -7,6 +7,7 @@
 
 from narwhals._compliant import DepthTrackingExpr, LazyExpr
 from narwhals._dask.expr_dt import DaskExprDateTimeNamespace
+from narwhals._dask.expr_list import DaskExprListNamespace
 from narwhals._dask.expr_str import DaskExprStringNamespace
 from narwhals._dask.utils import (
     add_row_index,
@@ -710,6 +711,10 @@ def str(self) -> DaskExprStringNamespace:
     def dt(self) -> DaskExprDateTimeNamespace:
         return DaskExprDateTimeNamespace(self)
 
+    @property
+    def list(self) -> DaskExprListNamespace:
+        return DaskExprListNamespace(self)
+
     any_value = not_implemented()
     filter = not_implemented(
         hint=length_changing_hint(
@@ -722,5 +727,4 @@ def dt(self) -> DaskExprDateTimeNamespace:
     last = not_implemented()
 
     # namespaces
-    list: Any = not_implemented()
     struct: Any = not_implemented()

src/narwhals/_dask/expr_list.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+
+from narwhals._compliant import LazyExprNamespace
+from narwhals._compliant.any_namespace import ListNamespace
+from narwhals._pandas_like.utils import is_dtype_pyarrow
+from narwhals._utils import not_implemented
+
+if TYPE_CHECKING:
+    import dask.dataframe.dask_expr as dx
+
+    from narwhals._dask.dataframe import Incomplete
+    from narwhals._dask.expr import DaskExpr
+    from narwhals.typing import NonNestedLiteral
+
+
+def _list_contains_partition(partition: pd.Series, item: NonNestedLiteral) -> pd.Series:
+    from narwhals._arrow.utils import list_contains
+
+    array: Incomplete = partition.array
+    result = pd.arrays.ArrowExtensionArray(list_contains(array._pa_array, item))
+    return pd.Series(result, index=partition.index, name=partition.name)
+
+
+class DaskExprListNamespace(LazyExprNamespace["DaskExpr"], ListNamespace["DaskExpr"]):
+    def contains(self, item: NonNestedLiteral) -> DaskExpr:
+        def func(expr: dx.Series) -> dx.Series:
+            if not is_dtype_pyarrow(expr.dtype):
+                msg = "Only pyarrow-backed lists are supported for Dask."
+                raise NotImplementedError(msg)
+            return expr.map_partitions(
+                _list_contains_partition, item, meta=(expr.name, "bool[pyarrow]")
+            )
+
+        return self.compliant._with_callable(func)
+
+    len = not_implemented()
+    unique = not_implemented()
+    get = not_implemented()
+    min = not_implemented()
+    max = not_implemented()
+    mean = not_implemented()
+    median = not_implemented()
+    sum = not_implemented()
+    sort = not_implemented()

src/narwhals/_dask/namespace.py

@@ -27,6 +27,7 @@
     Implementation,
     ensure_path_source,
     is_nested_literal,
+    isinstance_or_issubclass,
     not_implemented,
     validate_separators,
 )
@@ -86,6 +87,11 @@ def lit(self, value: NonNestedLiteral, dtype: IntoDType | None) -> DaskExpr:
         if is_nested_literal(value):
             msg = f"Nested structures are not supported for Dask backend, found {type(value).__name__}"
             raise NotImplementedError(msg)
+        if isinstance_or_issubclass(dtype, self._version.dtypes.List):
+            # Dask scalars drop the dtype, so this would be a String column:
+            # https://github.com/dask/dask/issues/11637
+            msg = "`List` literals are not supported for Dask backend."
+            raise NotImplementedError(msg)
 
         def func(df: DaskLazyFrame) -> list[dx.Series]:
             if dtype is not None:

src/narwhals/_dask/utils.py

@@ -4,6 +4,7 @@
 
 from narwhals._pandas_like.utils import (
     make_group_by_kwargs as pd_make_group_by_kwargs,
+    narwhals_to_native_arrow_dtype,
     select_columns_by_name,
 )
 from narwhals._utils import Implementation, Version, isinstance_or_issubclass
@@ -96,7 +97,6 @@ def validate_comparand(lhs: dx.Series, rhs: dx.Series) -> None:
     dtypes_v1.Duration: "timedelta64[ns]",
 }
 UNSUPPORTED_DTYPES = (
-    dtypes.List,
     dtypes.Struct,
     dtypes.Array,
     dtypes.Time,
@@ -127,6 +127,8 @@ def narwhals_to_native_dtype(dtype: IntoDType, version: Version) -> Any:
             return pd.CategoricalDtype(dtype.categories, ordered=True)  # type: ignore[arg-type]
         msg = "Can not cast / initialize Enum without categories present"
         raise ValueError(msg)
+    if isinstance_or_issubclass(dtype, dtypes.List):
+        return narwhals_to_native_arrow_dtype(dtype, Implementation.PANDAS, version)
     if issubclass(base_type, UNSUPPORTED_DTYPES):  # pragma: no cover
         msg = f"Converting to {base_type.__name__} dtype is not supported for Dask."
         raise NotImplementedError(msg)

src/narwhals/_pandas_like/series_list.py

@@ -15,6 +15,7 @@
     from typing import Literal
 
     from narwhals._pandas_like.series import PandasLikeSeries
+    from narwhals.typing import NonNestedLiteral
 
 
 class PandasLikeSeriesListNamespace(
@@ -35,6 +36,16 @@ def len(self) -> PandasLikeSeries:
         )
         return self.with_native(result.astype(dtype)).alias(self.native.name)
 
+    def contains(self, item: NonNestedLiteral) -> PandasLikeSeries:
+        self._raise_if_not_pyarrow_backend()
+
+        from narwhals._arrow.utils import list_contains
+
+        result_native = self.compliant._apply_pyarrow_compute_func(
+            self.native, partial(list_contains, item=item)
+        )
+        return self.with_native(result_native)
+
     def get(self, index: int) -> PandasLikeSeries:
         result = self.native.list[index]
         result.name = self.native.name
@@ -86,4 +97,3 @@ def sort(self, *, descending: bool, nulls_last: bool) -> PandasLikeSeries:
         return self.with_native(result_native)
 
     unique = not_implemented()
-    contains = not_implemented()

tests/conftest.py

@@ -192,16 +192,27 @@ def duckdb_lazy_constructor(obj: dict[str, Any]) -> NativeDuckDB:
     return duckdb.sql("select * from _df")
 
 
-def dask_lazy_p1_constructor(obj: Data) -> NativeDask:  # pragma: no cover
+def _dask_constructor(obj: Data, npartitions: int) -> NativeDask:  # pragma: no cover
     import dask.dataframe as dd
+    import pandas as pd
 
-    return cast("NativeDask", dd.from_dict(obj, npartitions=1))
+    frame = pd.DataFrame(obj)
+    if PANDAS_VERSION >= (2, 2) and find_spec("pyarrow"):
+        import pyarrow as pa
 
+        for name, values in obj.items():
+            if any(isinstance(value, list) for value in values):
+                # Dask's `convert-string` would turn object lists into strings.
+                frame[name] = pd.arrays.ArrowExtensionArray(pa.array(values))
+    return cast("NativeDask", dd.from_pandas(frame, npartitions=npartitions))
 
-def dask_lazy_p2_constructor(obj: Data) -> NativeDask:  # pragma: no cover
-    import dask.dataframe as dd
 
-    return cast("NativeDask", dd.from_dict(obj, npartitions=2))
+def dask_lazy_p1_constructor(obj: Data) -> NativeDask:  # pragma: no cover
+    return _dask_constructor(obj, npartitions=1)
+
+
+def dask_lazy_p2_constructor(obj: Data) -> NativeDask:  # pragma: no cover
+    return _dask_constructor(obj, npartitions=2)
 
 
 def pyarrow_table_constructor(obj: dict[str, Any]) -> pa.Table: