feat: allow inspecting the inner type / length of nw.Array (narwhals-…

…dev#1136) --------- Co-authored-by: Marco Gorelli <[email protected]>
DeaMariaLeon · Oct 6, 2024 · 3e0405d · 3e0405d
1 parent 5aa4e12
commit 3e0405d
Show file tree

Hide file tree

Showing 7 changed files with 90 additions and 9 deletions.
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
@@ -60,7 +60,9 @@ def native_to_narwhals_dtype(dtype: Any, dtypes: DTypes) -> DType:
     if pa.types.is_list(dtype) or pa.types.is_large_list(dtype):
         return dtypes.List(native_to_narwhals_dtype(dtype.value_type, dtypes))
     if pa.types.is_fixed_size_list(dtype):
-        return dtypes.Array()
+        return dtypes.Array(
+            native_to_narwhals_dtype(dtype.value_type, dtypes), dtype.list_size
+        )
     return dtypes.Unknown()  # pragma: no cover
 
 

diff --git a/narwhals/_duckdb/dataframe.py b/narwhals/_duckdb/dataframe.py
@@ -52,8 +52,11 @@ def map_duckdb_dtype_to_narwhals_dtype(duckdb_dtype: Any, dtypes: DTypes) -> DTy
         return dtypes.Struct()
     if match_ := re.match(r"(.*)\[\]$", duckdb_dtype):
         return dtypes.List(map_duckdb_dtype_to_narwhals_dtype(match_.group(1), dtypes))
-    if re.match(r"\w+\[\d+\]", duckdb_dtype):
-        return dtypes.Array()
+    if match_ := re.match(r"(\w+)\[(\d+)\]", duckdb_dtype):
+        return dtypes.Array(
+            map_duckdb_dtype_to_narwhals_dtype(match_.group(1), dtypes),
+            int(match_.group(2)),
+        )
     return dtypes.Unknown()
 
 

diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
@@ -285,7 +285,10 @@ def native_to_narwhals_dtype(column: Any, dtypes: DTypes) -> DType:
             arrow_native_to_narwhals_dtype(column.dtype.pyarrow_dtype.value_type, dtypes)
         )
     if dtype.startswith("fixed_size_list"):
-        return dtypes.Array()
+        return dtypes.Array(
+            arrow_native_to_narwhals_dtype(column.dtype.pyarrow_dtype.value_type, dtypes),
+            column.dtype.pyarrow_dtype.list_size,
+        )
     if dtype.startswith("struct"):
         return dtypes.Struct()
     if dtype == "object":

diff --git a/narwhals/_polars/utils.py b/narwhals/_polars/utils.py
@@ -8,6 +8,8 @@
     from narwhals.dtypes import DType
     from narwhals.typing import DTypes
 
+from narwhals.utils import parse_version
+
 
 def extract_native(obj: Any) -> Any:
     from narwhals._polars.dataframe import PolarsDataFrame
@@ -77,7 +79,12 @@ def native_to_narwhals_dtype(dtype: Any, dtypes: DTypes) -> DType:
     if dtype == pl.List:
         return dtypes.List(native_to_narwhals_dtype(dtype.inner, dtypes))
     if dtype == pl.Array:
-        return dtypes.Array()
+        if parse_version(pl.__version__) < (1, 0):  # pragma: no cover
+            return dtypes.Array(
+                native_to_narwhals_dtype(dtype.inner, dtypes), dtype.width
+            )
+        else:
+            return dtypes.Array(native_to_narwhals_dtype(dtype.inner, dtypes), dtype.size)
     return dtypes.Unknown()
 
 

diff --git a/narwhals/dtypes.py b/narwhals/dtypes.py
@@ -200,7 +200,35 @@ def __repr__(self) -> str:
         return f"{class_name}({self.inner!r})"
 
 
-class Array(DType): ...
+class Array(DType):
+    def __init__(self, inner: DType | type[DType], width: int | None = None) -> None:
+        self.inner = inner
+        if width is None:
+            error = "`width` must be specified when initializing an `Array`"
+            raise TypeError(error)
+        self.width = width
+
+    def __eq__(self, other: DType | type[DType]) -> bool:  # type: ignore[override]
+        # This equality check allows comparison of type classes and type instances.
+        # If a parent type is not specific about its inner type, we infer it as equal:
+        # > array[i64] == array[i64] -> True
+        # > array[i64] == array[f32] -> False
+        # > array[i64] == array      -> True
+
+        # allow comparing object instances to class
+        if type(other) is type and issubclass(other, self.__class__):
+            return True
+        elif isinstance(other, self.__class__):
+            return self.inner == other.inner
+        else:
+            return False
+
+    def __hash__(self) -> int:
+        return hash((self.__class__, self.inner, self.width))
+
+    def __repr__(self) -> str:
+        class_name = self.__class__.__name__
+        return f"{class_name}({self.inner!r}, {self.width})"
 
 
 class Date(TemporalType): ...
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
@@ -7,6 +7,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import pyarrow as pa
 import pytest
 
@@ -67,6 +68,42 @@ def test_list_valid() -> None:
     assert dtype in {nw.List(nw.List(nw.Int64))}
 
 
+def test_array_valid() -> None:
+    dtype = nw.Array(nw.Int64, 2)
+    assert dtype == nw.Array(nw.Int64, 2)
+    assert dtype == nw.Array
+    assert dtype != nw.Array(nw.Float32, 2)
+    assert dtype != nw.Duration
+    assert repr(dtype) == "Array(<class 'narwhals.dtypes.Int64'>, 2)"
+    dtype = nw.Array(nw.Array(nw.Int64, 2), 2)
+    assert dtype == nw.Array(nw.Array(nw.Int64, 2), 2)
+    assert dtype == nw.Array
+    assert dtype != nw.Array(nw.Array(nw.Float32, 2), 2)
+    assert dtype in {nw.Array(nw.Array(nw.Int64, 2), 2)}
+
+    with pytest.raises(
+        TypeError, match="`width` must be specified when initializing an `Array`"
+    ):
+        dtype = nw.Array(nw.Int64)
+
+
+@pytest.mark.skipif(
+    parse_version(pl.__version__) < (1,) or parse_version(pd.__version__) < (2, 2),
+    reason="`shape` is only available after 1.0",
+)
+def test_polars_2d_array() -> None:
+    df = pl.DataFrame(
+        {"a": [[[1, 2], [3, 4], [5, 6]]]}, schema={"a": pl.Array(pl.Int64, (3, 2))}
+    )
+    assert nw.from_native(df).collect_schema()["a"] == nw.Array(nw.Array(nw.Int64, 2), 3)
+    assert nw.from_native(df.to_arrow()).collect_schema()["a"] == nw.Array(
+        nw.Array(nw.Int64, 2), 3
+    )
+    assert nw.from_native(
+        df.to_pandas(use_pyarrow_extension_array=True)
+    ).collect_schema()["a"] == nw.Array(nw.Array(nw.Int64, 2), 3)
+
+
 def test_second_time_unit() -> None:
     s = pd.Series(np.array([np.datetime64("2020-01-01", "s")]))
     result = nw.from_native(s, series_only=True)

diff --git a/tests/frame/schema_test.py b/tests/frame/schema_test.py
@@ -213,22 +213,23 @@ def test_nested_dtypes() -> None:
         schema_overrides={"b": pl.Array(pl.Int64, 2)},
     ).to_pandas(use_pyarrow_extension_array=True)
     nwdf = nw.from_native(df)
+
     assert nwdf.schema == {"a": nw.List, "b": nw.Array, "c": nw.Struct}
     df = pl.DataFrame(
         {"a": [[1, 2]], "b": [[1, 2]], "c": [{"a": 1}]},
         schema_overrides={"b": pl.Array(pl.Int64, 2)},
     )
     nwdf = nw.from_native(df)
-    assert nwdf.schema == {"a": nw.List, "b": nw.Array, "c": nw.Struct}
+    assert nwdf.schema == {"a": nw.List, "b": nw.Array(nw.Int64, 2), "c": nw.Struct}
     df = pl.DataFrame(
         {"a": [[1, 2]], "b": [[1, 2]], "c": [{"a": 1}]},
         schema_overrides={"b": pl.Array(pl.Int64, 2)},
     ).to_arrow()
     nwdf = nw.from_native(df)
-    assert nwdf.schema == {"a": nw.List, "b": nw.Array, "c": nw.Struct}
+    assert nwdf.schema == {"a": nw.List, "b": nw.Array(nw.Int64, 2), "c": nw.Struct}
     df = duckdb.sql("select * from df")
     nwdf = nw.from_native(df)
-    assert nwdf.schema == {"a": nw.List, "b": nw.Array, "c": nw.Struct}
+    assert nwdf.schema == {"a": nw.List, "b": nw.Array(nw.Int64, 2), "c": nw.Struct}
 
 
 def test_nested_dtypes_ibis() -> None:  # pragma: no cover