Merge branch 'main' into oh-nodes

Author: dangotbanned

.pre-commit-config.yaml

@@ -22,10 +22,12 @@ repos:
       args: [--ignore-words-list=ser, --ignore-words-list=RightT]
       exclude: ^docs/api-completeness.md$
 -   repo: https://github.com/pycqa/flake8
-    rev: '7.3.0'  # todo: remove once https://github.com/astral-sh/ruff/issues/458 is addressed
+    rev: '7.3.0'
     hooks:
     -   id: flake8
-        additional_dependencies: [darglint==1.8.1, Flake8-pyproject]
+        # TODO(unassigned): replace with ruff once https://github.com/astral-sh/ruff/issues/458 is addressed.
+        name: darglint
+        alias: darglint
         entry: flake8 --select DAR --ignore DAR101,DAR201,DAR401,DAR402
         exclude: |
             (?x)^(
@@ -39,6 +41,15 @@ repos:
                 # Soft-deprecated, so less crucial to document so carefully
                 narwhals/stable/v1/.*
             )$
+        # Keep in sync with `flake8-typing-imports` so the same venv is reused.
+        additional_dependencies: [darglint==1.8.1, flake8-typing-imports==1.17.0]
+    -   id: flake8
+        # TODO(unassigned): replace with ruff once https://github.com/astral-sh/ruff/issues/2302 is addressed.
+        name: flake8-typing-imports
+        alias: flake8-typing-imports
+        entry: flake8 --select TYP --min-python-version=3.9.0
+        # Keep in sync with `darglint` so the same venv is reused.
+        additional_dependencies: [darglint==1.8.1, flake8-typing-imports==1.17.0]
 - repo: local
   hooks:
     - id: check-api-reference
@@ -52,6 +63,11 @@ repos:
       pass_filenames: false
       entry: python -m utils.sort_api_reference
       language: python
+    - id: check-slotted-classes
+      name: check-slotted-classes
+      pass_filenames: false
+      entry: python -m utils.check_slotted_classes
+      language: python
     - id: imports-are-banned
       name: import are banned (use `get_pandas` instead of `import pandas`)
       entry: python utils/import_check.py

narwhals/_arrow/namespace.py

@@ -23,7 +23,12 @@
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
-    from narwhals._arrow.typing import ArrayOrScalar, ChunkedArrayAny, Incomplete
+    from narwhals._arrow.typing import (
+        ArrayOrScalar,
+        ChunkedArrayAny,
+        Incomplete,
+        ScalarAny,
+    )
     from narwhals._compliant.typing import ScalarKwargs
     from narwhals._utils import Version
     from narwhals.typing import IntoDType, NonNestedLiteral
@@ -49,6 +54,11 @@ def _series(self) -> type[ArrowSeries]:
     def __init__(self, *, version: Version) -> None:
         self._version = version
 
+    def extract_native(
+        self, *series: ArrowSeries
+    ) -> Iterator[ChunkedArrayAny | ScalarAny]:
+        return (s.native[0] if s._broadcast else s.native for s in series)
+
     def len(self) -> ArrowExpr:
         # coverage bug? this is definitely hit
         return self._expr(  # pragma: no cover
@@ -83,10 +93,9 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries:
     def all_horizontal(self, *exprs: ArrowExpr, ignore_nulls: bool) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series: Iterator[ArrowSeries] = chain.from_iterable(e(df) for e in exprs)
-            align = self._series._align_full_broadcast
             if ignore_nulls:
                 series = (s.fill_null(True, None, None) for s in series)
-            return [reduce(operator.and_, align(*series))]
+            return [reduce(operator.and_, series)]
 
         return self._expr._from_callable(
             func=func,
@@ -100,10 +109,9 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
     def any_horizontal(self, *exprs: ArrowExpr, ignore_nulls: bool) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series: Iterator[ArrowSeries] = chain.from_iterable(e(df) for e in exprs)
-            align = self._series._align_full_broadcast
             if ignore_nulls:
                 series = (s.fill_null(False, None, None) for s in series)
-            return [reduce(operator.or_, align(*series))]
+            return [reduce(operator.or_, series)]
 
         return self._expr._from_callable(
             func=func,
@@ -118,8 +126,7 @@ def sum_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             it = chain.from_iterable(expr(df) for expr in exprs)
             series = (s.fill_null(0, strategy=None, limit=None) for s in it)
-            align = self._series._align_full_broadcast
-            return [reduce(operator.add, align(*series))]
+            return [reduce(operator.add, series)]
 
         return self._expr._from_callable(
             func=func,
@@ -135,11 +142,8 @@ def mean_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             expr_results = tuple(chain.from_iterable(expr(df) for expr in exprs))
-            align = self._series._align_full_broadcast
-            series = align(
-                *(s.fill_null(0, strategy=None, limit=None) for s in expr_results)
-            )
-            non_na = align(*(1 - s.is_null().cast(int_64) for s in expr_results))
+            series = [s.fill_null(0, strategy=None, limit=None) for s in expr_results]
+            non_na = [1 - s.is_null().cast(int_64) for s in expr_results]
             return [reduce(operator.add, series) / reduce(operator.add, non_na)]
 
         return self._expr._from_callable(
@@ -153,9 +157,7 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
 
     def min_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            align = self._series._align_full_broadcast
             init_series, *series = tuple(chain.from_iterable(expr(df) for expr in exprs))
-            init_series, *series = align(init_series, *series)
             native_series = reduce(
                 pc.min_element_wise, [s.native for s in series], init_series.native
             )
@@ -174,9 +176,7 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
 
     def max_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            align = self._series._align_full_broadcast
             init_series, *series = tuple(chain.from_iterable(expr(df) for expr in exprs))
-            init_series, *series = align(init_series, *series)
             native_series = reduce(
                 pc.max_element_wise, [s.native for s in series], init_series.native
             )
@@ -227,16 +227,13 @@ def concat_str(
         self, *exprs: ArrowExpr, separator: str, ignore_nulls: bool
     ) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            align = self._series._align_full_broadcast
-            compliant_series_list = align(
-                *(chain.from_iterable(expr(df) for expr in exprs))
-            )
-            name = compliant_series_list[0].name
+            series = list(chain.from_iterable(expr(df) for expr in exprs))
+            name = series[0].name
             null_handling: Literal["skip", "emit_null"] = (
                 "skip" if ignore_nulls else "emit_null"
             )
             it, separator_scalar = cast_to_comparable_string_types(
-                *(s.native for s in compliant_series_list), separator=separator
+                *self.extract_native(*series), separator=separator
             )
             # NOTE: stubs indicate `separator` must also be a `ChunkedArray`
             # Reality: `str` is fine

narwhals/_arrow/series.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Literal, cast, overload
 
 import pyarrow as pa
 import pyarrow.compute as pc
@@ -148,6 +148,16 @@ def _with_native(
             result._broadcast = self._broadcast
         return result
 
+    def _with_binary(self, op: Callable[..., ArrayOrScalar], other: Any) -> Self:
+        ser, other_native = extract_native(self, other)
+        preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True)
+        return self._with_native(
+            op(ser, other_native), preserve_broadcast=preserve_broadcast
+        ).alias(self.name)
+
+    def _with_binary_right(self, op: Callable[..., ArrayOrScalar], other: Any) -> Self:
+        return self._with_binary(lambda x, y: op(y, x), other).alias(self.name)
+
     @classmethod
     def from_iterable(
         cls,
@@ -214,106 +224,89 @@ def __narwhals_namespace__(self) -> ArrowNamespace:
         return ArrowNamespace(version=self._version)
 
     def __eq__(self, other: object) -> Self:  # type: ignore[override]
-        other = cast("PythonLiteral | ArrowSeries | None", other)
-        ser, rhs = extract_native(self, other)
-        return self._with_native(pc.equal(ser, rhs))
+        return self._with_binary(pc.equal, other)
 
     def __ne__(self, other: object) -> Self:  # type: ignore[override]
-        other = cast("PythonLiteral | ArrowSeries | None", other)
-        ser, rhs = extract_native(self, other)
-        return self._with_native(pc.not_equal(ser, rhs))
+        return self._with_binary(pc.not_equal, other)
 
     def __ge__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.greater_equal(ser, other))
+        return self._with_binary(pc.greater_equal, other)
 
     def __gt__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.greater(ser, other))
+        return self._with_binary(pc.greater, other)
 
     def __le__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.less_equal(ser, other))
+        return self._with_binary(pc.less_equal, other)
 
     def __lt__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.less(ser, other))
+        return self._with_binary(pc.less, other)
 
     def __and__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.and_kleene(ser, other))  # type: ignore[arg-type]
+        return self._with_binary(pc.and_kleene, other)
 
     def __rand__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.and_kleene(other, ser))  # type: ignore[arg-type]
+        return self._with_binary_right(pc.and_kleene, other)
 
     def __or__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.or_kleene(ser, other))  # type: ignore[arg-type]
+        return self._with_binary_right(pc.or_kleene, other)
 
     def __ror__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.or_kleene(other, ser))  # type: ignore[arg-type]
+        return self._with_binary_right(pc.or_kleene, other)
 
     def __add__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.add(ser, other))
+        return self._with_binary(pc.add, other)
 
     def __radd__(self, other: Any) -> Self:
-        return self + other
+        return self._with_binary_right(pc.add, other)
 
     def __sub__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.subtract(ser, other))
+        return self._with_binary(pc.subtract, other)
 
     def __rsub__(self, other: Any) -> Self:
-        return (self - other) * (-1)
+        return self._with_binary_right(pc.subtract, other)
 
     def __mul__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.multiply(ser, other))
+        return self._with_binary(pc.multiply, other)
 
     def __rmul__(self, other: Any) -> Self:
-        return self * other
+        return self._with_binary_right(pc.multiply, other)
 
     def __pow__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.power(ser, other))
+        return self._with_binary(pc.power, other)
 
     def __rpow__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.power(other, ser))
+        return self._with_binary_right(pc.power, other)
 
     def __floordiv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(floordiv_compat(ser, other))
+        return self._with_binary(floordiv_compat, other)
 
     def __rfloordiv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(floordiv_compat(other, ser))
+        return self._with_binary_right(floordiv_compat, other)
 
     def __truediv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.divide(*cast_for_truediv(ser, other)))  # type: ignore[type-var]
+        return self._with_binary(lambda x, y: pc.divide(*cast_for_truediv(x, y)), other)
 
     def __rtruediv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.divide(*cast_for_truediv(other, ser)))  # type: ignore[type-var]
+        return self._with_binary_right(
+            lambda x, y: pc.divide(*cast_for_truediv(x, y)), other
+        )
 
     def __mod__(self, other: Any) -> Self:
+        preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True)
         floor_div = (self // other).native
         ser, other = extract_native(self, other)
         res = pc.subtract(ser, pc.multiply(floor_div, other))
-        return self._with_native(res)
+        return self._with_native(res, preserve_broadcast=preserve_broadcast)
 
     def __rmod__(self, other: Any) -> Self:
+        preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True)
         floor_div = (other // self).native
         ser, other = extract_native(self, other)
         res = pc.subtract(other, pc.multiply(floor_div, ser))
-        return self._with_native(res)
+        return self._with_native(res, preserve_broadcast=preserve_broadcast)
 
     def __invert__(self) -> Self:
-        return self._with_native(pc.invert(self.native))
+        return self._with_native(pc.invert(self.native), preserve_broadcast=True)
 
     @property
     def _type(self) -> pa.DataType:

narwhals/_arrow/utils.py

@@ -445,8 +445,8 @@ def pad_series(
 
 
 def cast_to_comparable_string_types(
-    *chunked_arrays: ChunkedArrayAny, separator: str
-) -> tuple[Iterator[ChunkedArrayAny], ScalarAny]:
+    *chunked_arrays: ChunkedArrayAny | ScalarAny, separator: str
+) -> tuple[Iterator[ChunkedArrayAny | ScalarAny], ScalarAny]:
     # Ensure `chunked_arrays` are either all `string` or all `large_string`.
     dtype = (
         pa.string()  # (PyArrow default)