fix: correctly preserve arrow dtypes for pandas-like, improve concat_str performance (#3193)

Author: MarcoGorelli

narwhals/_pandas_like/namespace.py

@@ -332,9 +332,8 @@ def concat_str(
 
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
             expr_results = [s for _expr in exprs for s in _expr(df)]
-            align = self._series._align_full_broadcast
-            series = align(*(s.cast(string) for s in expr_results))
-            null_mask = align(*(s.is_null() for s in expr_results))
+            series = [s.cast(string) for s in expr_results]
+            null_mask = [s.is_null() for s in expr_results]
 
             if not ignore_nulls:
                 null_mask_result = reduce(operator.or_, null_mask)
@@ -345,15 +344,16 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 # NOTE: Trying to help `mypy` later
                 # error: Cannot determine type of "values"  [has-type]
                 values: list[PandasLikeSeries]
-                init_value, *values = [
+                init_value, *values = (
                     s.zip_with(~nm, "") for s, nm in zip_strict(series, null_mask)
-                ]
-                array_funcs = series[0]._array_funcs
-                sep_array = init_value.from_iterable(
-                    data=array_funcs.repeat(separator, len(init_value)),
-                    name="sep",
-                    index=init_value.native.index,
-                    context=self,
+                )
+                sep_array = init_value._with_native(
+                    init_value.__native_namespace__().Series(
+                        separator,
+                        name="sep",
+                        index=init_value.native.index,
+                        dtype=init_value.native.dtype,
+                    )
                 )
                 separators = (sep_array.zip_with(~nm, "") for nm in null_mask[:-1])
                 result = reduce(

narwhals/_pandas_like/series.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import operator
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, cast
 
 import numpy as np
 
@@ -304,6 +305,11 @@ def _scatter_in_place(self, indices: Self, values: Self) -> None:
             self.native.iloc[indices.native] = values_native
 
     def cast(self, dtype: IntoDType) -> Self:
+        if self.dtype == dtype and self.native.dtype != "object":
+            # Avoid dealing with pandas' type-system if we can. Note that it's only
+            # safe to do this if we're not starting with object dtype, see tests/expr_and_series/cast_test.py::test_cast_object_pandas
+            # for an example of why.
+            return self._with_native(self.native, preserve_broadcast=True)
         pd_dtype = narwhals_to_native_dtype(
             dtype,
             dtype_backend=get_dtype_backend(self.native.dtype, self._implementation),
@@ -387,103 +393,87 @@ def first(self) -> PythonLiteral:
     def last(self) -> PythonLiteral:
         return self.native.iloc[-1] if len(self.native) else None
 
+    def _with_binary(self, op: Callable[..., PandasLikeSeries], other: Any) -> Self:
+        ser, other_native = align_and_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[..., PandasLikeSeries], other: Any) -> Self:
+        return self._with_binary(lambda x, y: op(y, x), other).alias(self.name)
+
     def __eq__(self, other: object) -> Self:  # type: ignore[override]
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser == other).alias(self.name)
+        return self._with_binary(operator.eq, other)
 
     def __ne__(self, other: object) -> Self:  # type: ignore[override]
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser != other).alias(self.name)
+        return self._with_binary(operator.ne, other)
 
     def __ge__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser >= other).alias(self.name)
+        return self._with_binary(operator.ge, other)
 
     def __gt__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser > other).alias(self.name)
+        return self._with_binary(operator.gt, other)
 
     def __le__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser <= other).alias(self.name)
+        return self._with_binary(operator.le, other)
 
     def __lt__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser < other).alias(self.name)
+        return self._with_binary(operator.lt, other)
 
     def __and__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser & other).alias(self.name)
+        return self._with_binary(operator.and_, other)
 
     def __rand__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        ser = cast("pd.Series[Any]", ser)
-        return self._with_native(ser.__and__(other)).alias(self.name)
+        return self._with_binary_right(operator.and_, other)
 
     def __or__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser | other).alias(self.name)
+        return self._with_binary(operator.or_, other)
 
     def __ror__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        ser = cast("pd.Series[Any]", ser)
-        return self._with_native(ser.__or__(other)).alias(self.name)
+        return self._with_binary_right(operator.or_, other)
 
     def __add__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser + other).alias(self.name)
+        return self._with_binary(operator.add, other)
 
     def __radd__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__radd__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.add, other)
 
     def __sub__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser - other).alias(self.name)
+        return self._with_binary(operator.sub, other)
 
     def __rsub__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__rsub__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.sub, other)
 
     def __mul__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser * other).alias(self.name)
+        return self._with_binary(operator.mul, other)
 
     def __rmul__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__rmul__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.mul, other)
 
     def __truediv__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser / other).alias(self.name)
+        return self._with_binary(operator.truediv, other)
 
     def __rtruediv__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__rtruediv__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.truediv, other)
 
     def __floordiv__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser // other).alias(self.name)
+        return self._with_binary(operator.floordiv, other)
 
     def __rfloordiv__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__rfloordiv__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.floordiv, other)
 
     def __pow__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser**other).alias(self.name)
+        return self._with_binary(operator.pow, other)
 
     def __rpow__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__rpow__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.pow, other)
 
     def __mod__(self, other: Any) -> Self:
-        ser, other = align_and_extract_native(self, other)
-        return self._with_native(ser % other).alias(self.name)
+        return self._with_binary(operator.mod, other)
 
     def __rmod__(self, other: Any) -> Self:
-        _, other_native = align_and_extract_native(self, other)
-        return self._with_native(self.native.__rmod__(other_native)).alias(self.name)
+        return self._with_binary_right(operator.mod, other)
 
     # Unary
 

narwhals/_pandas_like/utils.py

@@ -469,7 +469,6 @@ def is_dtype_pyarrow(dtype: Any) -> TypeIs[pd.ArrowDtype]:
         None: "uint16",
     },
     dtypes.UInt8: {"pyarrow": "UInt8[pyarrow]", "numpy_nullable": "UInt8", None: "uint8"},
-    dtypes.String: {"pyarrow": "string[pyarrow]", "numpy_nullable": "string", None: str},
     dtypes.Boolean: {
         "pyarrow": "boolean[pyarrow]",
         "numpy_nullable": "boolean",
@@ -494,6 +493,25 @@ def narwhals_to_native_dtype(  # noqa: C901, PLR0912
         return pd_type
     if into_pd_type := NW_TO_PD_DTYPES_BACKEND.get(base_type):
         return into_pd_type[dtype_backend]
+    if issubclass(base_type, dtypes.String):
+        if dtype_backend == "pyarrow":
+            import pyarrow as pa  # ignore-banned-import
+
+            # Note: this is different from `string[pyarrow]`, even though the repr
+            # looks the same.
+            # >>> pd.DataFrame({'a':['foo']}, dtype='string[pyarrow]')['a'].str.len()
+            # 0    3
+            # Name: a, dtype: Int64
+            # >>> pd.DataFrame({'a':['foo']}, dtype=pd.ArrowDtype(pa.string()))['a'].str.len()
+            # 0    3
+            # Name: a, dtype: int32[pyarrow]
+            #
+            # `ArrowDType(pa.string())` is what `.convert_dtypes(dtype_backend='pyarrow')` converts to,
+            # so we use that here.
+            return pd.ArrowDtype(pa.string())
+        if dtype_backend == "numpy_nullable":
+            return "string"
+        return str
     if isinstance_or_issubclass(dtype, dtypes.Datetime):
         if is_pandas_or_modin(implementation) and PANDAS_VERSION < (
             2,
@@ -533,7 +551,7 @@ def narwhals_to_native_dtype(  # noqa: C901, PLR0912
         )
     if isinstance_or_issubclass(dtype, dtypes.Date):
         try:
-            import pyarrow as pa  # ignore-banned-import  # noqa: F401
+            import pyarrow as pa  # ignore-banned-import
         except ModuleNotFoundError as exc:  # pragma: no cover
             # BUG: Never re-raised?
             msg = "'pyarrow>=13.0.0' is required for `Date` dtype."

tests/expr_and_series/cast_test.py

@@ -13,7 +13,7 @@
     Constructor,
     ConstructorEager,
     assert_equal_data,
-    is_windows,
+    is_pyarrow_windows_no_tzdata,
     time_unit_compat,
 )
 
@@ -209,11 +209,13 @@ def test_cast_datetime_tz_aware(
         "dask" in str(constructor)
         or "duckdb" in str(constructor)
         or "cudf" in str(constructor)  # https://github.com/rapidsai/cudf/issues/16973
-        or ("pyarrow_table" in str(constructor) and is_windows())
         or "pyspark" in str(constructor)
         or "ibis" in str(constructor)
     ):
         request.applymarker(pytest.mark.xfail)
+    request.applymarker(
+        pytest.mark.xfail(is_pyarrow_windows_no_tzdata(constructor), reason="no tzdata")
+    )
 
     data = {
         "date": [
@@ -239,9 +241,11 @@ def test_cast_datetime_utc(
         "dask" in str(constructor)
         # https://github.com/eakmanrq/sqlframe/issues/406
         or "sqlframe" in str(constructor)
-        or ("pyarrow_table" in str(constructor) and is_windows())
     ):
         request.applymarker(pytest.mark.xfail)
+    request.applymarker(
+        pytest.mark.xfail(is_pyarrow_windows_no_tzdata(constructor), reason="no tzdata")
+    )
 
     data = {
         "date": [
@@ -394,3 +398,29 @@ def test_cast_typing_invalid() -> None:
 
     with pytest.raises(AttributeError):
         df.select(a.cast(nw.Array(nw.List, 2)))  # type: ignore[arg-type]
+
+
+@pytest.mark.skipif(PANDAS_VERSION < (2,), reason="too old for pyarrow")
+def test_pandas_pyarrow_dtypes() -> None:
+    s = nw.from_native(
+        pd.Series([123, None]).convert_dtypes(dtype_backend="pyarrow"), series_only=True
+    ).cast(nw.String)
+    result = s.str.len_chars().to_native()
+    assert result.dtype == "Int32[pyarrow]"
+
+    s = nw.from_native(
+        pd.Series([123, None], dtype="string[pyarrow]"), series_only=True
+    ).cast(nw.String)
+    result = s.str.len_chars().to_native()
+    assert result.dtype == "Int64"
+
+    s = nw.from_native(
+        pd.DataFrame({"a": ["foo", "bar"]}, dtype="string[pyarrow]")
+    ).select(nw.col("a").cast(nw.String))["a"]
+    assert s.to_native().dtype == "string[pyarrow]"
+
+
+def test_cast_object_pandas() -> None:
+    s = nw.from_native(pd.DataFrame({"a": [2, 3, None]}, dtype=object))["a"]
+    assert s[0] == 2
+    assert s.cast(nw.String)[0] == "2"

tests/frame/schema_test.py

@@ -387,6 +387,11 @@ def test_all_nulls_pandas() -> None:
 def test_schema_to_pandas(
     dtype_backend: DTypeBackend | Sequence[DTypeBackend] | None, expected: dict[str, Any]
 ) -> None:
+    if (
+        dtype_backend == "pyarrow"
+        or (isinstance(dtype_backend, list) and "pyarrow" in dtype_backend)
+    ) and PANDAS_VERSION < (1, 5):
+        pytest.skip()
     schema = nw.Schema(
         {
             "a": nw.Int64(),