Data quality and Airflow

Checking the quality of your data is essential to getting actionable insights from your data pipelines. Airflow offers many ways to orchestrate data quality checks directly from your DAGs.

This guide covers:

• Best practices and essential knowledge surrounding the planning of a data quality approach.
• When to implement data quality checks.
• How to choose a data quality check tool.
• How data lineage is connected to data quality.
• An in-depth look at two commonly used data quality check tools: SQL check operators and Great Expectations.
• An example comparing implementations of data quality checks using each of these tools.

Assumed knowledge

To get the most out of this guide, you should have knowledge of:

• What Airflow is and when to use it. See Introduction to Apache Airflow.
• Airflow operators. See Operators 101.
• Airflow connections. See Managing your Connections in Apache Airflow.
• Relational databases. See Relational database on Wikipedia.

Data quality essentials

What is considered good quality data is determined by the needs of your organization. Defining quality criteria for a given collection of datasets is often a task requiring collaboration with all data professionals involved in each step of the pipeline.

The following is a typical data quality check process:

• Assess where and in what formats relevant data is stored, and where it moves along your pipeline. A data lineage tool can assist with this process.
• Gather the data quality criteria from data professionals using the data.
• Determine the quality checks to run on different components of the dataset.
• Determine where those quality checks should occur in your pipeline.
• Determine what happens when a check fails.
• Choose a data quality tool.
• Write the checks.
• Test the checks.

Types of data quality checks

Data quality checks can be run on different components of the dataset:

• Column: Common checks include confirming uniqueness of a key column, restricting the amount of NULL values that are allowed, defining logical bounds for numeric or date columns, and defining valid options for categorical variables.
• Row: Common checks include comparing the values in several columns for each row, and ensuring that there are no mutually exclusive values present in the same row.
• Table: Common checks include checking against a template schema to verify that it has not changed, and checking that a table has a minimum row count.
• Over several tables: Common checks include data parity checks across several tables. You can also create logical checks like making sure a customer marked as inactive has no purchases listed in any of the other tables.

It is also important to distinguish between the two types of control in data quality:

• Static control: Runs checks on Spark or pandas DataFrames, or on existing tables in a database. These checks can be performed anywhere in a typical ETL pipeline and might halt the pipeline in case of a failed check to prevent further processing of data with quality issues.
• Flow control: Monitors incoming data on a record-by-record basis, which sometimes prevents single records that fail certain data quality checks from entering the database. This guide does not cover flow controls.

Where to run data quality checks

Data quality checks can be run in different locations within a data pipeline or an Airflow environment. For example, data quality checks can be placed in the following locations within an ETL pipeline:

• Before the transform step (data_quality_check_1)
• After the transform step (data_quality_check_2)
• Branching off from a step (data_quality_check_3)

The following DAG graph shows typical locations for data quality checks:

It’s common to use data quality checks (post_check_action_1 and post_check_action_2) with Airflow callbacks to alert data professionals of data quality issues through channels like email and Slack. You can also create a downstream task that runs only when all data quality checks are successful, which can be useful for reporting purposes.

When implementing data quality checks, consider how a check success or failure should influence downstream dependencies. Trigger Rules are especially useful for managing operator dependencies. It often makes sense to test your data quality checks in a dedicated DAG before you incorporate them into your pipelines.

When to implement data quality checks

When implementing data quality checks, it is important to consider the tradeoffs between the upfront work of implementing checks versus the cost of downstream issues caused by bad quality data.

You might need to implement data quality checks in the following circumstances:

• Downstream models serve customer-facing results that could be compromised in the case of a data quality issue.
• There are data quality concerns where an issue might not be immediately obvious to the data professional using the data.

Data quality tools

There are multiple open source tools that can be used to check data quality from an Airflow DAG. While this guide lists the most commonly used tools here, it focuses on the two tools that also integrate with OpenLineage:

• SQL check operators: A group of operators in the Common SQL provider package that you can use to define data quality checks using Python dictionaries and SQL from within your DAGs.
• Great Expectations: An open source data validation framework where checks are defined in JSON. Airflow offers a provider package including the GreatExpectationsOperator for easy integration.

Other tools that can be used for data quality checks include:

• Soda: An open source data validation framework that uses YAML to define checks which can be run in Airflow using the BashOperator. Soda also offers the ability to write any custom checks using SQL.
• dbt test: A testing framework for dbt models. The easiest way to run dbt models including dbt test within Airflow is with the Cosmos package.

Choosing a tool

Which tool you choose is determined by the needs and preferences of your organization. Astronomer recommends using SQL check operators if you want to:

• Write checks without needing to set up software in addition to Airflow.
• Write checks as Python dictionaries and in SQL.
• Use any SQL statement that returns a single row of booleans as a data quality check.
• Implement many different downstream dependencies depending on the outcome of different checks.
• Have full observability of which checks failed from within Airflow task logs, including the full SQL statements of failed checks.

Astronomer recommends using a data validation framework such as Great Expectations or Soda in the following circumstances:

• You want to collect the results of your data quality checks in a central place.
• You prefer to write checks in JSON (Great Expectations) or YAML (Soda).
• Most or all of your checks can be implemented by the predefined checks in the solution of your choice.
• You want to abstract your data quality checks from the DAG code.

Currently only SQL check operators and the GreatExpectationsOperator offer data lineage extraction through OpenLineage.

SQL check operators

SQL check operators execute a SQL statement that results in a set of booleans. A result of True leads to the check passing and the task being labeled as successful. A result of False, or any error when the statement is executed, leads to a failure of the task. Before using any of the operators, you need to define the connection to your data storage from the Airflow UI or with an external secrets manager.

The SQL check operators work with any backend solution that accepts SQL queries and supports Airflow, and differ in what kind of data quality checks they can perform and how they are defined. All SQL check operators are part of the Common SQL provider.

• SQLColumnCheckOperator: Defines checks on columns of a table using a Python dictionary.
• SQLTableCheckOperator: Runs aggregated and non-aggregated statements involving several columns of a table.
• SQLCheckOperator: Runs a SQL statement that returns a single row of booleans.
• SQLIntervalCheckOperator: Runs checks against historical data.

Astronomer recommends using the SQLColumnCheckOperator and SQLTableCheckOperator over the older SQLValueCheckOperator and SQLThresholdCheckOperator.

The logs from SQL check operators can be found in the regular Airflow task logs.

Great Expectations

Great Expectations is an open source data validation framework that allows the user to define data quality checks in JSON. The checks, also known as Expectation Suites, can be run in a DAG using the GreatExpectationsOperator from the Great Expectations provider. All currently available Expectations can be viewed on the Great Expectations website and creation of Custom Expectations is possible.

The easiest way to use Great Expectations with Airflow is to initialize a Great Expectations project in a directory accessible to your Airflow environment and using the automatic creation of a Checkpoint and Datasource from an Airflow connection by the GreatExpectationsOperator. This basic usage of the GreatExpectationsOperator does not need in-depth Great Expectations knowledge and full customization is possible.

When using Great Expectations, Airflow task logs show the results of the suite at the test-level in a JSON format. To get a detailed report on the checks that were run and their results, you can view the HTML files located in the great_expectations/uncommitted/data_docs/local_site/validations directory in any browser. These data docs can be generated to other backends as well as the local file.

Data lineage and data quality

In complex data ecosystems, lineage can be a powerful addition to data quality checks, especially for investigating what data from which origins caused a check to fail.

The SQL check operators will emit lineage metadata. The GreatExpectationsOperator will automatically trigger the OpenLineage action if an OpenLineage environment is recognized. If you are working with open source tools, you can view the resulting lineage using Marquez.

The output from the SQLColumnCheckOperator contains each individual check and whether or not it succeeded:

For the GreatExpectationsOperator, OpenLineage receives whether or not the whole Expectation Suite succeeded or failed:

Example: Comparing SQL check operators and Great Expectations

This example shows the steps necessary to perform the same set of data quality checks with SQL check operators and with Great Expectations.

The checks performed for both tools are:

• MY_DATE_COL has only unique values.
• MY_DATE_COL has only values between 2017-01-01 and 2022-01-01.
• MY_TEXT_COL has no null values.
• MY_TEXT_COL has at least 10 distinct values.
• MY_NUM_COL has a maximum value between 90 and 110.
• example_table has at least 10 rows.
• The value in each row of MY_COL_1 plus the value of the same row in MY_COL_2 is equal to 100.
• The most common value in MY_COL_3 is either val1 or val4.

Example: SQL check operators

The example DAG includes the following tasks:

• TheSQLColumnCheckOperator performs checks on several columns in the target table.
• The SQLTableCheckOperators performs checks on the whole table, involving one or more columns.
• The SQLCheckOperator makes sure the most common value of a categorical variable is one of two options.

While this example shows all the checks being written within the Python file defining the DAG, it is possible to modularize commonly used checks and SQL statements in separate files. If you’re using the Astro CLI, you can add the files to the include directory.

1
from airflow.decorators import dag
2
from pendulum import datetime
3
from airflow.operators.empty import EmptyOperator
4
from airflow.providers.common.sql.operators.sql import (
5
    SQLColumnCheckOperator,
6
    SQLTableCheckOperator,
7
    SQLCheckOperator,
8
)
9
10
DB_CONN = "snowflake_default"
11
DB = "database_name"
12
SCHEMA = "schema_name"
13
TABLE = "example_table"
14
15
16
@dag(
17
    start_date=datetime(2023, 7, 1),
18
    schedule="@daily",
19
    catchup=False,
20
    default_args={"conn_id": DB_CONN},
21
)
22
def example_dag_sql_check_operators():
23
    start = EmptyOperator(task_id="start")
24
    end = EmptyOperator(task_id="end")
25
26
    # SQLColumnCheckOperator example: runs checks on 3 columns:
27
    #   - MY_DATE_COL is checked to only contain unique values ("unique_check")
28
    #   - MY_TEXT_COL is checked to contain no NULL values
29
    #     and at least 10 distinct values
30
    #   - MY_NUM_COL is checked to have a maximum value between 90 and 110
31
    column_checks = SQLColumnCheckOperator(
32
        task_id="column_checks",
33
        table=TABLE,
34
        column_mapping={
35
            "MY_DATE_COL": {
36
                "unique_check": {"equal_to": 0},
37
            },
38
            "MY_TEXT_COL": {
39
                "distinct_check": {"geq_to": 10},
40
                "null_check": {"equal_to": 0},
41
            },
42
            "MY_NUM_COL": {"max": {"equal_to": 100, "tolerance": 0.1}},
43
        },
44
    )
45
46
    # SQLTableCheckOperator example: This operator performs three checks:
47
    #   - a row count check, making sure the table has >= 10 rows
48
    #   - a columns comparison check to see that the value in MY_COL_1 plus
49
    #   the value in MY_COL_2 is 100
50
    #   - a date between check to see that the value in MY_DATE_COL is between
51
    #   2017-01-01 and 2022-01-01.
52
    table_checks = SQLTableCheckOperator(
53
        task_id="table_checks",
54
        table=TABLE,
55
        checks={
56
            "my_row_count_check": {"check_statement": "COUNT(*) >= 10"},
57
            "my_column_comparison_check": {
58
                "check_statement": "MY_COL_1 + MY_COL_2 = 100"
59
            },
60
            "date_between_check": {
61
                "check_statement": "MY_DATE_COL BETWEEN '2017-01-01' AND '2022-01-01'"
62
            },
63
        },
64
    )
65
66
    # SQLCheckOperator example: ensure the most common categorical value in MY_COL_3
67
    # is one of a list of 2 options
68
    check_most_common_val_in_list = SQLCheckOperator(
69
        task_id="check_most_common_val_in_list",
70
        sql="""
71
            SELECT CASE
72
                WHEN (
73
                    SELECT {{ params.col }}
74
                    FROM {{ params.db_to_query }}.{{ params.schema }}.\
75
                            {{ params.table }}
76
                    GROUP BY {{ params.col }}
77
                    ORDER BY COUNT(*) DESC
78
                    LIMIT 1
79
                ) IN {{ params.options_tuple }}
80
                THEN 1
81
                ELSE 0
82
            END AS result;
83
            """,
84
        params={
85
            "db_to_query": DB,
86
            "schema": SCHEMA,
87
            "table": TABLE,
88
            "col": "MY_COL_3",
89
            "options_tuple": "('val1', 'val4')",
90
        },
91
    )
92
93
    (
94
        start
95
        >> [
96
            column_checks,
97
            table_checks,
98
            check_most_common_val_in_list,
99
        ]
100
        >> end
101
    )
102
103
104
example_dag_sql_check_operators()

Example: Great Expectations

The following example runs the same data quality checks as the SQL check operators example against the same database. After setting up the Great Expectations instance with at least the data context, the checks can be defined in a JSON file to form an Expectation Suite.

1
{
2
    "data_asset_type": null,
3
    "expectation_suite_name": "my_suite",
4
    "expectations": [
5
        {
6
            "expectation_context": {
7
                "description": "MY_DATE_COL has values between 2017-01-01 - 2022-01-01"
8
            },
9
            "expectation_type": "expect_column_values_to_be_between",
10
            "ge_cloud_id": null,
11
            "kwargs": {
12
                "column": "MY_DATE_COL",
13
                "max_value": "2022-01-01",
14
                "min_value": "2017-01-01"
15
            },
16
            "meta": {}
17
        },
18
        {
19
            "expectation_context": {
20
                "description": "MY_DATE_COL's values are all unique"
21
            },
22
            "expectation_type": "expect_column_values_to_be_unique",
23
            "ge_cloud_id": null,
24
            "kwargs": {
25
                "column": "MY_DATE_COL"
26
            },
27
            "meta": {}
28
        },
29
        {
30
            "expectation_context": {
31
                "description": "MY_TEXT_COL has at least 10 distinct values"
32
            },
33
            "expectation_type": "expect_column_unique_value_count_to_be_between",
34
            "ge_cloud_id": null,
35
            "kwargs": {
36
                "column": "MY_TEXT_COL",
37
                "min_value": 10
38
            },
39
            "meta": {}
40
        },
41
        {
42
            "expectation_context": {
43
                "description": "MY_TEXT_COL has no NULL values"
44
            },
45
            "expectation_type": "expect_column_values_to_not_be_null",
46
            "ge_cloud_id": null,
47
            "kwargs": {
48
                "column": "MY_TEXT_COL"
49
            },
50
            "meta": {}
51
        },
52
        {
53
            "expectation_context": {
54
                "description": "MY_NUM_COL has a maximum val between 90 and 110"
55
            },
56
            "expectation_type": "expect_column_max_to_be_between",
57
            "ge_cloud_id": null,
58
            "kwargs": {
59
                "column": "MY_NUM_COL",
60
                "min_value": 90,
61
                "max_value": 110
62
            },
63
            "meta": {}
64
        },
65
        {
66
            "expectation_context": {
67
                "description": "the table has at least 10 rows"
68
            },
69
            "expectation_type": "expect_table_row_count_to_be_between",
70
            "ge_cloud_id": null,
71
            "kwargs": {
72
                "min_value": 10
73
            },
74
            "meta": {}
75
        },
76
        {
77
            "expectation_context": {
78
                "description": "for each row MY_COL_1 plus MY_COL_2 is = 100"
79
            },
80
            "expectation_type": "expect_multicolumn_sum_to_equal",
81
            "ge_cloud_id": null,
82
            "kwargs": {
83
                "column_list": [
84
                    "MY_COL_1",
85
                    "MY_COL_2"
86
                ],
87
                "sum_total": 100
88
            },
89
            "meta": {}
90
        },
91
        {
92
            "expectation_context": {
93
                "description": "the most common value in MY_COL_3 is in the set ['val1', 'val4']"
94
            },
95
            "expectation_type": "expect_column_most_common_value_to_be_in_set",
96
            "ge_cloud_id": null,
97
            "kwargs": {
98
                "column": "MY_COL_3",
99
                "value_set": [
100
                    "val1",
101
                    "val4"
102
                ]
103
            },
104
            "meta": {}
105
        }
106
    ],
107
    "ge_cloud_id": null,
108
    "meta": {
109
        "great_expectations_version": "0.15.14"
110
    }
111
}

The corresponding DAG code shows how all the Expectations are run within one task using the GreatExpectationsOperator. Only the root directory of the data context, the schema and the data asset name have to be provided. For a step-by-step example and more information on the parameters of this operator see the Orchestrate Great Expectations with Airflow tutorial.

1
from airflow.decorators import dag
2
from pendulum import datetime
3
from great_expectations_provider.operators.great_expectations import (
4
    GreatExpectationsOperator,
5
)
6
7
DB_CONN = "snowflake_default"
8
SCHEMA = "schema_name"
9
TABLE = "example_table"
10
11
12
@dag(
13
    start_date=datetime(2023, 7, 1),
14
    schedule="@daily",
15
    catchup=False,
16
)
17
def gx_example_dag():
18
    GreatExpectationsOperator(
19
        task_id="gx_test",
20
        data_context_root_dir="/usr/local/airflow/include/great_expectations",
21
        conn_id=DB_CONN,
22
        schema=SCHEMA,
23
        data_asset_name=TABLE,
24
        expectation_suite_name="my_expectation_suite",
25
        do_xcom_push=False,
26
    )
27
28
29
gx_example_dag()

Conclusion

The growth of tools designed to perform data quality checks reflect the importance of ensuring data quality in production workflows. Commitment to data quality requires in-depth planning and collaboration between data professionals. What kind of data quality solution your organization needs will depend on your unique use case, which you can explore using the steps outlined in this guide. Special consideration should be given to the type of data quality checks and their location in the data pipeline.

Integrating Airflow with a data lineage tool can further enhance your ability to trace the origin of data that did not pass the checks you established.

This guide highlights two types of data quality tools and their use cases:

• SQL check operators offer a way to define your checks directly from within the DAG, with no other tools necessary.
• If you run many checks on different databases, you may benefit from using a more complex testing solution like Great Expectations or Soda.

No matter which tool is used, data quality checks can be orchestrated from within an Airflow DAG, which makes it possible to trigger downstream actions depending on the outcome of your checks.