cookbooks/Cookbook_1_Validate_data_during_ingestion_happy_path.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Cookbook 1: Validate data during ingestion (happy path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "This cookbook showcases a sample GX data validation workflow characteristic of data ingestion at the start of the data pipeline. Data is loaded into a Pandas dataframe, cleaned, validated, and then ingested into a Postgres database table.\n",
    "\n",
    "This cookbook explores the validation workflow first in a notebook setting, then embedded within an Airflow pipeline. Airflow pipelines are also referred to as directed acyclic graphs, or DAGs.\n",
    "\n",
    "This cookbook features a \"happy path\" scenario in which data passes validation and generates a successful pipeline run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Imports and constant definition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "This tutorial features the `great_expectations` library.\n",
    "\n",
    "The `tutorial_code` module contains helper functions used within this notebook and the associated Airflow pipeline.\n",
    "\n",
    "The `airflow_dags` submodule is included so that you can inspect the code used in the related Airflow DAG directly from this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "import inspect\n",
    "\n",
    "import great_expectations as gx\n",
    "import great_expectations.expectations as gxe\n",
    "import pandas as pd\n",
    "\n",
    "import tutorial_code as tutorial\n",
    "import airflow_dags.cookbook1_ingest_customer_data as dag"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Load raw customer data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "In this tutorial, you will clean and validate a dataset containing synthesized customer data. The data is loaded from a CSV file into a Pandas DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = pathlib.Path(\"/cookbooks/data/raw\")\n",
    "\n",
    "df_customers_raw = pd.read_csv(DATA_DIR / \"customers.csv\", encoding=\"unicode_escape\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Loaded {df_customers_raw.shape[0]} customer rows into dataframe.\\n\")\n",
    "\n",
    "display(df_customers_raw.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Examine destination table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "The customer data will be loaded into a Postgres table, `customers`. Examine the schema of the destination table and compare to the initial schema and contents of the raw customer data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "tutorial.db.get_table_schema(table_name=\"customers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Prior to running the Airflow pipeline, the Postgres `customers` table contains no data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "tutorial.db.get_table_row_count(table_name=\"customers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Clean customer data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "To clean the customer data, you will use a pre-prepared function, `clean_customer_data`. The cleaning code is displayed below, and then invoked to clean the raw customer data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pycat inspect.getsource(tutorial.cookbook1.clean_customer_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_customers = tutorial.cookbook1.clean_customer_data(df_customers_raw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Loaded {df_customers_raw.shape[0]} cleaned customer rows.\\n\")\n",
    "\n",
    "display(df_customers.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## GX validation workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "You will validate the cleaned customer data using GX prior to loading it into a Postgres database table. First, this section will explain an example of a simple a GX data validation workflow. Then, you'll apply that knowledge to validate the customer data.\n",
    "\n",
    "This tutorial contains concise explanations of GX components and workflows. For more detail, visit the [Introduction to GX Core](https://docs.greatexpectations.io/docs/core/introduction/) in the GX docs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "### Validate data with a single Expectation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "All GX workflows start with the creation of a **Data Context**. A Data Context is the Python object that serves as an entrypoint for the GX Core Python library, and it also manages the settings and metadata for your GX workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "context = gx.get_context()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "Next, you create a **Data Source**, **Data Asset**, and **Batch Definition**. You then use the Batch Definition to generate a **Batch** of data to validate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "```{admonition} Adding GX components to the Data Context\n",
    "GX components are unique on name. Once a component is created with the Data Context, adding another component with the same name will cause an error. To enable repeated execution of cookbook cells that add GX workflow components, you will see the following pattern:\n",
    "\n",
    "    try:\n",
    "        Add a new component(s) to the context\n",
    "    except:\n",
    "        Get component(s) from the context by name\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Data Source, Data Asset, and Batch Definition.\n",
    "try:\n",
    "    data_source = context.data_sources.add_pandas(\"pandas\")\n",
    "    data_asset = data_source.add_dataframe_asset(name=\"customer data\")\n",
    "    batch_definition = data_asset.add_batch_definition_whole_dataframe(\n",
    "        \"batch definition\"\n",
    "    )\n",
    "\n",
    "except:\n",
    "    data_source = context.data_sources.get(\"pandas\")\n",
    "    data_asset = data_source.get_asset(name=\"customer data\")\n",
    "    batch_definition = data_asset.get_batch_definition(\"batch definition\")\n",
    "\n",
    "# Get the Batch from the Batch Definition.\n",
    "batch = batch_definition.get_batch(batch_parameters={\"dataframe\": df_customers})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "An **Expectation** is a simple, declarative, verifiable assertion about your data. You can validate a Batch of data using an Expectation. Available Expectations can be easily found and instantiated using the `gxe` alias defined in the cookbook imports.\n",
    "\n",
    "First, create an Expectation that expects the columns in the customer data to match the provided ordered list of column names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "expectation = gxe.ExpectTableColumnsToMatchOrderedList(\n",
    "    column_list=[\"customer_id\", \"name\", \"dob\", \"city\", \"state\", \"zip\", \"country\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29",
   "metadata": {},
   "source": [
    "Next, validate your Batch using the Expectation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_result = batch.validate(expectation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {},
   "source": [
    "GX returns an `ExpectationValidationResult` object that provides metadata about the result of the validation and that can be accessed like a dictionary. The `ExpectationValidationResult` provides a variety of fields, most critically, the `success` field that indicates whether or not the Expectation passed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Results type: {type(validation_result)}\\n\")\n",
    "\n",
    "display(validation_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33",
   "metadata": {},
   "source": [
    "### Validate data with an Expectation Suite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34",
   "metadata": {},
   "source": [
    "Batches of data can also be validated with an **Expectation Suite**, which is a collection of Expectations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35",
   "metadata": {},
   "source": [
    "First, add a new Expectation Suite to the Data Context. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Expectation Suite.\n",
    "try:\n",
    "    expectation_suite = context.suites.add(\n",
    "        gx.ExpectationSuite(name=\"customer expectations\")\n",
    "    )\n",
    "except:\n",
    "    expectation_suite = context.suites.delete(name=\"customer expectations\")\n",
    "    expectation_suite = context.suites.add(\n",
    "        gx.ExpectationSuite(name=\"customer expectations\")\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37",
   "metadata": {},
   "source": [
    "Next, add Expectations to the Expectation Suite. Below, you will see Expectations that describe the required format of the customer data added to the Expectation Suite."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "expectations = [\n",
    "    gxe.ExpectTableColumnsToMatchOrderedList(\n",
    "        column_list=[\"customer_id\", \"name\", \"dob\", \"city\", \"state\", \"zip\", \"country\"]\n",
    "    ),\n",
    "    gxe.ExpectColumnValuesToBeOfType(column=\"customer_id\", type_=\"int\"),\n",
    "    *[\n",
    "        gxe.ExpectColumnValuesToBeOfType(column=x, type_=\"str\")\n",
    "        for x in [\"name\", \"city\", \"state\", \"zip\"]\n",
    "    ],\n",
    "    gxe.ExpectColumnValuesToMatchRegex(column=\"dob\", regex=r\"^\\d{4}-\\d{2}-\\d{2}$\"),\n",
    "    gxe.ExpectColumnValuesToBeInSet(\n",
    "        column=\"country\", value_set=[\"AU\", \"CA\", \"DE\", \"FR\", \"GB\", \"IT\", \"NL\", \"US\"]\n",
    "    ),\n",
    "]\n",
    "\n",
    "for expectation in expectations:\n",
    "    expectation_suite.add_expectation(expectation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39",
   "metadata": {},
   "source": [
    "Lastly, validate the Batch using the Expectation Suite."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate Batch using Expectation Suite.\n",
    "validation_result = batch.validate(expectation_suite)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41",
   "metadata": {},
   "source": [
    "When validating a Batch using an Expectation Suite, GX returns an `ExpectationSuiteValidationResult` object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(validation_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {},
   "source": [
    "Like the `ExpectationValidationResult` object, the `ExpectationSuiteValidationResult` object provides metadata about the result of the validation, but contains results for each of the individual Expectations that were run during the validation.\n",
    "\n",
    "* The `success` field indicates whether or not the validation passed. All individual Expectations in the Expectation Suite must pass for `success` to be `True`.\n",
    "* The `results` field contains indiviual results for each Expectation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Validation passed: {validation_result['success']}\\n\")\n",
    "\n",
    "display(validation_result[\"results\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45",
   "metadata": {},
   "source": [
    "## Integrate GX validation in the Airflow DAG"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46",
   "metadata": {},
   "source": [
    "You will use the `success` metadata of the GX validation result object to control the actions of the `cookbook1_validate_and_ingest_to_postgres` Airflow pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47",
   "metadata": {},
   "source": [
    "### Inspect DAG code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48",
   "metadata": {},
   "source": [
    "Examine the DAG code below that defines the `cookbook1_validate_and_ingest_to_postgres` pipeline. The DAG code checks the results of the GX validation before data is written to Postgres. If validation succeeds, the data is written to Postgres, but if validation fails, the pipeline will raise an error and halt.\n",
    "\n",
    "```\n",
    "# Halt pipeline with error if validation fails.\n",
    "if not validation_result[\"success\"]:\n",
    "    raise Exception(\"GX data validation failed.\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pycat inspect.getsource(dag)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50",
   "metadata": {},
   "source": [
    "### View the Airflow pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51",
   "metadata": {},
   "source": [
    "To view the `cookbook1_validate_and_ingest_to_postgres` pipeline in the Airflow UI, log into the locally running Airflow instance.\n",
    "\n",
    "1. Open [http://localhost:8080/](http://localhost:8080/) in a browser window.\n",
    "2. Log in with these credentials:\n",
    "  * Username: `admin`\n",
    "  * Password: `gx`\n",
    "\n",
    "You will see the pipeline under **DAGs** on login."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52",
   "metadata": {},
   "source": [
    "![Log in to tutorial Airflow UI](static/images/cookbook1_log_in_to_airflow_ui.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53",
   "metadata": {},
   "source": [
    "### Trigger the Airflow pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54",
   "metadata": {},
   "source": [
    "You can trigger the DAG from this notebook, using the provided convenience function in the cell below, or you can trigger the DAG manually in the Airflow UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "metadata": {},
   "outputs": [],
   "source": [
    "dag_run_id, dag_run_state = tutorial.airflow.trigger_airflow_dag(\n",
    "    \"cookbook1_validate_and_ingest_to_postgres\"\n",
    ")\n",
    "print(f\"DAG run {dag_run_id} is {dag_run_state}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56",
   "metadata": {},
   "source": [
    "To trigger the `cookbook1_validate_and_ingest_to_postgres` DAG from the Airflow UI, click the **Trigger DAG** button (with a play icon) under Actions. This will queue the DAG and it will execute shortly. The successful run is indicated by the run count inside the green circle under Runs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57",
   "metadata": {},
   "source": [
    "![Trigger the Airflow DAG](static/images/cookbook1_trigger_dag.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58",
   "metadata": {},
   "source": [
    "The `cookbook1_validate_and_ingest_to_postgres` DAG can be rerun multiple times; you can experiment with running it from this notebook or from the Airflow UI. The pipeline insert ignores into the Postgres `customers` table, meaning that it will not attempt to insert a row with the same primary key as an existing row."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59",
   "metadata": {},
   "source": [
    "### View pipeline results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60",
   "metadata": {},
   "source": [
    "Once the pipeline has been run, the `customers` table is populated with the cleaned customer data, a total of 15,266 rows. You can view the updated table count and a sampling of rows."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61",
   "metadata": {},
   "source": [
    "```{admonition} Results after using *Run All Cells*\n",
    "Airflow typically takes a few seconds to execute the `cookbook1_validate_and_ingest_to_postgres` DAG. If you ran this notebook using the JupyterLab *Run All Cells* or *Restart Kernel and Run All Cells* feature, there might not have been sufficient time for the queued Airflow DAG to execute prior to running the table count cell below.\n",
    "\n",
    "If you see `0` for the `customers` table count below and no sample rows, try re-running the cells.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62",
   "metadata": {},
   "outputs": [],
   "source": [
    "tutorial.db.get_table_row_count(table_name=\"customers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.read_sql_query(\n",
    "    \"select * from customers limit 10\", con=tutorial.db.get_local_postgres_engine()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64",
   "metadata": {},
   "source": [
    "It can also be helpful to view the pipeline logs to investigate the details of a successful (or unsuccessful run). To examine these logs in the Airflow UI:\n",
    "1. On the DAGs screen, click on the run(s) of interest under Runs.\n",
    "2. Click the name of the individual run you want to examine. This will load the DAG execution details.\n",
    "3. Click the Graph tab, and then the `cookbook1_validate_and_ingest_to_postgres` task box on the visual rendering.\n",
    "4. Click the Logs tab to load the DAG logs.\n",
    "\n",
    "You can see in the screen capture below that the logs reflect the row insertion print statement that was included in the DAG code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65",
   "metadata": {},
   "source": [
    "![Check logs for successful pipeline run](static/images/cookbook1_check_pipeline_logs.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67",
   "metadata": {},
   "source": [
    "This cookbook has walked you through the process of validating data using GX and integrating the data validation workflow in an Airflow pipeline.\n",
    "\n",
    "Future cookbooks will explore additional scenarios in which pipeline validation fails, the pipeline is halted, and invalid data is automatically handled in the pipeline execution."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}