feat: Implement list.contains for compatible types

narwhals_daft/dataframe.py

@@ -19,6 +19,7 @@
 from narwhals.exceptions import (
     ColumnNotFoundError,
     DuplicateError,
+    InvalidOperationError,
     MultiOutputExpressionError,
 )
 from narwhals.typing import CompliantLazyFrame
@@ -112,6 +113,11 @@ def _evaluate_window_expr(
             raise MultiOutputExpressionError(msg)
         return result[0]
 
+    def _native_dtype(self, expr: Expression) -> daft.DataType:
+        """Resolve the dtype of `expr` against this frame's schema without executing it."""
+        name = generate_temporary_column_name(8, self.columns)
+        return self.native.select(expr.alias(name)).schema()[name].dtype
+
     @property
     def columns(self) -> list[str]:
         if self._cached_columns is None:
@@ -187,6 +193,8 @@ def select(self, *exprs: DaftExpr) -> DaftLazyFrame:
                     f"{e!s}\n\nHint: Did you mean one of these columns: {self.columns}?"
                 )
                 raise ColumnNotFoundError(msg) from e
+            if "TypeError" in str(e):
+                raise InvalidOperationError(str(e)) from None
             raise
 
     def with_columns(self, *exprs: DaftExpr) -> DaftLazyFrame:

narwhals_daft/expr.py

@@ -345,6 +345,30 @@ def _with_elementwise(
             version=self._version,
         )
 
+    def _with_elementwise_frame(
+        self, call: Callable[[DaftLazyFrame, Expression], Expression], /
+    ) -> DaftExpr:
+        """Like `_with_elementwise`, but `call` also receives the frame.
+
+        Useful when building the native expression needs the frame's schema.
+        """
+
+        def func(df: DaftLazyFrame) -> list[Expression]:
+            return [call(df, expr) for expr in self(df)]
+
+        def window_function(
+            df: DaftLazyFrame, window_inputs: WindowInputs
+        ) -> list[Expression]:
+            return [call(df, expr) for expr in self.window_function(df, window_inputs)]
+
+        return self.__class__(
+            func,
+            window_function,
+            evaluate_output_names=self._evaluate_output_names,
+            alias_output_names=self._alias_output_names,
+            version=self._version,
+        )
+
     def _with_binary(self, op: Callable[..., Expression], other: DaftExpr) -> DaftExpr:
         return self.__class__(
             self._callable_to_eval_series(op, other=other),

narwhals_daft/expr_list.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import daft.functions as F
-from daft import lit
-from narwhals._utils import not_implemented
+from daft import DataType, lit
 from narwhals.compliant import ListNamespace
 
 if TYPE_CHECKING:
     from daft import Expression
+    from narwhals.typing import NonNestedLiteral
 
+    from narwhals_daft.dataframe import DaftLazyFrame
     from narwhals_daft.expr import DaftExpr
 
 
@@ -76,4 +77,50 @@ def func(expr: Expression) -> Expression:
 
         return self.compliant._with_elementwise(func)
 
-    contains = not_implemented()
+    def contains(self, item: NonNestedLiteral) -> DaftExpr:
+        if item is None:
+            return self.compliant._with_elementwise(
+                lambda expr: F.list_count(expr, "null") > lit(0)
+            )
+        item_dtype = DataType.infer_from_object(item)
+
+        def func(df: DaftLazyFrame, expr: Expression) -> Expression:
+            # `list_contains` requires `item` to have exactly the element type of
+            # the list, while a Python number only ever infers to Int64 or Float64,
+            # so cast it, when that is lossless, to the element type resolved
+            # against the frame's schema (plan-time only). Any other mismatch
+            # (e.g. a `Timestamp[us]` item on a `Timestamp[ns]` list) raises in
+            # Daft, like it does in Polars.
+            list_dtype = df._native_dtype(expr)
+            if not (list_dtype.is_list() or list_dtype.is_fixed_size_list()):
+                return F.list_contains(expr, lit(item))
+            element_dtype = list_dtype.dtype
+            if not (element_dtype.is_numeric() and item_dtype.is_numeric()):
+                return F.list_contains(expr, lit(item))
+            if not _is_representable(item, element_dtype):
+                # e.g. `1.5` or `300` can never be an element of an Int8 list.
+                return F.when(expr.not_null(), lit(False))
+            return F.list_contains(expr, lit(item).cast(element_dtype))
+
+        return self.compliant._with_elementwise_frame(func)
+
+
+_INTEGER_BOUNDS: dict[str, tuple[int, int]] = {
+    "Int8": (-(1 << 7), (1 << 7) - 1),
+    "Int16": (-(1 << 15), (1 << 15) - 1),
+    "Int32": (-(1 << 31), (1 << 31) - 1),
+    "Int64": (-(1 << 63), (1 << 63) - 1),
+    "UInt8": (0, (1 << 8) - 1),
+    "UInt16": (0, (1 << 16) - 1),
+    "UInt32": (0, (1 << 32) - 1),
+    "UInt64": (0, (1 << 64) - 1),
+}
+
+
+def _is_representable(item: Any, element_dtype: DataType) -> bool:
+    """Whether the numeric `item` survives a cast to an integer `element_dtype`."""
+    if not element_dtype.is_integer():
+        return True
+    lower, upper = _INTEGER_BOUNDS[str(element_dtype)]
+    is_integral = isinstance(item, int) or item.is_integer()
+    return is_integral and lower <= item <= upper

narwhals_daft/utils.py

@@ -123,9 +123,9 @@ def narwhals_to_native_dtype(  # noqa: PLR0912,C901,PLR0911
         return DataType.time("ns")
     if dtype == dtypes.Binary:
         return DataType.binary()
-    if dtype == dtypes.Decimal:  # pragma: no cover
-        msg = "Casting to Decimal is not supported yet."
-        raise NotImplementedError(msg)
+    if isinstance_or_issubclass(dtype, dtypes.Decimal):
+        decimal = dtypes.Decimal() if isinstance(dtype, type) else dtype
+        return DataType.decimal128(decimal.precision, decimal.scale)
     if isinstance_or_issubclass(dtype, dtypes.Datetime):
         return DataType.timestamp(dtype.time_unit, dtype.time_zone)
     if isinstance_or_issubclass(dtype, dtypes.Duration):

run_tests.py

@@ -30,7 +30,6 @@
     "test_concat_str_with_lit",
     "test_concat_str_with_lit_and_nulls",
     "test_concat_vertical",
-    "test_contains_expr",
     "test_convert_time_zone",
     "test_convert_time_zone_from_none",
     "test_corr_constant_column",