pip install pyspark-utils-library

pip install -e .

pip install pyspark-utils-library[fuzzy]

pip install pyspark-utils-library[dev]

from pyspark_utils.transformations import clean_column_names, add_metadata_columns
from pyspark_utils.data_quality import DQValidator
from pyspark_utils.scd import scd2_merge
from pyspark_utils.logging_utils import PipelineLogger

# Clean and enrich a raw DataFrame
df = spark.read.table("bronze.raw_orders")
df = clean_column_names(df)
df = add_metadata_columns(df, source="erp_system")

# Validate data quality
report = (
    DQValidator(df)
    .check_nulls(["order_id", "customer_id"])
    .check_unique(["order_id"])
    .check_range("amount", min_val=0.0)
    .validate()
)

# Log results
logger = PipelineLogger(pipeline_name="orders_silver")
logger.log_dq_report(report)

from pyspark_utils.transformations import clean_column_names, rename_columns_bulk

# Normalize all column names to lowercase snake_case
df = clean_column_names(df)
# "Order ID" → "order_id", "customer-name" → "customer_name"

# Bulk rename with a mapping dict
df = rename_columns_bulk(df, {"old_col": "new_col", "amt": "amount"})

from pyspark_utils.transformations import cast_columns, type_cast_dataframe

# Cast specific columns
df = cast_columns(df, {"amount": "double", "quantity": "integer"})

# Auto-cast all columns based on a target schema
df = type_cast_dataframe(df, target_schema)

from pyspark_utils.transformations import flatten_struct, explode_array

# Flatten nested structs into top-level columns (address.city → address_city)
df = flatten_struct(df)

# Explode an array column with position index
df = explode_array(df, array_column="line_items", position_alias="item_index")

from pyspark_utils.transformations import pivot_aggregate, unpivot

# Pivot: rows → columns
df_pivoted = pivot_aggregate(
    df,
    group_columns=["customer_id"],
    pivot_column="metric_name",
    agg_column="metric_value",
    agg_function="sum",
)

# Unpivot: columns → rows
df_unpivoted = unpivot(
    df,
    id_columns=["customer_id"],
    value_columns=["revenue", "cost", "profit"],
    var_col_name="metric",
    val_col_name="value",
)

from pyspark_utils.transformations import (
    add_metadata_columns,
    hash_columns,
    add_surrogate_key,
)

# Add ETL metadata: _load_timestamp, _source, _filename
df = add_metadata_columns(df, source="salesforce")

# Create a SHA-256 hash column from selected columns
df = hash_columns(df, columns=["name", "email"], output_column="row_hash")

# Add a monotonically increasing surrogate key
df = add_surrogate_key(df, key_column="sk_customer")

from pyspark_utils.transformations import (
    trim_strings,
    coalesce_columns,
    filter_nulls,
    standardize_timestamps,
)

# Trim whitespace from all string columns
df = trim_strings(df)

# Coalesce multiple columns into one (first non-null wins)
df = coalesce_columns(df, columns=["phone_mobile", "phone_work", "phone_home"], output="phone")

# Filter rows where any/all specified columns are null
df = filter_nulls(df, columns=["order_id", "amount"], how="any")

# Parse and standardize mixed timestamp formats to a consistent format
df = standardize_timestamps(df, columns=["created_at", "updated_at"])

from pyspark_utils.data_quality import DQValidator

report = (
    DQValidator(df)
    .check_nulls(["order_id", "customer_id"])
    .check_unique(["order_id"])
    .check_range("amount", min_val=0.0, max_val=1_000_000.0)
    .check_regex("email", r"^[\w.+-]+@[\w-]+\.[\w.]+$")
    .validate()
)

print(report)
# DQReport(total=5, passed=4, failed=1, warnings=0, ...)

report = (
    DQValidator(orders_df)
    .check_referential_integrity(
        column="customer_id",
        reference_df=customers_df,
        reference_column="id",
    )
    .validate()
)

from pyspark.sql.types import StructType, StructField, StringType, DoubleType

expected_schema = StructType([
    StructField("order_id", StringType(), False),
    StructField("amount", DoubleType(), True),
])

report = (
    DQValidator(df)
    .check_schema(expected_schema)
    .validate()
)

report = (
    DQValidator(df)
    .check_custom(
        name="positive_margin",
        rule_fn=lambda d: d.filter(F.col("margin") < 0).count() == 0,
        description="All margins must be non-negative",
    )
    .validate()
)

report = (
    DQValidator(df)
    .check_row_count(min_count=1000, max_count=10_000_000)
    .validate()
)

from pyspark_utils.data_quality import DQValidator

# Raises ValueError immediately on first failure
DQValidator(df).check_nulls(["order_id"]).validate_or_raise()

report = DQValidator(df).check_nulls(["order_id"]).validate()

# Structured access
print(report.passed)      # Number of passed checks
print(report.failed)      # Number of failed checks
print(report.is_healthy)  # True if zero failures

# Iterate individual results
for result in report.results:
    print(f"{result.rule_name}: {result.status} — {result.detail}")

# Serialize to JSON for logging
print(report.to_json())

from pyspark_utils.scd import scd2_merge

scd2_merge(
    spark=spark,
    source_df=new_customers,
    target_table="gold.dim_customer",
    key_columns=["customer_id"],
    tracked_columns=["name", "email", "address"],
)

from pyspark_utils.scd import detect_changes

changes = detect_changes(
    source_df=new_data,
    target_df=spark.read.table("gold.dim_customer"),
    key_columns=["customer_id"],
    tracked_columns=["name", "email"],
)

print(f"New records: {changes['inserts'].count()}")
print(f"Changed records: {changes['updates'].count()}")
print(f"Unchanged records: {changes['unchanged'].count()}")

from pyspark_utils.scd import create_history_table

create_history_table(
    spark=spark,
    source_df=initial_load,
    target_table="gold.dim_product",
    key_columns=["product_id"],
)

from pyspark_utils.scd import close_expired_records

close_expired_records(
    spark=spark,
    target_table="gold.dim_customer",
    key_column_values={"customer_id": ["CUST-001", "CUST-002"]},
)

from pyspark_utils.scd import scd1_overwrite

scd1_overwrite(
    spark=spark,
    source_df=updated_products,
    target_table="gold.dim_product",
    key_columns=["product_id"],
)

from pyspark_utils.schema_utils import compare_schemas

result = compare_schemas(old_schema, new_schema)
print(result.added)           # ["new_column"]
print(result.removed)         # ["deprecated_col"]
print(result.type_changed)    # {"amount": ("string", "double")}
print(result.is_compatible)   # False (breaking changes detected)

from pyspark_utils.schema_utils import evolve_schema

# Adds missing columns (as null) and reorders to match target
evolved_df = evolve_schema(df, target_schema, fill_new_with_null=True)

from pyspark_utils.schema_utils import schema_to_ddl, ddl_to_schema

# Convert StructType to DDL string (useful for CREATE TABLE)
ddl_string = schema_to_ddl(df.schema)
# "order_id STRING NOT NULL, amount DOUBLE, created_at TIMESTAMP"

# Parse DDL back to StructType
schema = ddl_to_schema("order_id STRING, amount DOUBLE")

from pyspark_utils.schema_utils import validate_schema_compatibility

is_ok = validate_schema_compatibility(
    source_schema=incoming_df.schema,
    target_schema=spark.read.table("silver.orders").schema,
)
# Returns True if source can safely write to target

from pyspark_utils.schema_utils import merge_schemas

# Union of all fields from multiple schemas
merged = merge_schemas([schema_a, schema_b, schema_c])

from pyspark_utils.logging_utils import PipelineLogger

logger = PipelineLogger(pipeline_name="orders_silver")
logger.info("Pipeline started", extra={"batch_date": "2026-01-15"})
logger.warning("Null rate above threshold", extra={"column": "email", "pct": 0.12})
logger.error("Schema mismatch detected")

logger.track_metric("row_count", 125_000)
logger.track_metric("null_pct_order_id", 0.001)
logger.track_metric("processing_time_seconds", 42.5)

# Retrieve tracked metrics
metrics = logger.get_metrics()
# {"row_count": 125000, "null_pct_order_id": 0.001, ...}

# Alert if a metric exceeds bounds
logger.alert_on_threshold("null_pct_order_id", max_val=0.05)
# No alert — 0.001 < 0.05

logger.track_metric("null_pct_email", 0.12)
logger.alert_on_threshold("null_pct_email", max_val=0.05)
# WARNING: null_pct_email (0.12) exceeds max threshold 0.05

# Persist all logs to a Delta table
logger = PipelineLogger(
    pipeline_name="orders_silver",
    log_to_delta_table="audit.pipeline_logs",
)
# Logs are buffered and flushed to Delta on close or explicit flush

logger.info("Processing batch")
logger.flush_to_delta()  # Explicit flush

with PipelineLogger(pipeline_name="orders_silver") as logger:
    logger.info("Starting")
    # ... pipeline logic ...
    logger.info("Done")
# Automatically flushes on exit

from pyspark_utils.dedup import window_dedup

deduped = window_dedup(
    df,
    key_columns=["customer_id"],
    order_columns=[("updated_at", "desc")],
    keep="first",  # Keep the most recent per customer
)

from pyspark_utils.dedup import hash_dedup

deduped = hash_dedup(
    df,
    columns=["name", "email", "phone"],  # Hash only these columns
)

from pyspark_utils.dedup import fuzzy_dedup

deduped = fuzzy_dedup(
    df,
    key_column="customer_name",
    threshold=2,  # Max edit distance to consider a match
    group_columns=["city"],  # Reduce comparison space
)

from pyspark_utils.dedup import dedup_report

report = dedup_report(df, key_columns=["customer_id"])
report.show()
# +-------------+-----+
# | customer_id | count|
# +-------------+-----+
# | CUST-001    |    3 |
# | CUST-042    |    2 |
# +-------------+-----+

from pyspark_utils.diff import df_diff

result = df_diff(
    left=df_expected,
    right=df_actual,
    key_columns=["order_id"],
)

result["added"].show()      # Rows in actual but not expected
result["removed"].show()    # Rows in expected but not actual
result["changed"].show()    # Rows with same key but different values
result["unchanged"].show()  # Identical rows

summary = result["summary"]
print(summary)
# DiffSummary(left=1000, right=1005, added=10, removed=5, changed=3, unchanged=992)

from pyspark_utils.diff import column_diff

changes = column_diff(
    left=df_before,
    right=df_after,
    key_columns=["customer_id"],
)
changes.show()
# +-------------+--------+----------+-----------+
# | customer_id | column | old_value| new_value |
# +-------------+--------+----------+-----------+
# | CUST-001    | email  | old@x.co | new@x.co  |
# | CUST-001    | phone  | 555-0001 | 555-0099  |
# +-------------+--------+----------+-----------+

from pyspark_utils.diff import schema_diff

result = schema_diff(df_v1, df_v2)
print(result.added)          # Columns added in v2
print(result.removed)        # Columns removed in v2
print(result.type_changed)   # Columns with type changes

from pyspark_utils.diff import row_count_comparison

result = row_count_comparison(
    left=bronze_df,
    right=silver_df,
    labels=("bronze", "silver"),
)
print(result)
# RowCountResult(bronze=100000, silver=99850, difference=-150, pct_change=-0.15)

from pyspark_utils.diff import value_distribution_diff

dist = value_distribution_diff(
    left=df_january,
    right=df_february,
    column="order_status",
)
dist.show()
# +-----------+----------+-----------+--------+
# | value     | left_pct | right_pct | change |
# +-----------+----------+-----------+--------+
# | completed |     0.85 |      0.82 |  -0.03 |
# | pending   |     0.10 |      0.14 |  +0.04 |
# | cancelled |     0.05 |      0.04 |  -0.01 |
# +-----------+----------+-----------+--------+

from pyspark_utils.transformations import clean_column_names, trim_strings, add_metadata_columns
from pyspark_utils.data_quality import DQValidator
from pyspark_utils.dedup import window_dedup
from pyspark_utils.logging_utils import PipelineLogger

with PipelineLogger(pipeline_name="orders_bronze_to_silver") as logger:
    # Read bronze
    df = spark.read.table("bronze.raw_orders")
    logger.track_metric("bronze_row_count", df.count())

    # Transform
    df = clean_column_names(df)
    df = trim_strings(df)
    df = add_metadata_columns(df, source="erp")

    # Deduplicate
    df = window_dedup(
        df,
        key_columns=["order_id"],
        order_columns=[("_load_timestamp", "desc")],
    )

    # Validate
    report = (
        DQValidator(df)
        .check_nulls(["order_id", "customer_id", "amount"])
        .check_unique(["order_id"])
        .check_range("amount", min_val=0.0)
        .validate()
    )
    logger.log_dq_report(report)
    logger.track_metric("silver_row_count", df.count())

    # Write silver
    df.write.format("delta").mode("overwrite").saveAsTable("silver.orders")
    logger.info("Silver table written successfully")

from pyspark_utils.scd import scd2_merge
from pyspark_utils.schema_utils import compare_schemas
from pyspark_utils.logging_utils import PipelineLogger

with PipelineLogger(pipeline_name="dim_customer_scd2") as logger:
    # Read silver
    silver_df = spark.read.table("silver.customers")

    # Check for schema drift
    target_schema = spark.read.table("gold.dim_customer").schema
    diff = compare_schemas(target_schema, silver_df.schema)
    if not diff.is_compatible:
        logger.warning("Schema drift detected", extra={"added": diff.added})

    # Merge with SCD2
    scd2_merge(
        spark=spark,
        source_df=silver_df,
        target_table="gold.dim_customer",
        key_columns=["customer_id"],
        tracked_columns=["name", "email", "address", "phone"],
    )
    logger.info("SCD2 merge completed")

from pyspark_utils.data_quality import DQValidator

# Use validate_or_raise to halt the pipeline on failure
DQValidator(df) \
    .check_nulls(["order_id"]) \
    .check_unique(["order_id"]) \
    .check_row_count(min_count=1000) \
    .check_range("amount", min_val=0.0) \
    .validate_or_raise()

# If we reach here, all checks passed
df.write.format("delta").mode("overwrite").saveAsTable("silver.orders")