Fix reference MapType dataset comparisons

src/databricks/labs/dqx/check_funcs.py

@@ -2982,7 +2982,12 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) ->
         ref_df = _get_ref_df(ref_df_name, ref_table, ref_dfs, spark)
 
         # map type columns must be skipped as they cannot be compared with eqNullSafe
-        map_type_columns = {field.name for field in df.schema.fields if isinstance(field.dataType, types.MapType)}
+        map_type_columns = {
+            field.name
+            for schema in (df.schema, ref_df.schema)
+            for field in schema.fields
+            if isinstance(field.dataType, types.MapType)
+        }
 
         # columns to compare: present in both df and ref_df, not in PK, not excluded, not map type
         compare_columns = [

tests/integration/test_dataset_checks.py

@@ -2831,6 +2831,26 @@ def test_dataset_compare_ref_as_table_and_skip_map_col(spark: SparkSession, set_
     assertDataFrameEqual(actual, expected)
 
 
+@pytest.mark.parametrize(
+    "schema, value, ref_schema, ref_value",
+    [
+        ("id int, value string", "source", "id int, value map<string, string>", {"key": "reference"}),
+        ("id int, value map<string, string>", {"key": "source"}, "id int, value string", "reference"),
+    ],
+)
+def test_dataset_compare_skips_map_col_from_either_schema(
+    spark: SparkSession, schema: str, value: Any, ref_schema: str, ref_value: Any
+):
+    df = spark.createDataFrame([[1, value]], schema)
+    ref_df = spark.createDataFrame([[1, ref_value]], ref_schema)
+    condition, apply = compare_datasets(columns=["id"], ref_columns=["id"], ref_df_name="ref_df")
+
+    actual = apply(df, spark, {"ref_df": ref_df}).select(*df.columns, condition)
+    expected = spark.createDataFrame([[1, value, None]], f"{schema}, {get_column_name_or_alias(condition)} string")
+
+    assertDataFrameEqual(actual, expected)
+
+
 def test_dataset_compare_with_no_columns_to_compare_and_check_missing(spark: SparkSession):
     schema = "id long"