Data pipeline

Data Pipeline Patterns

PySpark DataFrames

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType

spark = SparkSession.builder \
    .appName("FeaturePipeline") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

# Read with schema inference disabled — always define schema explicitly
schema = StructType([
    StructField("user_id", StringType(), nullable=False),
    StructField("amount", DoubleType(), nullable=True),
    StructField("ts", StringType(), nullable=False),
])

df = spark.read.schema(schema).parquet("s3://bucket/events/")

# Transformations
result = (df
    .filter(F.col("amount").isNotNull())
    .withColumn("date", F.to_date("ts"))
    .groupBy("user_id", "date")
    .agg(
        F.sum("amount").alias("daily_spend"),
        F.count("*").alias("tx_count"),
    )
    .withColumn("avg_tx", F.col("daily_spend") / F.col("tx_count"))
)

result.write.mode("overwrite").partitionBy("date").parquet("s3://bucket/features/")

Kafka Consumer (Python)

from confluent_kafka import Consumer, KafkaError
import json

conf = {
    "bootstrap.servers": "kafka:9092",
    "group.id": "feature-consumer",
    "auto.offset.reset": "earliest",
    "enable.auto.commit": False,
}

consumer = Consumer(conf)
consumer.subscribe(["user-events"])

try:
    while True:
        msg = consumer.poll(timeout=1.0)
        if msg is None:
            continue
        if msg.error():
            if msg.error().code() == KafkaError._PARTITION_EOF:
                continue
            raise KafkaError(msg.error())

        record = json.loads(msg.value())
        process_record(record)
        consumer.commit(asynchronous=False)
finally:
    consumer.close()

Airflow DAG

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
from datetime import datetime, timedelta

default_args = {
    "owner": "ml-team",
    "retries": 2,
    "retry_delay": timedelta(minutes=5),
    "email_on_failure": True,
}

with DAG(
    dag_id="feature_pipeline",
    default_args=default_args,
    schedule_interval="0 2 * * *",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["ml", "features"],
) as dag:

    extract = SparkSubmitOperator(
        task_id="extract_events",
        application="jobs/extract_events.py",
        conf={"spark.executor.memory": "4g"},
    )

    transform = SparkSubmitOperator(
        task_id="compute_features",
        application="jobs/compute_features.py",
        application_args=["--date", "{{ ds }}"],
    )

    validate = PythonOperator(
        task_id="validate_features",
        python_callable=validate_feature_counts,
        op_kwargs={"date": "{{ ds }}"},
    )

    extract >> transform >> validate

Prefect Flow

from prefect import flow, task
from prefect.tasks import task_input_hash
from datetime import timedelta

@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=1))
def extract(date: str) -> pd.DataFrame:
    return pd.read_parquet(f"s3://bucket/raw/{date}/")

@task(retries=3, retry_delay_seconds=30)
def transform(df: pd.DataFrame) -> pd.DataFrame:
    return df.dropna().assign(revenue=df["price"] * df["qty"])

@task
def load(df: pd.DataFrame, date: str) -> None:
    df.to_parquet(f"s3://bucket/features/{date}/", index=False)

@flow(name="feature-pipeline", log_prints=True)
def feature_pipeline(date: str = "2024-01-01"):
    raw = extract(date)
    features = transform(raw)
    load(features, date)
    print(f"Loaded {len(features)} rows for {date}")

Data Quality Checks

import great_expectations as gx

context = gx.get_context()
datasource = context.sources.add_pandas("my_source")
asset = datasource.add_dataframe_asset("features")
batch = asset.build_batch_request(dataframe=df)

suite = context.add_expectation_suite("feature_suite")
validator = context.get_validator(batch_request=batch, expectation_suite=suite)

validator.expect_column_values_to_not_be_null("user_id")
validator.expect_column_values_to_be_between("daily_spend", min_value=0)
validator.expect_column_pair_values_to_be_equal("tx_count", "expected_count")
validator.save_expectation_suite()

results = validator.validate()
if not results["success"]:
    raise ValueError(f"Data quality check failed: {results}")

Key Patterns

• Partition pruning: always filter on partition columns early
• Broadcast joins: use F.broadcast(small_df) for lookup tables < 100MB
• Checkpoint: call df.checkpoint() before wide transformations to cut lineage
• Schema evolution: use mergeSchema=True only when intentional
• Idempotency: write to staging first, then atomic move/swap