fix: More aligning modin, pandas dtypes (#2958)

Co-authored-by: FBruzzesi <[email protected]>

Author: dangotbanned

narwhals/_pandas_like/utils.py

@@ -129,6 +129,16 @@
     "ns": "nanosecond",
 }
 
+PANDAS_VERSION = Implementation.PANDAS._backend_version()
+"""Static backend version for `pandas`.
+
+Always available if we reached here, due to a module-level import.
+"""
+
+
+def is_pandas_or_modin(implementation: Implementation) -> bool:
+    return implementation in {Implementation.PANDAS, Implementation.MODIN}
+
 
 def align_and_extract_native(
     lhs: PandasLikeSeries, rhs: PandasLikeSeries | object
@@ -499,10 +509,9 @@ def narwhals_to_native_dtype(  # noqa: C901, PLR0912, PLR0915
         # or at least, convert_dtypes(dtype_backend='pyarrow') doesn't
         # convert to it?
         return "category"
-    backend_version = implementation._backend_version()
     if isinstance_or_issubclass(dtype, dtypes.Datetime):
         # Pandas does not support "ms" or "us" time units before version 2.0
-        if implementation is Implementation.PANDAS and backend_version < (
+        if is_pandas_or_modin(implementation) and PANDAS_VERSION < (
             2,
         ):  # pragma: no cover
             dt_time_unit = "ns"
@@ -515,7 +524,7 @@ def narwhals_to_native_dtype(  # noqa: C901, PLR0912, PLR0915
         tz_part = f", {tz}" if (tz := dtype.time_zone) else ""
         return f"datetime64[{dt_time_unit}{tz_part}]"
     if isinstance_or_issubclass(dtype, dtypes.Duration):
-        if implementation is Implementation.PANDAS and backend_version < (
+        if is_pandas_or_modin(implementation) and PANDAS_VERSION < (
             2,
         ):  # pragma: no cover
             du_time_unit = "ns"
@@ -545,10 +554,7 @@ def narwhals_to_native_dtype(  # noqa: C901, PLR0912, PLR0915
     if isinstance_or_issubclass(
         dtype, (dtypes.Struct, dtypes.Array, dtypes.List, dtypes.Time, dtypes.Binary)
     ):
-        if implementation in {
-            Implementation.PANDAS,
-            Implementation.MODIN,
-        } and Implementation.PANDAS._backend_version() >= (2, 2):
+        if is_pandas_or_modin(implementation) and PANDAS_VERSION >= (2, 2):
             try:
                 import pandas as pd
                 import pyarrow as pa  # ignore-banned-import  # noqa: F401

tests/expr_and_series/cast_test.py

@@ -59,17 +59,15 @@
 SPARK_LIKE_INCOMPATIBLE_COLUMNS = {"e", "f", "g", "h", "o", "p"}
 DUCKDB_INCOMPATIBLE_COLUMNS = {"o"}
 IBIS_INCOMPATIBLE_COLUMNS = {"o"}
+MODIN_XFAIL_COLUMNS = {"o", "k"}
 
 
 @pytest.mark.filterwarnings("ignore:casting period[M] values to int64:FutureWarning")
-def test_cast(constructor: Constructor, request: pytest.FixtureRequest) -> None:
+def test_cast(constructor: Constructor) -> None:
     if "pyarrow_table_constructor" in str(constructor) and PYARROW_VERSION <= (
         15,
     ):  # pragma: no cover
         pytest.skip()
-    if "modin_constructor" in str(constructor):
-        # TODO(unassigned): in modin, we end up with `'<U0'` dtype
-        request.applymarker(pytest.mark.xfail)
 
     if "pyspark" in str(constructor):
         incompatible_columns = SPARK_LIKE_INCOMPATIBLE_COLUMNS  # pragma: no cover
@@ -107,8 +105,18 @@ def test_cast(constructor: Constructor, request: pytest.FixtureRequest) -> None:
     }
     cast_map = {c: t for c, t in cast_map.items() if c not in incompatible_columns}
 
-    result = df.select(*[nw.col(col_).cast(dtype) for col_, dtype in cast_map.items()])
-    assert dict(result.collect_schema()) == cast_map
+    result = df.select(
+        *[nw.col(col_).cast(dtype) for col_, dtype in cast_map.items()]
+    ).collect_schema()
+
+    for (key, ltype), rtype in zip(result.items(), cast_map.values()):
+        if "modin_constructor" in str(constructor) and key in MODIN_XFAIL_COLUMNS:
+            # TODO(unassigned): in modin we end up with `'<U0'` dtype
+            # This block will act similarly to an xfail i.e. if we fix the issue, the
+            # assert will fail
+            assert ltype != rtype
+        else:
+            assert ltype == rtype, f"types differ for column {key}: {ltype}!={rtype}"
 
 
 def test_cast_series(
@@ -118,17 +126,14 @@ def test_cast_series(
         15,
     ):  # pragma: no cover
         request.applymarker(pytest.mark.xfail)
-    if "modin_constructor" in str(constructor_eager):
-        # TODO(unassigned): in modin, we end up with `'<U0'` dtype
-        request.applymarker(pytest.mark.xfail)
+
     df = (
         nw.from_native(constructor_eager(DATA))
         .select(nw.col(key).cast(value) for key, value in SCHEMA.items())
         .lazy()
         .collect()
     )
-
-    expected = {
+    cast_map = {
         "a": nw.Int32,
         "b": nw.Int16,
         "c": nw.Int8,
@@ -146,25 +151,16 @@ def test_cast_series(
         "o": nw.String,
         "p": nw.Duration,
     }
-    result = df.select(
-        df["a"].cast(nw.Int32),
-        df["b"].cast(nw.Int16),
-        df["c"].cast(nw.Int8),
-        df["d"].cast(nw.Int64),
-        df["e"].cast(nw.UInt32),
-        df["f"].cast(nw.UInt16),
-        df["g"].cast(nw.UInt8),
-        df["h"].cast(nw.UInt64),
-        df["i"].cast(nw.Float32),
-        df["j"].cast(nw.Float64),
-        df["k"].cast(nw.String),
-        df["l"].cast(nw.Datetime),
-        df["m"].cast(nw.Int8),
-        df["n"].cast(nw.Int8),
-        df["o"].cast(nw.String),
-        df["p"].cast(nw.Duration),
-    )
-    assert result.schema == expected
+    result = df.select(df[col_].cast(dtype) for col_, dtype in cast_map.items()).schema
+
+    for (key, ltype), rtype in zip(result.items(), cast_map.values()):
+        if "modin_constructor" in str(constructor_eager) and key in MODIN_XFAIL_COLUMNS:
+            # TODO(unassigned): in modin we end up with `'<U0'` dtype
+            # This block will act similarly to an xfail i.e. if we fix the issue, the
+            # assert will fail
+            assert ltype != rtype
+        else:
+            assert ltype == rtype, f"types differ for column {key}: {ltype}!={rtype}"
 
 
 def test_cast_string() -> None: