Simple Example - Single Method Decorated#

This article shows how to get your hands dirty end-to-end with dltflow. In this sample, we’ll be going through the following steps:

Code#

Base Pipeline Code#

All pipelines leverage a base_pipeline.py file for the example that provides a common PipelineBase class that each pipeline inherits from. This is provided to help standardize how things like spark, logging, and configuration are initialized in pipelines. This pattern largely follows dbx’s documentation and task templating guide.

import sys
import pathlib
from logging import Logger
from typing import Any, Dict
from argparse import ArgumentParser

import yaml
from pyspark.sql import SparkSession


class PipelineBase:
    def __init__(self, spark: SparkSession, init_conf: dict = None):
        self.spark = self._get_spark(spark)
        self.logger = self._prepare_logger()
        self.conf = self._provide_config() if not init_conf else init_conf

    @staticmethod
    def _get_spark(spark: SparkSession):
        if not spark:
            spark = builder = (
                SparkSession.builder.master("local[1]")
                .appName("dltflow-examples")
                .getOrCreate()
            )
        return spark

    @staticmethod
    def _get_conf_file():
        """Uses the arg parser to extract the config location from cli."""
        p = ArgumentParser()
        p.add_argument("--conf-file", required=False, type=str)
        namespace = p.parse_known_args(sys.argv[1:])[0]
        return namespace.conf_file

    @staticmethod
    def _read_config(conf_file) -> dict[str, Any]:
        config = yaml.safe_load(pathlib.Path(conf_file).read_text())
        return config

    def _provide_config(self):
        """Orchestrates getting configuration."""
        self.logger.info("Reading configuration from --conf-file job option")
        conf_file = self._get_conf_file()
        if not conf_file:
            self.logger.info(
                "No conf file was provided, setting configuration to empty dict."
                "Please override configuration in subclass init method"
            )
            return {}
        else:
            self.logger.info(
                f"Conf file was provided, reading configuration from {conf_file}"
            )
            return self._read_config(conf_file)

    def _prepare_logger(self) -> Logger:  # pragma: no cover
        """Sets up the logger and ensures our job uses the log4j provided by spark."""
        log4j_logger = self.spark._jvm.org.apache.log4j  # noqa
        return log4j_logger.LogManager.getLogger(self.__class__.__name__)

Now that we understand that what base_pipeline.py does, lets get into our sample code.

Example Pipeline Code#

For this example, we’ll have a very simple pipeline. It will:

  • Import a DLTMetaMixin from dltflow.quality and will tell our sample pipeline to inherit from it.

  • Create sample data in memory.

  • We’ll transform that data by adding two integer columns together.

You should see that there are no direct calls to dlt. This is the beauty and intentional simplicity dltflow. It does not want to get in your way. Rather, it really wants you to focus on your transformation logic to help keep your code simple and easy to share with other team members.

from dltflow.quality import DLTMetaMixin
from pyspark.sql import DataFrame, SparkSession

from .base_pipeline import PipelineBase


class MySimplePipeline(PipelineBase, DLTMetaMixin):
    def __init__(self, spark: SparkSession, init_conf: dict = None):
        super().__init__(spark=spark, init_conf=init_conf)

    def read(self) -> DataFrame:
        return self.spark.createDataFrame([(1, 2), (5, 8), (1, 5)], ["col1", "col2"])

    def transform(self, df: DataFrame) -> DataFrame:
        return df.withColumn("result", df["col1"] + df["col2"])

    def write(self, df):
        df.write.saveAsTable("my_table")

    def orchestrate(self):
        df = self.read()
        df = self.transform(df)
        return df

Configuration#

Now that we have our example code, we need to write our configuration to tell the DLTMetaMixin how wrap our codebase.

Under the hood, dltflow uses pydantic to create validation for configuration. When working with dltflow, it requires your configuration to adhere to a specific structure. Namely, file should have the following sections:

  • reader: This is helpful for telling your pipeline where to read data from.

  • writer: Used to define where your data is written to after being processed.

  • dlt: Defines how dlt will be used in the project. We use this to dynamically wrap your code with dlt commands.

With this brief overview out of the way, lets review our configuration for this sample.

reader:
  ...
writer:
  adl_processed_path: "a/fake/path"
  schema: fake_schema
  table_name: table_name
dlt:
  func_name: "orchestrate"
  kind: table
  expectation_action: "allow"
  expectations:
    - name: "check_addition"
      constraint: "result < 10"

The dlt section has the following keys, though this configuration can also be a list of dlt configs.

  • func_name: The name of the function/method we want dlt to decorate.

  • kind: Tells dlt if this query should be materialized as a table or view

  • expectation_action: Tells dlt how to handle the expectations. drop, fail, and allow are all supported.

  • expectations: These are a list of constraints we want to apply to our data.

Workflow Spec#

Now that we’ve gone through the code and configuration, we need to start defining the workflow that we want to deploy to Databricks so that our pipeline can be registered as a DLT Pipeline. This structure largely follows the Databricks Pipeline API with the addition of a tasks key. This key is used during deployment for transitioning your python module into a Notebook that can be deployed as a DLT Pipeline.

dev:
  workflows:
    - name: "dltflow-simple-pipeline_with_expectations"
      storage: "/mnt/igdatalake/experiment/dltflow-samples/dlt/simple"
      target: "dltflow-samples"
      development: "true"
      edition: "ADVANCED"
      continuous: "false"
      clusters:
        - label: "default"
          node_type_id: Standard_DS3_v2"
          autoscale:
            min_workers: 1
            max_workers: 2
            mode: "ENHANCED"
      pipeline_type: "WORKSPACE"
      data_sampling: false
      tasks:
        items:
          - python_file: "pipelines/simple_dlt_pipeline_with_expectations.py"
            parameters:
              - "--conf"
              - "conf/simple_dlt_pipeline_with_expectations.yml"
        dependencies:
          - whl: "/dbfs/private-site-packages/dltflow-0.0.1b0-py3-none-any.whl"
          - pypi:
              package: "pyspark"

Deployment#

We’re at the final step of this simple example. The last piece of the puzzle here is that we need to deploy our assets to a Databricks workspace. To do so, we’ll use the dltflow cli.

# bin/sh

dltflow deploy-py-dlt --deployment-file ../workflows/simple_dlt_pipeline_with_expectations.yml \
  --environment dev --as-individual