Add option to ignore metadata when comparing DataFrames

Author: MrPowers

chispa/dataframe_comparer.py

@@ -10,7 +10,7 @@ class DataFramesNotEqualError(Exception):
 
 
 def assert_df_equality(df1, df2, ignore_nullable=False, transforms=None, allow_nan_equality=False,
-                       ignore_column_order=False, ignore_row_order=False, underline_cells=False):
+                       ignore_column_order=False, ignore_row_order=False, underline_cells=False, ignore_metadata=False):
     if transforms is None:
         transforms = []
     if ignore_column_order:
@@ -19,7 +19,7 @@ def assert_df_equality(df1, df2, ignore_nullable=False, transforms=None, allow_n
         transforms.append(lambda df: df.sort(df.columns))
     df1 = reduce(lambda acc, fn: fn(acc), transforms, df1)
     df2 = reduce(lambda acc, fn: fn(acc), transforms, df2)
-    assert_schema_equality(df1.schema, df2.schema, ignore_nullable)
+    assert_schema_equality(df1.schema, df2.schema, ignore_nullable, ignore_metadata)
     if allow_nan_equality:
         assert_generic_rows_equality(
             df1.collect(), df2.collect(), are_rows_equal_enhanced, [True], underline_cells=underline_cells)

chispa/schema_comparer.py

@@ -8,13 +8,37 @@ class SchemasNotEqualError(Exception):
    pass
 
 
-def assert_schema_equality(s1, s2, ignore_nullable=False):
-    if ignore_nullable:
-        assert_schema_equality_ignore_nullable(s1, s2)
-    else:
+def assert_schema_equality(s1, s2, ignore_nullable=False, ignore_metadata=False):
+    if not ignore_nullable and not ignore_metadata:
         assert_basic_schema_equality(s1, s2)
+    else:
+        assert_schema_equality_full(s1, s2, ignore_nullable, ignore_metadata)
+
 
+def assert_schema_equality_full(s1, s2, ignore_nullable=False, ignore_metadata=False):
+    def inner(s1, s2, ignore_nullable, ignore_metadata):
+        if len(s1) != len(s2):
+            return False
+        zipped = list(six.moves.zip_longest(s1, s2))
+        for sf1, sf2 in zipped:
+            if not are_structfields_equal(sf1, sf2, ignore_nullable, ignore_metadata):
+                return False
+        return True
 
+    if not inner(s1, s2, ignore_nullable, ignore_metadata):
+        t = PrettyTable(["schema1", "schema2"])
+        zipped = list(six.moves.zip_longest(s1, s2))
+        for sf1, sf2 in zipped:
+            if are_structfields_equal(sf1, sf2, True):
+                t.add_row([blue(sf1), blue(sf2)])
+            else:
+                t.add_row([sf1, sf2])
+        raise SchemasNotEqualError("\n" + t.get_string())
+
+
+# deprecate this
+# perhaps it is a little faster, but do we really need this?
+# I think schema equality operations are really fast to begin with
 def assert_basic_schema_equality(s1, s2):
     if s1 != s2:
         t = PrettyTable(["schema1", "schema2"])
@@ -27,8 +51,10 @@ def assert_basic_schema_equality(s1, s2):
         raise SchemasNotEqualError("\n" + t.get_string())
 
 
+
+# deprecate this.  ignore_nullable should be a flag.
 def assert_schema_equality_ignore_nullable(s1, s2):
-    if are_schemas_equal_ignore_nullable(s1, s2) == False:
+    if not are_schemas_equal_ignore_nullable(s1, s2):
         t = PrettyTable(["schema1", "schema2"])
         zipped = list(six.moves.zip_longest(s1, s2))
         for sf1, sf2 in zipped:
@@ -39,6 +65,7 @@ def assert_schema_equality_ignore_nullable(s1, s2):
         raise SchemasNotEqualError("\n" + t.get_string())
 
 
+# deprecate this.  ignore_nullable should be a flag.
 def are_schemas_equal_ignore_nullable(s1, s2):
     if len(s1) != len(s2):
         return False
@@ -49,21 +76,25 @@ def are_schemas_equal_ignore_nullable(s1, s2):
     return True
 
 
-def are_structfields_equal(sf1, sf2, ignore_nullability=False):
-    if ignore_nullability:
+# "ignore_nullability" should be "ignore_nullable" for consistent terminology
+def are_structfields_equal(sf1, sf2, ignore_nullability=False, ignore_metadata=False):
+    if not ignore_nullability and not ignore_metadata:
+        return sf1 == sf2
+    else:
         if sf1 is None or sf2 is None:
             if sf1 is None and sf2 is None:
                 return True
             else:
                 return False
         if sf1.name != sf2.name:
             return False
+        if not ignore_metadata and sf1.metadata != sf2.metadata:
+            return False
         else:
             return are_datatypes_equal_ignore_nullable(sf1.dataType, sf2.dataType)
-    else:
-        return sf1 == sf2
 
 
+# deprecate this
 def are_datatypes_equal_ignore_nullable(dt1, dt2):
     """Checks if datatypes are equal, descending into structs and arrays to
     ignore nullability.

tests/test_column_comparer.py

@@ -1,6 +1,6 @@
 import pytest
 
-from .spark import *
+from .spark import spark
 from chispa import *
 
 

tests/test_dataframe_comparer.py

@@ -5,6 +5,7 @@
 from chispa.dataframe_comparer import are_dfs_equal
 from chispa.schema_comparer import SchemasNotEqualError
 import math
+from pyspark.sql.types import StringType, IntegerType, StructType, StructField
 
 
 def describe_assert_df_equality():
@@ -17,7 +18,6 @@ def it_throws_with_schema_mismatches():
             assert_df_equality(df1, df2)
 
 
-
     def it_can_work_with_different_row_orders():
         data1 = [(1, "jose"), (2, "li")]
         df1 = spark.createDataFrame(data1, ["num", "name"])
@@ -114,6 +114,46 @@ def it_does_not_consider_nan_values_equal_by_default():
             assert_df_equality(df1, df2, allow_nan_equality=False)
 
 
+    def it_can_ignore_metadata():
+        rows_data = [("jose", 1), ("li", 2), ("luisa", 3)]
+        schema1 = StructType(
+            [
+                StructField("name", StringType(), True, {"hi": "no"}),
+                StructField("age", IntegerType(), True),
+            ]
+        )
+        schema2 = StructType(
+            [
+                StructField("name", StringType(), True, {"hi": "whatever"}),
+                StructField("age", IntegerType(), True),
+            ]
+        )
+        df1 = spark.createDataFrame(rows_data, schema1)
+        df2 = spark.createDataFrame(rows_data, schema2)
+        assert_df_equality(df1, df2, ignore_metadata=True)
+
+
+    def it_catches_mismatched_metadata():
+        rows_data = [("jose", 1), ("li", 2), ("luisa", 3)]
+        schema1 = StructType(
+            [
+                StructField("name", StringType(), True, {"hi": "no"}),
+                StructField("age", IntegerType(), True),
+            ]
+        )
+        schema2 = StructType(
+            [
+                StructField("name", StringType(), True, {"hi": "whatever"}),
+                StructField("age", IntegerType(), True),
+            ]
+        )
+        df1 = spark.createDataFrame(rows_data, schema1)
+        df2 = spark.createDataFrame(rows_data, schema2)
+        with pytest.raises(SchemasNotEqualError) as e_info:
+            assert_df_equality(df1, df2)
+
+
+
 def describe_are_dfs_equal():
     def it_returns_false_with_schema_mismatches():
         data1 = [(1, "jose"), (2, "li"), (3, "laura")]

tests/test_schema_comparer.py

@@ -124,7 +124,7 @@ def it_returns_false_when_columns_have_different_order():
         assert are_schemas_equal_ignore_nullable(s1, s2) == False
 
 
-def describe_are_structfield_types_equal_ignore_nullable():
+def describe_are_structfields_equal():
     def it_returns_true_when_only_nullable_flag_is_different_within_array_element():
         s1 = StructField("coords", ArrayType(DoubleType(), True), True)
         s2 = StructField("coords", ArrayType(DoubleType(), False), True)
@@ -159,3 +159,22 @@ def it_returns_true_when_different_nullability_within_struct():
         s1 = StructField("coords", StructType([StructField("hello", DoubleType(), True)]), True)
         s2 = StructField("coords", StructType([StructField("hello", DoubleType(), False)]), True)
         assert are_structfields_equal(s1, s2, True) == True
+    def it_returns_false_when_metadata_differs():
+        s1 = StructField("coords", StringType(), True, {"hi": "whatever"})
+        s2 = StructField("coords", StringType(), True, {"hi": "no"})
+        assert are_structfields_equal(s1, s2, ignore_nullability=True, ignore_metadata=False) is False
+
+    def it_allows_metadata_to_be_ignored():
+        s1 = StructField("coords", StringType(), True, {"hi": "whatever"})
+        s2 = StructField("coords", StringType(), True, {"hi": "no"})
+        assert are_structfields_equal(s1, s2, ignore_nullability=False, ignore_metadata=True) is True
+
+    def it_allows_nullability_and_metadata_to_be_ignored():
+        s1 = StructField("coords", StringType(), True, {"hi": "whatever"})
+        s2 = StructField("coords", StringType(), False, {"hi": "no"})
+        assert are_structfields_equal(s1, s2, ignore_nullability=True, ignore_metadata=True) is True
+
+    def it_returns_true_when_metadata_is_the_same():
+        s1 = StructField("coords", StringType(), True, {"hi": "whatever"})
+        s2 = StructField("coords", StringType(), True, {"hi": "whatever"})
+        assert are_structfields_equal(s1, s2, ignore_nullability=True, ignore_metadata=False) is True