Merge branch 'main' into feat/with-row-index-by

Author: FBruzzesi

.github/workflows/downstream_tests.yml

@@ -309,22 +309,26 @@ jobs:
           git clone https://github.com/plotly/plotly.py --depth=1
           cd plotly.py
           git log
-      - name: install-basics
-        run: uv pip install --upgrade tox virtualenv setuptools --system
       - name: install-deps
         run: |
           cd plotly.py
-          uv pip install -r test_requirements/requirements_core.txt -r test_requirements/requirements_optional.txt --system
-          uv pip install -e . --system
+          uv venv -p ${{ matrix.python-version }}
+          uv sync --extra dev_optional
       - name: install-narwhals-dev
         run: |
-          uv pip uninstall narwhals --system
-          uv pip install -e . --system
+          cd plotly.py
+          . .venv/bin/activate
+          uv pip uninstall narwhals
+          uv pip install -e ./..
       - name: show-deps
-        run: uv pip freeze
+        run: |
+          cd plotly.py
+          . .venv/bin/activate
+          uv pip freeze
       - name: Run pytest on plotly express
         run: |
           cd plotly.py
+          . .venv/bin/activate
           pytest tests/test_optional/test_px
 
   hierarchicalforecast:

docs/api-reference/expr_str.md

@@ -18,5 +18,6 @@
         - to_datetime
         - to_lowercase
         - to_uppercase
+        - zfill
       show_source: false
       show_bases: false

docs/api-reference/series_str.md

@@ -18,5 +18,6 @@
         - to_datetime
         - to_lowercase
         - to_uppercase
+        - zfill
       show_source: false
       show_bases: false

narwhals/_arrow/dataframe.py

@@ -8,7 +8,7 @@
 import pyarrow.compute as pc
 
 from narwhals._arrow.series import ArrowSeries
-from narwhals._arrow.utils import align_series_full_broadcast, native_to_narwhals_dtype
+from narwhals._arrow.utils import native_to_narwhals_dtype
 from narwhals._compliant import EagerDataFrame
 from narwhals._expression_parsing import ExprKind
 from narwhals._utils import (
@@ -74,7 +74,9 @@
     PromoteOptions: TypeAlias = Literal["none", "default", "permissive"]
 
 
-class ArrowDataFrame(EagerDataFrame["ArrowSeries", "ArrowExpr", "pa.Table"]):
+class ArrowDataFrame(
+    EagerDataFrame["ArrowSeries", "ArrowExpr", "pa.Table", "ChunkedArrayAny"]
+):
     def __init__(
         self,
         native_dataframe: pa.Table,
@@ -330,7 +332,8 @@ def select(self: ArrowDataFrame, *exprs: ArrowExpr) -> ArrowDataFrame:
                 self.native.__class__.from_arrays([]), validate_column_names=False
             )
         names = [s.name for s in new_series]
-        reshaped = align_series_full_broadcast(*new_series)
+        align = new_series[0]._align_full_broadcast
+        reshaped = align(*new_series)
         df = pa.Table.from_arrays([s.native for s in reshaped], names=names)
         return self._with_native(df, validate_column_names=True)
 

narwhals/_arrow/namespace.py

@@ -12,10 +12,7 @@
 from narwhals._arrow.expr import ArrowExpr
 from narwhals._arrow.selectors import ArrowSelectorNamespace
 from narwhals._arrow.series import ArrowSeries
-from narwhals._arrow.utils import (
-    align_series_full_broadcast,
-    cast_to_comparable_string_types,
-)
+from narwhals._arrow.utils import cast_to_comparable_string_types
 from narwhals._compliant import CompliantThen, EagerNamespace, EagerWhen
 from narwhals._expression_parsing import (
     combine_alias_output_names,
@@ -26,12 +23,14 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from narwhals._arrow.typing import Incomplete
+    from narwhals._arrow.typing import ArrayOrScalar, ChunkedArrayAny, Incomplete
     from narwhals._utils import Version
     from narwhals.typing import IntoDType, NonNestedLiteral
 
 
-class ArrowNamespace(EagerNamespace[ArrowDataFrame, ArrowSeries, ArrowExpr, pa.Table]):
+class ArrowNamespace(
+    EagerNamespace[ArrowDataFrame, ArrowSeries, ArrowExpr, pa.Table, "ChunkedArrayAny"]
+):
     @property
     def _dataframe(self) -> type[ArrowDataFrame]:
         return ArrowDataFrame
@@ -86,7 +85,8 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries:
     def all_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series = chain.from_iterable(expr(df) for expr in exprs)
-            return [reduce(operator.and_, align_series_full_broadcast(*series))]
+            align = self._series._align_full_broadcast
+            return [reduce(operator.and_, align(*series))]
 
         return self._expr._from_callable(
             func=func,
@@ -100,7 +100,8 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
     def any_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series = chain.from_iterable(expr(df) for expr in exprs)
-            return [reduce(operator.or_, align_series_full_broadcast(*series))]
+            align = self._series._align_full_broadcast
+            return [reduce(operator.or_, align(*series))]
 
         return self._expr._from_callable(
             func=func,
@@ -115,7 +116,8 @@ 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)
-            return [reduce(operator.add, align_series_full_broadcast(*series))]
+            align = self._series._align_full_broadcast
+            return [reduce(operator.add, align(*series))]
 
         return self._expr._from_callable(
             func=func,
@@ -131,12 +133,11 @@ def mean_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             expr_results = list(chain.from_iterable(expr(df) for expr in exprs))
-            series = align_series_full_broadcast(
+            align = self._series._align_full_broadcast
+            series = align(
                 *(s.fill_null(0, strategy=None, limit=None) for s in expr_results)
             )
-            non_na = align_series_full_broadcast(
-                *(1 - s.is_null().cast(int_64) for s in expr_results)
-            )
+            non_na = align(*(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(
@@ -150,8 +151,9 @@ 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 = list(chain.from_iterable(expr(df) for expr in exprs))
-            init_series, *series = align_series_full_broadcast(init_series, *series)
+            init_series, *series = align(init_series, *series)
             native_series = reduce(
                 pc.min_element_wise, [s.native for s in series], init_series.native
             )
@@ -175,8 +177,9 @@ 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 = list(chain.from_iterable(expr(df) for expr in exprs))
-            init_series, *series = align_series_full_broadcast(init_series, *series)
+            init_series, *series = align(init_series, *series)
             native_series = reduce(
                 pc.max_element_wise, [s.native for s in series], init_series.native
             )
@@ -232,7 +235,8 @@ def concat_str(
         self, *exprs: ArrowExpr, separator: str, ignore_nulls: bool
     ) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            compliant_series_list = align_series_full_broadcast(
+            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
@@ -263,23 +267,20 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
         )
 
 
-class ArrowWhen(EagerWhen[ArrowDataFrame, ArrowSeries, ArrowExpr]):
+class ArrowWhen(EagerWhen[ArrowDataFrame, ArrowSeries, ArrowExpr, "ChunkedArrayAny"]):
     @property
     def _then(self) -> type[ArrowThen]:
         return ArrowThen
 
     def _if_then_else(
-        self, when: ArrowSeries, then: ArrowSeries, otherwise: ArrowSeries | None, /
-    ) -> ArrowSeries:
-        if otherwise is None:
-            when, then = align_series_full_broadcast(when, then)
-            res_native = pc.if_else(
-                when.native, then.native, pa.nulls(len(when.native), then.native.type)
-            )
-        else:
-            when, then, otherwise = align_series_full_broadcast(when, then, otherwise)
-            res_native = pc.if_else(when.native, then.native, otherwise.native)
-        return then._with_native(res_native)
+        self,
+        when: ChunkedArrayAny,
+        then: ChunkedArrayAny,
+        otherwise: ArrayOrScalar | NonNestedLiteral,
+        /,
+    ) -> ChunkedArrayAny:
+        otherwise = pa.nulls(len(when), then.type) if otherwise is None else otherwise
+        return pc.if_else(when, then, otherwise)
 
 
 class ArrowThen(CompliantThen[ArrowDataFrame, ArrowSeries, ArrowExpr], ArrowExpr): ...

narwhals/_arrow/series.py

@@ -32,7 +32,7 @@
     validate_backend_version,
 )
 from narwhals.dependencies import is_numpy_array_1d
-from narwhals.exceptions import InvalidOperationError
+from narwhals.exceptions import InvalidOperationError, ShapeError
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Mapping, Sequence
@@ -192,6 +192,25 @@ def from_numpy(cls, data: Into1DArray, /, *, context: _FullContext) -> Self:
             data if is_numpy_array_1d(data) else [data], context=context
         )
 
+    @classmethod
+    def _align_full_broadcast(cls, *series: Self) -> Sequence[Self]:
+        lengths = [len(s) for s in series]
+        max_length = max(lengths)
+        fast_path = all(_len == max_length for _len in lengths)
+        if fast_path:
+            return series
+        reshaped = []
+        for s in series:
+            if s._broadcast:
+                compliant = s._with_native(pa.repeat(s.native[0], max_length))
+            elif (actual_len := len(s)) != max_length:
+                msg = f"Expected object of length {max_length}, got {actual_len}."
+                raise ShapeError(msg)
+            else:
+                compliant = s
+            reshaped.append(compliant)
+        return reshaped
+
     def __narwhals_namespace__(self) -> ArrowNamespace:
         from narwhals._arrow.namespace import ArrowNamespace
 

narwhals/_arrow/series_str.py

@@ -3,12 +3,14 @@
 import string
 from typing import TYPE_CHECKING
 
+import pyarrow as pa
 import pyarrow.compute as pc
 
 from narwhals._arrow.utils import ArrowSeriesNamespace, lit, parse_datetime_format
 
 if TYPE_CHECKING:
     from narwhals._arrow.series import ArrowSeries
+    from narwhals._arrow.typing import Incomplete
 
 
 class ArrowSeriesStringNamespace(ArrowSeriesNamespace):
@@ -60,3 +62,36 @@ def to_uppercase(self) -> ArrowSeries:
 
     def to_lowercase(self) -> ArrowSeries:
         return self.with_native(pc.utf8_lower(self.native))
+
+    def zfill(self, width: int) -> ArrowSeries:
+        binary_join: Incomplete = pc.binary_join_element_wise
+        native = self.native
+        hyphen, plus = lit("-"), lit("+")
+        first_char, remaining_chars = self.slice(0, 1).native, self.slice(1, None).native
+
+        # Conditions
+        less_than_width = pc.less(pc.utf8_length(native), lit(width))
+        starts_with_hyphen = pc.equal(first_char, hyphen)
+        starts_with_plus = pc.equal(first_char, plus)
+
+        conditions = pc.make_struct(
+            pc.and_(starts_with_hyphen, less_than_width),
+            pc.and_(starts_with_plus, less_than_width),
+            less_than_width,
+        )
+
+        # Cases
+        padded_remaining_chars = pc.utf8_lpad(remaining_chars, width - 1, padding="0")
+
+        result = pc.case_when(
+            conditions,
+            binary_join(
+                pa.repeat(hyphen, len(native)), padded_remaining_chars, ""
+            ),  # starts with hyphen and less than width
+            binary_join(
+                pa.repeat(plus, len(native)), padded_remaining_chars, ""
+            ),  # starts with plus and less than width
+            pc.utf8_lpad(native, width=width, padding="0"),  # less than width
+            native,
+        )
+        return self.with_native(result)

narwhals/_arrow/utils.py

@@ -8,10 +8,9 @@
 
 from narwhals._compliant.series import _SeriesNamespace
 from narwhals._utils import isinstance_or_issubclass
-from narwhals.exceptions import ShapeError
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator, Mapping, Sequence
+    from collections.abc import Iterable, Iterator, Mapping
 
     from typing_extensions import TypeAlias, TypeIs
 
@@ -261,31 +260,6 @@ def extract_native(
     return lhs.native, rhs if isinstance(rhs, pa.Scalar) else lit(rhs)
 
 
-def align_series_full_broadcast(*series: ArrowSeries) -> Sequence[ArrowSeries]:
-    # Ensure all of `series` are of the same length.
-    lengths = [len(s) for s in series]
-    max_length = max(lengths)
-    fast_path = all(_len == max_length for _len in lengths)
-
-    if fast_path:
-        return series
-
-    reshaped = []
-    for s in series:
-        if s._broadcast:
-            value = s.native[0]
-            if s._backend_version < (13,) and hasattr(value, "as_py"):
-                value = value.as_py()
-            reshaped.append(s._with_native(pa.array([value] * max_length, type=s._type)))
-        else:
-            if (actual_len := len(s)) != max_length:
-                msg = f"Expected object of length {max_length}, got {actual_len}."
-                raise ShapeError(msg)
-            reshaped.append(s)
-
-    return reshaped
-
-
 def floordiv_compat(left: ArrayOrScalar, right: ArrayOrScalar, /) -> Any:
     # The following lines are adapted from pandas' pyarrow implementation.
     # Ref: https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L124-L154

narwhals/_compliant/any_namespace.py

@@ -79,6 +79,7 @@ def split(self, by: str) -> CompliantT_co: ...
     def to_datetime(self, format: str | None) -> CompliantT_co: ...
     def to_lowercase(self) -> CompliantT_co: ...
     def to_uppercase(self) -> CompliantT_co: ...
+    def zfill(self, width: int) -> CompliantT_co: ...
 
 
 class StructNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):

narwhals/_compliant/dataframe.py

@@ -13,6 +13,7 @@
     EagerSeriesT,
     NativeExprT,
     NativeFrameT,
+    NativeSeriesT,
 )
 from narwhals._translate import (
     ArrowConvertible,
@@ -392,11 +393,11 @@ def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | N
 class EagerDataFrame(
     CompliantDataFrame[EagerSeriesT, EagerExprT, NativeFrameT, "DataFrame[NativeFrameT]"],
     CompliantLazyFrame[EagerExprT, NativeFrameT, "DataFrame[NativeFrameT]"],
-    Protocol[EagerSeriesT, EagerExprT, NativeFrameT],
+    Protocol[EagerSeriesT, EagerExprT, NativeFrameT, NativeSeriesT],
 ):
     def __narwhals_namespace__(
         self,
-    ) -> EagerNamespace[Self, EagerSeriesT, EagerExprT, NativeFrameT]: ...
+    ) -> EagerNamespace[Self, EagerSeriesT, EagerExprT, NativeFrameT, NativeSeriesT]: ...
 
     def to_narwhals(self) -> DataFrame[NativeFrameT]:
         return self._version.dataframe(self, level="full")
@@ -444,10 +445,14 @@ def _numpy_column_names(
     ) -> list[str]:
         return list(columns or (f"column_{x}" for x in range(data.shape[1])))
 
-    def _gather(self, rows: SizedMultiIndexSelector[Any]) -> Self: ...
+    def _gather(self, rows: SizedMultiIndexSelector[NativeSeriesT]) -> Self: ...
     def _gather_slice(self, rows: _SliceIndex | range) -> Self: ...
-    def _select_multi_index(self, columns: SizedMultiIndexSelector[Any]) -> Self: ...
-    def _select_multi_name(self, columns: SizedMultiNameSelector[Any]) -> Self: ...
+    def _select_multi_index(
+        self, columns: SizedMultiIndexSelector[NativeSeriesT]
+    ) -> Self: ...
+    def _select_multi_name(
+        self, columns: SizedMultiNameSelector[NativeSeriesT]
+    ) -> Self: ...
     def _select_slice_index(self, columns: _SliceIndex | range) -> Self: ...
     def _select_slice_name(self, columns: _SliceName) -> Self: ...
     def __getitem__(  # noqa: C901, PLR0912