chispa/tests/test_dataframe_comparer.py at d8b4e518b3fb255bc5b50d95abf4a1e1095fa0a2 · MrPowers/chispa · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
from __future__ import annotations

import math

import pytest
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType, StringType, StructField, StructType

from chispa import DataFramesNotEqualError, assert_approx_df_equality, assert_df_equality
from chispa.dataframe_comparer import are_dfs_equal
from chispa.schema_comparer import SchemasNotEqualError


def describe_assert_df_equality():
    def it_throws_with_schema_mismatches(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li"), (3, "laura")]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        with pytest.raises(SchemasNotEqualError):
            assert_df_equality(df1, df2)

    def it_can_work_with_different_row_orders(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [(2, "li"), (1, "jose")]
        df2 = spark.createDataFrame(data2, ["num", "name"])
        assert_df_equality(df1, df2, transforms=[lambda df: df.sort(df.columns)])

    def it_can_work_with_different_row_orders_with_a_flag(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [(2, "li"), (1, "jose")]
        df2 = spark.createDataFrame(data2, ["num", "name"])
        assert_df_equality(df1, df2, ignore_row_order=True)

    def it_can_work_with_struct_columns_and_ignore_row_order(spark: SparkSession):
        data1 = [((1, "jose"),), ((2, "li"),)]
        df1 = spark.createDataFrame(data1, ["person"])
        data2 = [((2, "li"),), ((1, "jose"),)]
        df2 = spark.createDataFrame(data2, ["person"])
        assert_df_equality(df1, df2, ignore_row_order=True)

    def it_can_work_with_mixed_columns_and_ignore_row_order(spark: SparkSession):
        data1 = [((1, "jose"), 100), ((2, "li"), 200)]
        df1 = spark.createDataFrame(data1, ["person", "score"])
        data2 = [((2, "li"), 200), ((1, "jose"), 100)]
        df2 = spark.createDataFrame(data2, ["person", "score"])
        assert_df_equality(df1, df2, ignore_row_order=True)

    def it_can_work_with_nested_struct_columns_and_ignore_row_order(spark: SparkSession):
        data1 = [(((1, "jose"), 30),), (((2, "li"), 40),)]
        df1 = spark.createDataFrame(data1, ["nested_person"])
        data2 = [(((2, "li"), 40),), (((1, "jose"), 30),)]
        df2 = spark.createDataFrame(data2, ["nested_person"])
        assert_df_equality(df1, df2, ignore_row_order=True)

    def it_can_work_with_different_row_and_column_orders(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [("li", 2), ("jose", 1)]
        df2 = spark.createDataFrame(data2, ["name", "num"])
        assert_df_equality(df1, df2, ignore_row_order=True, ignore_column_order=True)

    def it_raises_for_row_insensitive_with_diff_content(spark: SparkSession):
        data1 = [(1, "XXXX"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [(2, "li"), (1, "jose")]
        df2 = spark.createDataFrame(data2, ["num", "name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_df_equality(df1, df2, transforms=[lambda df: df.sort(df.columns)])

    def it_throws_with_schema_column_order_mismatch(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [("jose", 1), ("li", 1)]
        df2 = spark.createDataFrame(data2, ["name", "num"])
        with pytest.raises(SchemasNotEqualError):
            assert_df_equality(df1, df2)

    def it_does_not_throw_on_schema_column_order_mismatch_with_transforms(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [("jose", 1), ("li", 2)]
        df2 = spark.createDataFrame(data2, ["name", "num"])
        assert_df_equality(df1, df2, transforms=[lambda df: df.select(sorted(df.columns))])

    def it_throws_with_schema_mismatch(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li")]
        df1 = spark.createDataFrame(data1, ["num", "different_name"])
        data2 = [("jose", 1), ("li", 2)]
        df2 = spark.createDataFrame(data2, ["name", "num"])
        with pytest.raises(SchemasNotEqualError):
            assert_df_equality(df1, df2, transforms=[lambda df: df.select(sorted(df.columns))])

    def it_throws_with_content_mismatches(spark: SparkSession):
        data1 = [("jose", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_df_equality(df1, df2)

    def it_throws_with_length_mismatches(spark: SparkSession):
        data1 = [("jose", "jose"), ("li", "li"), ("laura", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("jose", "jose"), ("li", "li")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_df_equality(df1, df2)

    def it_can_consider_nan_values_equal(spark: SparkSession):
        data1 = [(float("nan"), "jose"), (2.0, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [(float("nan"), "jose"), (2.0, "li")]
        df2 = spark.createDataFrame(data2, ["num", "name"])
        assert_df_equality(df1, df2, allow_nan_equality=True)

    def it_does_not_consider_nan_values_equal_by_default(spark: SparkSession):
        data1 = [(float("nan"), "jose"), (2.0, "li")]
        df1 = spark.createDataFrame(data1, ["num", "name"])
        data2 = [(float("nan"), "jose"), (2.0, "li")]
        df2 = spark.createDataFrame(data2, ["num", "name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_df_equality(df1, df2, allow_nan_equality=False)

    def it_can_consider_nan_values_equal_in_array_fields(spark: SparkSession):
        data1 = [([1.0, float("nan"), 3.0], "jose"), ([4.0, 5.0], "li")]
        df1 = spark.createDataFrame(data1, ["nums", "name"])
        data2 = [([1.0, float("nan"), 3.0], "jose"), ([4.0, 5.0], "li")]
        df2 = spark.createDataFrame(data2, ["nums", "name"])
        assert_df_equality(df1, df2, allow_nan_equality=True)

    def it_raises_when_array_nan_positions_are_different_with_allow_nan_equality(spark: SparkSession):
        data1 = [([1.0, float("nan"), 3.0], "jose"), ([4.0, 5.0], "li")]
        df1 = spark.createDataFrame(data1, ["nums", "name"])
        data2 = [([float("nan"), 1.0, 3.0], "jose"), ([4.0, 5.0], "li")]
        df2 = spark.createDataFrame(data2, ["nums", "name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_df_equality(df1, df2, allow_nan_equality=True)

    def it_does_not_consider_nan_values_equal_in_array_fields_by_default(spark: SparkSession):
        data1 = [([1.0, float("nan"), 3.0], "jose"), ([4.0, 5.0], "li")]
        df1 = spark.createDataFrame(data1, ["nums", "name"])
        data2 = [([1.0, float("nan"), 3.0], "jose"), ([4.0, 5.0], "li")]
        df2 = spark.createDataFrame(data2, ["nums", "name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_df_equality(df1, df2, allow_nan_equality=False)

    def it_can_ignore_metadata(spark: SparkSession):
        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(spark: SparkSession):
        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):
            assert_df_equality(df1, df2)

    def it_can_ignore_columns(spark: SparkSession):
        data1 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "boo"), ("luisa", "boo")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        assert_df_equality(df1, df2, ignore_columns=["expected_name"])

    def it_throws_when_dfs_are_not_same_with_ignored_columns(spark: SparkSession):
        data1 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "boo"), ("luisa", "boo")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        with pytest.raises(DataFramesNotEqualError):
            assert assert_df_equality(df1, df2, ignore_columns=["name"])

    def it_works_when_sorting_and_dropping_columns(spark: SparkSession):
        data1 = [("b", "jose", 10), ("a", "jose", 20)]
        df1 = spark.createDataFrame(data1, ["ignore_me", "name", "score"])
        data2 = [("a", "jose", 10), ("b", "jose", 20)]
        df2 = spark.createDataFrame(data2, ["ignore_me", "name", "score"])
        assert_df_equality(df1, df2, ignore_columns=["ignore_me"], ignore_row_order=True)


def describe_are_dfs_equal():
    def it_returns_false_with_schema_mismatches(spark: SparkSession):
        data1 = [(1, "jose"), (2, "li"), (3, "laura")]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        assert are_dfs_equal(df1, df2) is False

    def it_returns_false_with_content_mismatches(spark: SparkSession):
        data1 = [("jose", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        assert are_dfs_equal(df1, df2) is False

    def it_returns_true_when_dfs_are_same(spark: SparkSession):
        data1 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        assert are_dfs_equal(df1, df2) is True


def describe_assert_approx_df_equality():
    def it_throws_with_content_mismatch(spark: SparkSession):
        data1 = [(1.0, "jose"), (1.1, "li"), (1.2, "laura"), (1.0, None)]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [(1.0, "jose"), (1.05, "li"), (1.0, "laura"), (None, "hi")]
        df2 = spark.createDataFrame(data2, ["num", "expected_name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_approx_df_equality(df1, df2, 0.1)

    def it_throws_with_with_length_mismatch(spark: SparkSession):
        data1 = [(1.0, "jose"), (1.1, "li"), (1.2, "laura"), (None, None)]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [(1.0, "jose"), (1.05, "li")]
        df2 = spark.createDataFrame(data2, ["num", "expected_name"])
        with pytest.raises(DataFramesNotEqualError):
            assert_approx_df_equality(df1, df2, 0.1)

    def it_does_not_throw_with_no_mismatch(spark: SparkSession):
        data1 = [(1.0, "jose"), (1.1, "li"), (1.2, "laura"), (None, None)]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [(1.0, "jose"), (1.05, "li"), (1.2, "laura"), (None, None)]
        df2 = spark.createDataFrame(data2, ["num", "expected_name"])
        assert_approx_df_equality(df1, df2, 0.1)

    def it_does_not_throw_with_different_row_col_order(spark: SparkSession):
        data1 = [(1.0, "jose"), (1.1, "li"), (1.2, "laura"), (None, None)]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [("li", 1.05), ("laura", 1.2), (None, None), ("jose", 1.0)]
        df2 = spark.createDataFrame(data2, ["expected_name", "num"])
        assert_approx_df_equality(df1, df2, 0.1, ignore_row_order=True, ignore_column_order=True)

    def it_does_not_throw_with_nan_values(spark: SparkSession):
        data1 = [
            (1.0, "jose"),
            (1.1, "li"),
            (1.2, "laura"),
            (None, None),
            (float("nan"), "buk"),
        ]
        df1 = spark.createDataFrame(data1, ["num", "expected_name"])
        data2 = [
            (1.0, "jose"),
            (1.05, "li"),
            (1.2, "laura"),
            (None, None),
            (math.nan, "buk"),
        ]
        df2 = spark.createDataFrame(data2, ["num", "expected_name"])
        assert_approx_df_equality(df1, df2, 0.1, allow_nan_equality=True)

    def it_can_ignore_columns(spark: SparkSession):
        data1 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "boo"), ("luisa", "boo")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        assert_approx_df_equality(df1, df2, 0.1, ignore_columns=["expected_name"])

    def it_throws_when_dfs_are_not_same_with_ignored_columns(spark: SparkSession):
        data1 = [("bob", "jose"), ("li", "li"), ("luisa", "laura")]
        df1 = spark.createDataFrame(data1, ["name", "expected_name"])
        data2 = [("bob", "jose"), ("li", "boo"), ("luisa", "boo")]
        df2 = spark.createDataFrame(data2, ["name", "expected_name"])
        with pytest.raises(DataFramesNotEqualError):
            assert assert_approx_df_equality(df1, df2, 0.1, ignore_columns=["name"])

    def it_can_ignore_metadata(spark: SparkSession):
        schema1 = StructType([
            StructField("num", IntegerType(), True, {"comment": "a"}),
            StructField("name", StringType(), True),
        ])
        schema2 = StructType([
            StructField("num", IntegerType(), True, {"comment": "b"}),
            StructField("name", StringType(), True),
        ])
        df1 = spark.createDataFrame([(1, "jose"), (2, "li")], schema=schema1)
        df2 = spark.createDataFrame([(1, "jose"), (2, "li")], schema=schema2)
        assert_approx_df_equality(df1, df2, 0.1, ignore_metadata=True)

    def it_does_not_throw_with_struct_columns_and_ignore_row_order(spark: SparkSession):
        data1 = [((1.0, "jose"),), ((1.1, "li"),)]
        df1 = spark.createDataFrame(data1, ["person"])
        data2 = [((1.1, "li"),), ((1.0, "jose"),)]
        df2 = spark.createDataFrame(data2, ["person"])
        assert_approx_df_equality(df1, df2, 0.1, ignore_row_order=True)