Add map datatype to the Ibis engine implementation (#2206)

* Add map datatype to the Ibis engine implementation

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Limit generic bases that get looked up in registry

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Revert to map implementation with default datatype

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Unit test equivalence of newly-added map datatypes

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Fix the module name lookup logic for generic bases

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Undo reversion to default datatypes implementation

This reverts commit b058433.

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Test checking all parametrizations of `Map` dtypes

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

* Add test cases corresponding to scenarios in issue

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

---------

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>

Author: deepyaman

pandera/engines/engine.py

@@ -260,11 +260,13 @@ def dtype(cls: "Engine", data_type: Any) -> DataType:
             or ((NamedTuple,) if _is_namedtuple(data_type) else ())
             or typing_inspect.get_generic_bases(data_type)
         )
-        if datatype_generic_bases:
-            equivalent_data_type = None
-            for base in datatype_generic_bases:
-                equivalent_data_type = registry.get_equivalent(base)
-                break
+        if datatype_generic_bases and inspect.getmodule(
+            base := datatype_generic_bases[0]
+        ).__name__ in {  # type: ignore[union-attr]
+            *sys.stdlib_module_names,
+            "typing_extensions",
+        }:
+            equivalent_data_type = registry.get_equivalent(base)
             if equivalent_data_type is None:
                 raise TypeError(
                     f"Type '{data_type}' not understood by {cls.__name__}."

pandera/engines/ibis_engine.py

@@ -58,7 +58,7 @@ def check(
         data_container: ibis.Table | None = None,
     ) -> Union[bool, Iterable[bool]]:
         try:
-            return self.type == pandera_dtype.type
+            return self.type == Engine.dtype(pandera_dtype).type
         except TypeError:
             return False
 
@@ -484,3 +484,38 @@ def from_parametrized_dtype(cls, ibis_dtype: dt.Interval):
         """Convert a :class:`dt.Interval` to a Pandera
         :class:`~pandera.engines.ibis_engine.Timedelta`."""
         return cls(unit=ibis_dtype.unit)
+
+
+###############################################################################
+# nested
+###############################################################################
+
+
+@Engine.register_dtype(
+    equivalents=[
+        dict,
+        dt.Map,
+    ]
+)
+@immutable(init=True)
+class Map(DataType):
+    """Semantic representation of a :class:`dt.Map`."""
+
+    type: dt.Map
+
+    def __init__(
+        self,
+        key_type: dt.DataType | None = None,
+        value_type: dt.DataType | None = None,
+    ):
+        if key_type is not None and value_type is not None:
+            object.__setattr__(self, "type", dt.Map(key_type, value_type))
+
+    @classmethod
+    def from_parametrized_dtype(cls, ibis_dtype: dt.Map):
+        """Convert a :class:`dt.Map` to a Pandera
+        :class:`~pandera.engines.ibis_engine.Map`."""
+        return cls(
+            key_type=ibis_dtype.key_type,
+            value_type=ibis_dtype.value_type,
+        )

pandera/engines/polars_engine.py

@@ -623,7 +623,7 @@ def __init__(
         elif shape is not None:
             kwargs["shape"] = shape
 
-        if inner:
+        if inner is not None:
             object.__setattr__(self, "type", pl.Array(inner=inner, **kwargs))
 
     @classmethod
@@ -655,7 +655,7 @@ def __init__(
         self,
         inner: PolarsDataType | None = None,
     ) -> None:
-        if inner:
+        if inner is not None:
             object.__setattr__(self, "type", pl.List(inner=inner))
 
     @classmethod
@@ -674,7 +674,7 @@ def __init__(
         self,
         fields: Union[Sequence[pl.Field], SchemaDict] | None = None,
     ) -> None:
-        if fields:
+        if fields is not None:
             object.__setattr__(self, "type", pl.Struct(fields=fields))
 
     @classmethod

tests/ibis/test_ibis_dtypes.py

@@ -5,9 +5,11 @@
 import pytest
 from hypothesis import given, settings
 from hypothesis import strategies as st
+from ibis import _
 from polars.testing import assert_frame_equal
 from polars.testing.parametric import dataframes
 
+import pandera.ibis as pa
 from pandera.engines import ibis_engine as ie
 
 NUMERIC_TYPES = [
@@ -22,6 +24,10 @@
     ie.String,
 ]
 
+SPECIAL_TYPES = [
+    ie.Map,
+]
+
 ALL_TYPES = NUMERIC_TYPES + TEMPORAL_TYPES + OTHER_TYPES
 
 
@@ -71,7 +77,7 @@ def test_coerce_cast(from_dtype, to_dtype, strategy, data):
         assert dtype == to_dtype.type
 
 
-@pytest.mark.parametrize("dtype", ALL_TYPES)
+@pytest.mark.parametrize("dtype", ALL_TYPES + SPECIAL_TYPES)
 def test_check_not_equivalent(dtype):
     """Test that check() rejects non-equivalent dtypes."""
     if str(ie.Engine.dtype(dtype)) == "string":
@@ -82,7 +88,7 @@ def test_check_not_equivalent(dtype):
     assert not actual_dtype.check(expected_dtype)
 
 
-@pytest.mark.parametrize("dtype", ALL_TYPES)
+@pytest.mark.parametrize("dtype", ALL_TYPES + SPECIAL_TYPES)
 def test_check_equivalent(dtype):
     """Test that check() accepts equivalent dtypes."""
     actual_dtype = ie.Engine.dtype(dtype)
@@ -130,3 +136,57 @@ def test_ibis_decimal_from_parametrized_dtype(ibis_dtype, expected_dtype):
 
     assert pandera_dtype.precision == expected_dtype.precision
     assert pandera_dtype.scale == expected_dtype.scale
+
+
+@pytest.mark.parametrize("key_dtype", ALL_TYPES)
+@pytest.mark.parametrize("value_dtype", ALL_TYPES)
+def test_ibis_map_nested_type(key_dtype, value_dtype):
+    ibis_dtype = dt.Map(key_dtype.type(), value_dtype.type())
+    pandera_dtype = ie.Engine.dtype(ibis_dtype)
+
+    assert pandera_dtype.check(ibis_dtype)
+    assert pandera_dtype.check(pandera_dtype)
+
+
+def test_ibis_map_type():
+    # https://github.com/unionai-oss/pandera/issues/2201
+    data = {
+        "id": ["01", "02", "03"],
+        "key_col": [
+            [1, 2, 3],
+            [
+                1,
+                2,
+            ],
+            None,
+        ],
+        "value_col": [
+            ["value_1A", "value_1B", "value_1C"],
+            ["value_2A", "value_2B"],
+            None,
+        ],
+    }
+    df = (
+        ibis.memtable(data)
+        .mutate(dict_col=ibis.map(_.key_col, _.value_col))
+        .drop("key_col", "value_col")
+    )
+
+    class ValidateSchema(pa.DataFrameModel):
+        id: str = pa.Field(nullable=False)
+        dict_col: ibis.dtype("map<int64, string>") = pa.Field(nullable=True)
+
+    ValidateSchema.validate(df)
+
+    class ValidateSchema(pa.DataFrameModel):
+        id: str = pa.Field(nullable=False)
+        dict_col: dt.Map = pa.Field(nullable=True)
+
+    ValidateSchema.validate(df)
+
+    class ValidateSchema(pa.DataFrameModel):
+        id: str = pa.Field(nullable=False)
+        dict_col: ibis.dtype("map<string, string>") = pa.Field(nullable=True)
+
+    with pytest.raises(pa.errors.SchemaError):
+        ValidateSchema.validate(df)

tests/polars/test_polars_dtypes.py

@@ -354,7 +354,7 @@ def test_polars_decimal_from_parametrized_dtype(polars_dtype, expected_dtype):
 )
 @given(st.integers(min_value=2, max_value=10))
 @settings(max_examples=5)
-def test_polars_nested_array_type_check(inner_dtype_cls, width):
+def test_polars_array_nested_type(inner_dtype_cls, width):
     polars_dtype = pl.Array(inner_dtype_cls(), width)
     pandera_dtype = pe.Engine.dtype(polars_dtype)