Batch Data Source in Machine Learning

import pandas as pd
import pyarrow.parquet as pq

# Read only specific columns and partitions from a partitioned Parquet dataset
# Dataset partitioned by date: s3://data-lake/events/date=2025-01-01/
dataset = pq.ParquetDataset(
    "s3://data-lake/events/",
    filters=[("date", ">=", "2025-01-01"), ("date", "<", "2025-02-01")],
    use_legacy_dataset=False,
)

# Read only the columns needed for feature engineering
df = dataset.read(
    columns=["user_id", "event_type", "timestamp", "amount", "category"],
).to_pandas()

print(f"Loaded {len(df):,} rows, {df.memory_usage(deep=True).sum() / 1e6:.1f} MB")

# Validate schema expectations
expected_columns = {"user_id", "event_type", "timestamp", "amount", "category"}
assert set(df.columns) == expected_columns, f"Schema mismatch: {set(df.columns) - expected_columns}"
assert df["amount"].dtype == "float64", "Amount column type changed!"

# Tag this extraction for reproducibility
extraction_metadata = {
    "source": "s3://data-lake/events/",
    "date_range": "2025-01-01 to 2025-01-31",
    "row_count": len(df),
    "columns": list(df.columns),
    "checksum": pd.util.hash_pandas_object(df).sum(),
}

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from datetime import datetime
import json

spark = SparkSession.builder \
    .appName("batch-data-source") \
    .config("spark.sql.sources.partitionOverwriteMode", "dynamic") \
    .config("spark.sql.parquet.mergeSchema", "true") \
    .getOrCreate()

# --- Incremental Loading with High-Water Mark ---
def get_high_water_mark(state_path: str) -> str:
    """Retrieve the last successful extraction timestamp."""
    try:
        state = spark.read.json(state_path).first()
        return state["high_water_mark"]
    except Exception:
        return "1970-01-01T00:00:00"  # First run: extract everything

def update_high_water_mark(state_path: str, new_mark: str):
    """Atomically update the high-water mark after successful extraction."""
    state_df = spark.createDataFrame(
        [{"high_water_mark": new_mark, "updated_at": datetime.utcnow().isoformat()}]
    )
    state_df.coalesce(1).write.mode("overwrite").json(state_path)

# Read the high-water mark
state_path = "s3://ml-pipeline/state/events_hwm/"
hwm = get_high_water_mark(state_path)

# Incremental extraction: read only new records
df = spark.read.parquet("s3://data-lake/events/") \
    .filter(F.col("updated_at") > hwm) \
    .select("user_id", "event_type", "timestamp", "amount", "category")

# Schema validation
expected_schema = ["user_id", "event_type", "timestamp", "amount", "category"]
assert df.columns == expected_schema, f"Schema drift detected: {df.columns}"

# Cache and compute metrics
df.cache()
row_count = df.count()
max_ts = df.agg(F.max("updated_at")).first()[0]

print(f"Extracted {row_count:,} incremental rows since {hwm}")

# Write to staging area in optimized format
df.write \
    .mode("append") \
    .partitionBy("event_type") \
    .parquet("s3://ml-pipeline/staging/events/")

# Update high-water mark only after successful write
if row_count > 0:
    update_high_water_mark(state_path, max_ts)

import pandas as pd
from sqlalchemy import create_engine, text
from contextlib import contextmanager
from typing import Iterator

# Connection pooling for production database reads
engine = create_engine(
    "postgresql://reader:[email protected]:5432/analytics",
    pool_size=5,
    max_overflow=10,
    pool_timeout=30,
    pool_recycle=3600,  # Recycle connections every hour
)

@contextmanager
def read_connection():
    """Context manager for safe connection handling."""
    conn = engine.connect()
    try:
        yield conn
    finally:
        conn.close()

def extract_in_chunks(
    query: str,
    chunk_size: int = 50_000,
    params: dict = None,
) -> Iterator[pd.DataFrame]:
    """Extract data in memory-efficient chunks.
    
    Critical: Always read from a READ REPLICA, never from
    the primary database serving production traffic.
    """
    with read_connection() as conn:
        for chunk in pd.read_sql(
            text(query),
            conn,
            chunksize=chunk_size,
            params=params,
        ):
            # Validate each chunk's schema
            yield chunk

# Usage: extract user features for ML training
query = """
    SELECT user_id, signup_date, total_orders, avg_order_value,
           last_active_at, city, is_premium
    FROM user_features
    WHERE last_active_at >= :start_date
      AND last_active_at < :end_date
    ORDER BY user_id
"""

chunks = []
for chunk in extract_in_chunks(
    query,
    chunk_size=100_000,
    params={"start_date": "2025-01-01", "end_date": "2025-02-01"},
):
    chunks.append(chunk)
    print(f"  Extracted chunk: {len(chunk):,} rows")

df = pd.concat(chunks, ignore_index=True)
print(f"Total extracted: {len(df):,} rows")

# Save to Parquet for downstream pipeline
df.to_parquet("staging/user_features_202501.parquet", index=False)

# dvc.yaml - Pipeline definition with versioned data
# stages:
#   extract:
#     cmd: python extract_training_data.py
#     deps:
#       - extract_training_data.py
#     params:
#       - extract.start_date
#       - extract.end_date
#       - extract.source_table
#     outs:
#       - data/training/raw/

# extract_training_data.py
import yaml
import pandas as pd
from pathlib import Path

# Load parameters from params.yaml (tracked by DVC)
with open("params.yaml") as f:
    params = yaml.safe_load(f)["extract"]

print(f"Extracting from {params['source_table']}")
print(f"Date range: {params['start_date']} to {params['end_date']}")

# Extract data (using any connector -- DB, S3, etc.)
df = pd.read_parquet(
    f"s3://data-lake/{params['source_table']}/",
    filters=[
        ("date", ">=", params["start_date"]),
        ("date", "<", params["end_date"]),
    ],
)

# Save locally -- DVC will track and version this output
output_path = Path("data/training/raw/")
output_path.mkdir(parents=True, exist_ok=True)
df.to_parquet(output_path / "training_data.parquet", index=False)

# Print metrics for DVC tracking
print(f"Rows extracted: {len(df):,}")
print(f"Columns: {list(df.columns)}")
print(f"Size: {(output_path / 'training_data.parquet').stat().st_size / 1e6:.1f} MB")

# --- CLI Usage ---
# dvc init
# dvc run -n extract -d extract_training_data.py \
#         -p extract.start_date,extract.end_date,extract.source_table \
#         -o data/training/raw/ \
#         python extract_training_data.py
#
# dvc push  # Push data to remote storage (S3, GCS, Azure)
# git add dvc.yaml dvc.lock params.yaml
# git commit -m "Extract training data for Jan 2025"
#
# To reproduce later:
# git checkout <commit-hash>
# dvc pull
# dvc repro

from pyspark.sql import SparkSession
from delta.tables import DeltaTable

spark = SparkSession.builder \
    .appName("delta-batch-read") \
    .config("spark.jars.packages", "io.delta:delta-spark_2.12:3.1.0") \
    .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") \
    .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog") \
    .getOrCreate()

# Read the latest version of a Delta table
df_latest = spark.read.format("delta").load("s3://data-lake/transactions/")

# Time travel: read the table as it existed at a specific version
df_v42 = spark.read.format("delta") \
    .option("versionAsOf", 42) \
    .load("s3://data-lake/transactions/")

# Time travel: read the table as of a specific timestamp
df_jan = spark.read.format("delta") \
    .option("timestampAsOf", "2025-01-15T00:00:00Z") \
    .load("s3://data-lake/transactions/")

# Schema evolution: the table may have added columns over time
# Delta Lake handles this transparently -- new columns appear as NULL
# for older records
print(f"Current schema: {df_latest.schema.fieldNames()}")
print(f"V42 schema:     {df_v42.schema.fieldNames()}")

# Read incremental changes between versions (Change Data Feed)
changes = spark.read.format("delta") \
    .option("readChangeFeed", "true") \
    .option("startingVersion", 40) \
    .option("endingVersion", 42) \
    .load("s3://data-lake/transactions/")

# Filter to only new inserts (skip updates and deletes)
new_records = changes.filter(changes["_change_type"] == "insert")
print(f"New records between v40 and v42: {new_records.count():,}")

# batch_source_config.yaml
source:
  type: parquet
  path: s3://data-lake/events/
  format_options:
    merge_schema: true
    predicate_pushdown: true

extraction:
  mode: incremental  # or 'full'
  incremental_key: updated_at
  state_store: s3://ml-pipeline/state/events/
  partition_filter:
    date: ">=2025-01-01"
    region: ["IN", "US", "EU"]

schema:
  validation: strict  # strict | warn | skip
  expected_columns:
    - name: user_id
      type: string
      nullable: false
    - name: amount
      type: float64
      nullable: true
    - name: event_type
      type: string
      nullable: false

output:
  format: parquet
  compression: snappy
  staging_path: s3://ml-pipeline/staging/events/
  partition_by: [event_type, date]

versioning:
  enabled: true
  tool: dvc  # or 'lakefs', 'delta'
  tag_template: "events-{date}-{run_id}"

retry:
  max_attempts: 3
  backoff_seconds: [10, 60, 300]
  idempotent: true