What you will be able to do
- Shape PySpark logic as small DataFrame-in, DataFrame-out functions that a unit test can pin down
- Chain and parameterise those functions with DataFrame.transform
- Compare an actual DataFrame against an expected one with pyspark.testing.assertDataFrameEqual and know its defaults
- Check a DataFrame's schema, and spot type and nullability differences that schema inference introduces
Key concept
Transformation as a testable function — Write each piece of PySpark logic as a function that takes a DataFrame and returns a new DataFrame. A test can then give it small fake input and compare what it returns with an expected DataFrame, and production code chains the same functions together with DataFrame.transform.
1.Design functions a test can pin down
A unit test checks one self-contained unit of code, usually a function, and you run it early and often. The aim is to find wrong assumptions before they reach a scheduled job. In PySpark the unit is almost always a transformation: it receives a DataFrame and returns another one. Before any assertion helper is useful, the function has to be written so a test can predict its result exactly.
Databricks is direct about this. A function should return a single predictable outcome of a single data type. An existence check returns true or false. It never returns false in one case and the found object in another. A row count returns a non-negative whole number, never "the count, or false if nothing matched". If a function can return two different types, every test and every caller has to handle both cases.
The second rule is about inputs. Unit tests should not run against production data. This matters most for functions that add, remove or change data. The documented approach is to build fake data that looks as much like production as possible, usually with spark.createDataFrame and an explicit StructType schema, and run the function against that. Files that are not notebooks also have to create their own SparkSession. Databricks notebooks create one for you.
Checkpoint 1 of 5· Check yourself
You are writing numRowsInColumnForValue(df, columnName, columnValue) so that it is easy to unit test. Which return contract fits Databricks' guidance?
A testable function returns one predictable outcome of one data type. For a count, that means a non-negative whole number every time.
“To return the number of rows that exist, the function should return a non-negative, whole number.”Source: docs.databricks.com
Sources1
2.Chain testable steps with DataFrame.transform
After each step is its own function, you need a clean way to compose those functions in production code without wrapping them inside each other. DataFrame.transform does this. It takes a function plus any extra arguments, calls the function on the DataFrame, and returns a new DataFrame. The reference calls it concise syntax for chaining custom transformations. For testing, the benefit is that each link in the chain is a plain function. A test can call it directly on a two-row input, while the job calls it through .transform().
| Parameter | Type | What it carries |
|---|---|---|
| func | function | a function that takes and returns a DataFrame |
| *args | any | Positional arguments to pass to func |
| **kwargs | any | Keyword arguments to pass to func |
from pyspark.sql import functions as sf
df = spark.createDataFrame([(1, 1.0), (2, 2.0)], ["int", "float"])
def cast_all_to_int(input_df):
return input_df.select([sf.col(c).cast("int") for c in input_df.columns])
def sort_columns_asc(input_df):
return input_df.select(*sorted(input_df.columns))
df.transform(cast_all_to_int).transform(sort_columns_asc).show()Two columns, float then int, because sort_columns_asc orders the column names alphabetically. Both columns hold the integers 1 and 2, because cast_all_to_int cast everything to int first.
A step that needs configuration takes extra parameters. transform passes them through to the function, either by position or by keyword:
def add_n(input_df, n):
cols = [(sf.col(c) + n).alias(c) for c in input_df.columns]
return input_df.select(cols)
df.transform(add_n, 1).transform(add_n, n=10).show()Checkpoint 2 of 5· Exam question
A data engineer wrote a pytest test that calls `assertDataFrameEqual(actual_df, expected_df)` to validate a transformation that must preserve input row order for a downstream windowing calculation. The test passes even though a manual review shows the rows in `actual_df` appear in a different order than in `expected_df`. What causes this behavior, and how should the engineer fix the test?
Correct answer: A — `assertDataFrameEqual` defaults to `checkRowOrder=False`, so it sorts rows before comparing content and ignores their original sequence; passing `checkRowOrder=True` makes the assertion sensitive to row order.
- A. `assertDataFrameEqual` ships with `checkRowOrder=False` by default, meaning rows are effectively compared as a multiset rather than a sequence. Explicitly passing `checkRowOrder=True` is the documented way to make the assertion fail when row ordering changes.
- B. The function does inspect every cell value, not just row counts and names, so this description of its behavior is wrong. `ignoreColumnOrder` controls whether columns must appear in the same left-to-right position, which is unrelated to row sequence.
- C. There is no cross-run caching of assertion outcomes in PySpark's testing utilities; each call independently re-evaluates both DataFrames. Clearing `.pytest_cache` would have no effect on this comparison.
- D. The function raises `PySparkAssertionError` rather than only warning when DataFrames differ, so a genuine mismatch would stop the test. Running `assertSchemaEqual` separately would confirm schema equality but would not reveal anything about row ordering or values.
Sources2
3.Compare whole results with assertDataFrameEqual
A function like add_n produces a whole DataFrame. You could check single cells with collect() and assert, but each check covers only one cell. pyspark.testing.assertDataFrameEqual compares an actual DataFrame with an expected one in a single call. Its signature shows how lenient the comparison is when you leave the defaults alone.
| Parameter | Default |
|---|---|
| checkRowOrder | False |
| rtol | 1e-05 |
| atol | 1e-08 |
| ignoreNullable | True |
| ignoreColumnOrder | False |
| ignoreColumnName | False |
| ignoreColumnType | False |
| maxErrors | None |
| showOnlyDiff | False |
| includeDiffRows | False |
Databricks uses this helper in its own test examples. In the one below, test_spark is a test SparkSession, and test_pipeline.run executes the named output tables against mock input. status.is_success confirms that the run worked before anything is compared. The test then reads the output table and builds the expected DataFrame inline with createDataFrame. It passes only a list of column names, so each column's type is inferred from the data.
# Test 5: Full DataFrame comparison with assertDataFrameEqual
def test_counts_full_dataframe(test_spark):
mock_users(test_spark)
status = test_pipeline.run(test_spark, set(["catalog.schema.users", "catalog.schema.counts"]))
assert status.is_success
result = test_spark.table("catalog.schema.counts")
expected = test_spark.createDataFrame(
[("admin", 2, 1), ("user", 2, 1)],
schema=["user_type", "total_count", "count_valid_emails"]
)
assertDataFrameEqual(result, expected)Checkpoint 3 of 5· Check yourself
In the assertDataFrameEqual signature, which of these defaults is correct?
ignoreNullable defaults to True. checkRowOrder, ignoreColumnType and ignoreColumnOrder all default to False.
“ignoreNullable=True, ignoreColumnOrder=False, ignoreColumnName=False, ignoreColumnType=False”Source: spark.apache.org
Checkpoint 4 of 5· Exam question
Which statement accurately describes the default behavior of `assertSchemaEqual` when comparing two `StructType` schemas in a PySpark unit test?
Correct answer: A — By default, nullable flags are ignored while column names, order, and data types must match exactly, so relaxing a field from non-nullable to nullable alone will not fail the test unless `ignoreNullable=False` is passed.
- A. `assertSchemaEqual` defaults to `ignoreNullable=True`, so nullable differences alone are tolerated, while column names, positional order, and data types are still compared and must match. This is the correct description of its out-of-the-box behavior.
- B. The default is the opposite: `ignoreNullable=True` means nullable-only differences are tolerated, not treated as mismatches. `ignoreColumnOrder` also controls positional comparison, not nullable handling, so this option misattributes both defaults.
- C. Data types are compared by default, not ignored; a genuine `IntegerType` vs `StringType` difference would raise a schema mismatch. `ignoreColumnType` must be explicitly set to `True` to obtain the behavior described here.
- D. Column names are compared by default, not ignored; `ignoreColumnName` must be explicitly set to `True` to switch to purely positional, type-only comparison. Describing this as the default reverses the actual behavior.
4.Assert on schemas, not just rows
A row comparison only helps if you know which types you expected. DataFrame.schema returns the schema as a StructType: a list of StructFields, each with a name, a type and a nullable flag. Schema tests compare against that object.
The sources for this lesson have no reference page for pyspark.testing.assertSchemaEqual, so this lesson does not describe its parameters. Below is the schema material the sources do document, and it is what schema assertions catch in practice.
df = spark.createDataFrame(
[(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])
df.schema
# StructType([StructField('age', LongType(), True),
# StructField('name', StringType(), True)])df = spark.createDataFrame(
[(14, "Tom"), (23, "Alice"), (16, "Bob")],
"age INT, name STRING")
df.schema
# StructType([StructField('age', IntegerType(), True),
# StructField('name', StringType(), True)])The values are the same, but the schemas differ. This is why the way you build expected data matters. If you pass only a list of column names, types are inferred from the data. If you pass a DataType or a DDL string, you fix the types yourself. The nullable flag is the third part of each field. It is a real schema difference, but assertDataFrameEqual ignores it by default (ignoreNullable=True).
Some tests only need the column names. The Databricks pipeline example below compares the set of field names and nothing else:
# Test 2: Schema validation
def test_users_schema(test_spark):
mock_users(test_spark)
status = test_pipeline.run(test_spark, set(["catalog.schema.users"]))
assert status.is_success
result = test_spark.table("catalog.schema.users")
expected_fields = {"user_id", "email", "name", "user_type"}
actual_fields = set(f.name for f in result.schema.fields)
assert expected_fields == actual_fieldsCheckpoint 5 of 5· Check yourself
A colleague changes user_id in the users table from IntegerType to LongType. Which of these does Test 2 above still pass?
Test 2 builds a set of f.name values, so types and nullability are not part of the comparison. To catch a type change, the test has to compare the types as well as the names.
“actual_fields = set(f.name for f in result.schema.fields)”Source: docs.databricks.com
Sources6
Exam traps
Each one states something that sounds right. Open it to see what is actually true.
1.assertDataFrameEqual fails when the actual rows come back in a different order from the expected rows.Why is that wrong?
checkRowOrder defaults to False, so row order is not compared unless you change that parameter.
2.df1.sameSemantics(df2) is a reliable way to assert that two outputs contain the same data.Why is that wrong?
sameSemantics compares logical query plans. It can return False even when both DataFrames return the same results, and it is a developer API, not a testing helper.
3.An expected DataFrame built from Python ints with only column names will have IntegerType columns.Why is that wrong?
With only column names, types are inferred, and Python ints are inferred as LongType. You get IntegerType only if you declare it, for example with "age INT, name STRING".
Covered in Assert on schemas, not just rows
Sources
Every claim above is drawn from one of these pages, quoted as it was written on the date shown.
- 1.https://docs.databricks.com/aws/en/notebooks/testingOfficial docs
“a function should return a single predictable outcome and be of a single data type.”
↩︎ Design functions a test can pin down“you should run unit tests against non-production data.”
↩︎ Design functions a test can pin down“To return the number of rows that exist, the function should return a non-negative, whole number.”
↩︎ Checkpoint - 2.
“Concise syntax for chaining custom transformations.”
↩︎ Chain testable steps with DataFrame.transform“a function that takes and returns a DataFrame.”
↩︎ Key concept“Keyword arguments to pass to func.”
↩︎ Prediction - 3.https://docs.databricks.com/aws/en/ldp/unit-testingOfficial docs
“Full DataFrame comparison with assertDataFrameEqual”
↩︎ Compare whole results with assertDataFrameEqual“actual_fields = set(f.name for f in result.schema.fields)”
↩︎ Checkpoint - 4.https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.testing.assertDataFrameEqual.htmlSecondary source
“assertDataFrameEqual(actual, expected, checkRowOrder=False, rtol=1e-05, atol=1e-08, ignoreNullable=True”
↩︎ Compare whole results with assertDataFrameEqual“assertDataFrameEqual(actual, expected, checkRowOrder=False, rtol=1e-05, atol=1e-08”
↩︎ Exam trap 1“ignoreNullable=True, ignoreColumnOrder=False, ignoreColumnName=False, ignoreColumnType=False”
↩︎ Checkpoint - 5.https://docs.databricks.com/aws/en/pyspark/reference/classes/sparksession/createDataFrameOfficial docs
“When a list of column names is provided, the type of each column is inferred from data.”
↩︎ Compare whole results with assertDataFrameEqual - 6.
“Returns the schema of this DataFrame as a StructType.”
↩︎ Assert on schemas, not just rows“StructType([StructField('age', LongType(), True),”
↩︎ Exam trap 3
Also cited
“can still return False on the DataFrames that return the same results”
↩︎ Exam trap 2