fix(typing): Resolve all mypy & pyright errors for _arrow

narwhals/_arrow/dataframe.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from functools import partial
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterator
@@ -36,9 +37,6 @@
 
     import pandas as pd
     import polars as pl
-    from pyarrow._stubs_typing import (  # pyright: ignore[reportMissingModuleSource]
-        Indices,
-    )
     from pyarrow._stubs_typing import Order  # pyright: ignore[reportMissingModuleSource]
     from typing_extensions import Self
     from typing_extensions import TypeAlias
@@ -47,6 +45,8 @@
     from narwhals._arrow.group_by import ArrowGroupBy
     from narwhals._arrow.namespace import ArrowNamespace
     from narwhals._arrow.series import ArrowSeries
+    from narwhals._arrow.typing import Indices
+    from narwhals._arrow.typing import Mask
     from narwhals.dtypes import DType
     from narwhals.typing import SizeUnit
     from narwhals.typing import _1DArray
@@ -133,7 +133,7 @@ def __len__(self: Self) -> int:
         return len(self._native_frame)
 
     def row(self: Self, index: int) -> tuple[Any, ...]:
-        return tuple(col[index] for col in self._native_frame)
+        return tuple(col[index] for col in self._native_frame.itercolumns())
 
     @overload
     def rows(self: Self, *, named: Literal[True]) -> list[dict[str, Any]]: ...
@@ -165,7 +165,7 @@ def iter_rows(
             for i in range(0, num_rows, buffer_size):
                 yield from df[i : i + buffer_size].to_pylist()
 
-    def get_column(self: Self, name: str) -> ArrowSeries:
+    def get_column(self: Self, name: str) -> ArrowSeries[Any]:
         from narwhals._arrow.series import ArrowSeries
 
         if not isinstance(name, str):
@@ -185,7 +185,7 @@ def __array__(self: Self, dtype: Any, copy: bool | None) -> _2DArray:
     @overload
     def __getitem__(  # type: ignore[overload-overlap, unused-ignore]
         self: Self, item: str | tuple[slice | Sequence[int] | _1DArray, int | str]
-    ) -> ArrowSeries: ...
+    ) -> ArrowSeries[Any]: ...
     @overload
     def __getitem__(
         self: Self,
@@ -214,7 +214,7 @@ def __getitem__(
                 slice | Sequence[int] | _1DArray, slice | Sequence[int] | Sequence[str]
             ]
         ),
-    ) -> ArrowSeries | Self:
+    ) -> ArrowSeries[Any] | Self:
         if isinstance(item, tuple):
             item = tuple(list(i) if is_sequence_but_not_str(i) else i for i in item)  # pyright: ignore[reportAssignmentType]
 
@@ -345,7 +345,7 @@ def aggregate(self: Self, *exprs: ArrowExpr) -> Self:
         return self.select(*exprs)
 
     def select(self: Self, *exprs: ArrowExpr) -> Self:
-        new_series: list[ArrowSeries] = evaluate_into_exprs(self, *exprs)
+        new_series: list[ArrowSeries[Any]] = evaluate_into_exprs(self, *exprs)
         if not new_series:
             # return empty dataframe, like Polars does
             return self._from_native_frame(
@@ -357,7 +357,7 @@ def select(self: Self, *exprs: ArrowExpr) -> Self:
 
     def with_columns(self: Self, *exprs: ArrowExpr) -> Self:
         native_frame = self._native_frame
-        new_columns: list[ArrowSeries] = evaluate_into_exprs(self, *exprs)
+        new_columns: list[ArrowSeries[Any]] = evaluate_into_exprs(self, *exprs)
 
         length = len(self)
         columns = self.columns
@@ -497,14 +497,16 @@ def to_numpy(self: Self) -> _2DArray:
         return arr
 
     @overload
-    def to_dict(self: Self, *, as_series: Literal[True]) -> dict[str, ArrowSeries]: ...
+    def to_dict(
+        self: Self, *, as_series: Literal[True]
+    ) -> dict[str, ArrowSeries[Any]]: ...
 
     @overload
     def to_dict(self: Self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ...
 
     def to_dict(
         self: Self, *, as_series: bool
-    ) -> dict[str, ArrowSeries] | dict[str, list[Any]]:
+    ) -> dict[str, ArrowSeries[Any]] | dict[str, list[Any]]:
         df = self._native_frame
 
         names_and_values = zip(df.column_names, df.columns)
@@ -532,9 +534,9 @@ def with_row_index(self: Self, name: str) -> Self:
             df.append_column(name, row_indices).select([name, *cols])
         )
 
-    def filter(self: Self, predicate: ArrowExpr | list[bool]) -> Self:
+    def filter(self: Self, predicate: ArrowExpr | list[bool | None]) -> Self:
         if isinstance(predicate, list):
-            mask_native = predicate
+            mask_native: Mask = predicate
         else:
             # `[0]` is safe as the predicate's expression only returns a single column
             mask = evaluate_into_exprs(self, predicate)[0]
@@ -703,7 +705,7 @@ def write_csv(self: Self, file: str | Path | BytesIO | None) -> str | None:
         pa_csv.write_csv(pa_table, file)
         return None
 
-    def is_unique(self: Self) -> ArrowSeries:
+    def is_unique(self: Self) -> ArrowSeries[Any]:
         from narwhals._arrow.series import ArrowSeries
 
         col_token = generate_temporary_column_name(n_bytes=8, columns=self.columns)
@@ -745,17 +747,14 @@ def unique(
 
             agg_func = agg_func_map[keep]
             col_token = generate_temporary_column_name(n_bytes=8, columns=self.columns)
-            keep_idx = (
+            keep_idx_native = (
                 df.append_column(col_token, pa.array(np.arange(len(self))))
                 .group_by(subset)
                 .aggregate([(col_token, agg_func)])
                 .column(f"{col_token}_{agg_func}")
             )
-
-            return self._from_native_frame(
-                pc.take(df, keep_idx),  # type: ignore[call-overload, unused-ignore]
-                validate_column_names=False,
-            )
+            indices = cast("Indices", keep_idx_native)
+            return self._from_native_frame(df.take(indices), validate_column_names=False)
 
         keep_idx = self.simple_select(*subset).is_unique()
         plx = self.__narwhals_namespace__()
@@ -804,30 +803,28 @@ def unpivot(
         on_: list[str] = (
             [c for c in self.columns if c not in index_] if on is None else on
         )
-
-        promote_kwargs: dict[Literal["promote_options"], PromoteOptions] = (
-            {"promote_options": "permissive"}
+        concat = (
+            partial(pa.concat_tables, promote_options="permissive")
             if self._backend_version >= (14, 0, 0)
-            else {}
+            else pa.concat_tables
         )
         names = [*index_, variable_name, value_name]
         return self._from_native_frame(
-            pa.concat_tables(
+            concat(
                 [
                     pa.Table.from_arrays(
                         [
                             *(native_frame.column(idx_col) for idx_col in index_),
                             cast(
-                                "pa.ChunkedArray",
+                                "pa.ChunkedArray[Any]",
                                 pa.array([on_col] * n_rows, pa.string()),
                             ),
                             native_frame.column(on_col),
                         ],
                         names=names,
                     )
                     for on_col in on_
-                ],
-                **promote_kwargs,
+                ]
             )
         )
         # TODO(Unassigned): Even with promote_options="permissive", pyarrow does not

narwhals/_arrow/expr.py

@@ -29,12 +29,12 @@
     from narwhals.utils import Version
 
 
-class ArrowExpr(CompliantExpr[ArrowSeries]):
+class ArrowExpr(CompliantExpr[ArrowSeries[Any]]):
     _implementation: Implementation = Implementation.PYARROW
 
     def __init__(
         self: Self,
-        call: Callable[[ArrowDataFrame], Sequence[ArrowSeries]],
+        call: Callable[[ArrowDataFrame], Sequence[ArrowSeries[Any]]],
         *,
         depth: int,
         function_name: str,
@@ -48,7 +48,7 @@ def __init__(
         self._depth = depth
         self._function_name = function_name
         self._depth = depth
-        self._evaluate_output_names = evaluate_output_names
+        self._evaluate_output_names = evaluate_output_names  # pyright: ignore[reportAttributeAccessIssue]
         self._alias_output_names = alias_output_names
         self._backend_version = backend_version
         self._version = version
@@ -57,7 +57,7 @@ def __init__(
     def __repr__(self: Self) -> str:  # pragma: no cover
         return f"ArrowExpr(depth={self._depth}, function_name={self._function_name}, "
 
-    def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
+    def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries[Any]]:
         return self._call(df)
 
     @classmethod
@@ -69,7 +69,7 @@ def from_column_names(
     ) -> Self:
         from narwhals._arrow.series import ArrowSeries
 
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries[Any]]:
             try:
                 return [
                     ArrowSeries(
@@ -106,7 +106,7 @@ def from_column_indices(
     ) -> Self:
         from narwhals._arrow.series import ArrowSeries
 
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries[Any]]:
             return [
                 ArrowSeries(
                     df._native_frame[column_index],
@@ -370,7 +370,7 @@ def clip(self: Self, lower_bound: Any | None, upper_bound: Any | None) -> Self:
         )
 
     def over(self: Self, keys: list[str]) -> Self:
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries[Any]]:
             output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
             if overlap := set(output_names).intersection(keys):
                 # E.g. `df.select(nw.all().sum().over('a'))`. This is well-defined,
@@ -406,7 +406,7 @@ def map_batches(
         function: Callable[[Any], Any],
         return_dtype: DType | None,
     ) -> Self:
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries[Any]]:
             input_series_list = self._call(df)
             output_names = [input_series.name for input_series in input_series_list]
             result = [function(series) for series in input_series_list]

narwhals/_arrow/group_by.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterator
+from typing import cast
 
 import pyarrow as pa
 import pyarrow.compute as pc
@@ -68,7 +69,7 @@ def agg(self: Self, *exprs: ArrowExpr) -> ArrowDataFrame:
             )
             raise ValueError(msg)
 
-        aggs: list[tuple[str, str, pc.FunctionOptions | None]] = []
+        aggs: list[Any] = []
         expected_pyarrow_column_names: list[str] = self._keys.copy()
         new_column_names: list[str] = self._keys.copy()
 
@@ -91,7 +92,7 @@ def agg(self: Self, *exprs: ArrowExpr) -> ArrowDataFrame:
 
             function_name = re.sub(r"(\w+->)", "", expr._function_name)
             if function_name in {"std", "var"}:
-                option = pc.VarianceOptions(ddof=expr._kwargs["ddof"])
+                option: Any = pc.VarianceOptions(ddof=expr._kwargs["ddof"])
             elif function_name in {"len", "n_unique"}:
                 option = pc.CountOptions(mode="all")
             elif function_name == "count":
@@ -139,15 +140,26 @@ def agg(self: Self, *exprs: ArrowExpr) -> ArrowDataFrame:
 
     def __iter__(self: Self) -> Iterator[tuple[Any, ArrowDataFrame]]:
         col_token = generate_temporary_column_name(n_bytes=8, columns=self._df.columns)
-        null_token = "__null_token_value__"  # noqa: S105
+        null_token: str = "__null_token_value__"  # noqa: S105
 
         table = self._df._native_frame
-        key_values = pc.binary_join_element_wise(
-            *[pc.cast(table[key], pa.string()) for key in self._keys],
-            "",
-            null_handling="replace",
-            null_replacement=null_token,
+        # NOTE: stubs fail in multiple places for `ChunkedArray`
+        it = cast(
+            "Iterator[pa.StringArray]",
+            (table[key].cast(pa.string()) for key in self._keys),
         )
+        if TYPE_CHECKING:
+            # NOTE: stubs indicate `separator` would get appended to the end, instead of between elements
+            key_values = pc.binary_join_element_wise(
+                *it, null_handling="replace", null_replacement=null_token
+            )
+        else:
+            key_values = pc.binary_join_element_wise(
+                *it,
+                "",
+                null_handling="replace",
+                null_replacement=null_token,
+            )
         table = table.add_column(i=0, field_=col_token, column=key_values)
 
         yield from (