Data Quality and Validation Questions

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, lit, array, struct, expr
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DoubleType, TimestampType

spark = SparkSession.builder.appName("enforce-schema-job").getOrCreate()

# Example target schema (define per your contract)
target_schema = StructType([
    StructField("id", StringType(), nullable=False),
    StructField("event_time", TimestampType(), nullable=False),
    StructField("score", DoubleType(), nullable=True),
    StructField("age", IntegerType(), nullable=True),
    StructField("text", StringType(), nullable=True)
])

input_path = "s3://bucket/input/parquet/"
clean_output = "s3://bucket/output/clean/"
quarantine_output = "s3://bucket/output/quarantine/"
report_output = "s3://bucket/output/report/"

# 1) Read dataset (schema inference or read as is)
df = spark.read.parquet(input_path)

# 2) Fail-fast: check required columns present
required_cols = [f.name for f in target_schema.fields if not f.nullable]
missing = [c for c in required_cols if c not in df.columns]
if missing:
    raise RuntimeError(f"Missing required columns: {missing}")  # fail-fast

# 3) For each field, attempt cast and capture coercion failures
coercion_flags = []
coerced_cols = []
for field in target_schema.fields:
    name = field.name
    tgt = field.dataType
    # casted column
    cast_col = col(name).cast(tgt).alias(f"{name}__cast")
    coerced_cols.append(cast_col)
    # detect failure:
    # If original is null and target nullable -> not an error.
    # If original not null and casted is null -> coercion failure.
    flag = when(
        col(name).isNotNull() & col(f"{name}__cast").isNull(), lit(1)
    ).otherwise(lit(0)).alias(f"{name}__error")
    coercion_flags.append(flag)

# Build a select that includes original, casted cols, and error flags
select_exprs = []
for field in target_schema.fields:
    name = field.name
    select_exprs.append(col(name))
# append coerced and flags via expression (need to add cast expressions since casted columns are referenced)
for field in target_schema.fields:
    name = field.name
    select_exprs.append(col(name).cast(field.dataType).alias(f"{name}__cast"))
    select_exprs.append(
        when(col(name).isNotNull() & col(name).cast(field.dataType).isNull(), lit(1)).otherwise(lit(0)).alias(f"{name}__error")
    )

df_checked = df.select(*select_exprs)

# 4) Separate quarantined rows (any error flag == 1) and clean rows
error_flag_cols = [f"{f.name}__error" for f in target_schema.fields]
any_error_expr = " + ".join(error_flag_cols) + " > 0"
quarantined = df_checked.filter(expr(any_error_expr))
clean = df_checked.filter(~expr(any_error_expr))

# 5) Write quarantine with metadata: original row + per-column error flags
quarantined_with_meta = quarantined.withColumn("error_columns", array(*[col(c) for c in error_flag_cols]))
quarantined_with_meta.write.mode("append").parquet(quarantine_output)

# 6) Produce final clean DataFrame with casted columns renamed back to schema names
final_cols = [col(f"{f.name}__cast").alias(f.name) for f in target_schema.fields]
clean_final = clean.select(*final_cols)
clean_final.write.mode("overwrite").parquet(clean_output)

# 7) Summary report: counts per column of coercion failures
agg_exprs = []
for f in target_schema.fields:
    agg_exprs.append(expr(f"sum({f.name}__error) as {f.name}_error_count"))
report = df_checked.agg(*agg_exprs)
report.coalesce(1).write.mode("overwrite").json(report_output)

spark.stop()