How would you implement and validate a weighted aggregation in PySpark?

Instruction: Compute a weighted average by group with supported Python DataFrame APIs. Define null, zero-weight, invalid-input, and partition-merge behavior, then explain when a grouped pandas UDF would be justified.

Context: Assesses choosing built-in distributed aggregates before custom Python UDFs, mergeable sufficient statistics, input contracts, and skew/memory trade-offs.

Updated

Example answer

First ask whether the aggregation can be expressed with built-ins. A weighted average can: retain (sum(value × weight), sum(weight)), merge those sums across partitions, and divide once at the end. Averaging partition averages is generally wrong. JVM Aggregator and encoder examples are not Python DataFrame APIs.

For this exercise, null pairs are ignored; nonfinite paired values or weights and negative weights are rejected; and groups with no positive total weight return SQL NULL. Finite negative values are allowed. Zero-weight rows contribute zero. The complete PySpark 3.5+/4.x example uses double precision:

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.master("local[2]").appName("weighted-average").getOrCreate()
df = spark.createDataFrame(
    [("a", 10.0, 1.0), ("a", 20.0, 3.0), ("a", None, 9.0),
     ("b", 5.0, 0.0), ("c", None, 2.0), ("d", -5.0, 2.0)],
    "group string, value double, weight double",
).repartition(2)

def weighted_average(df):
    value, weight = F.col("value"), F.col("weight")
    pair = value.isNotNull() & weight.isNotNull()
    invalid = pair & (
        F.isnan(value) | F.isnan(weight)
        | (F.abs(value) == float("inf")) | (F.abs(weight) == float("inf"))
        | (weight < 0)
    )
    if df.where(invalid).limit(1).count():
        raise ValueError("paired values must be finite; weights must be nonnegative")
    totals = df.groupBy("group").agg(
        F.sum(F.when(pair, value * weight)).alias("weighted_sum"),
        F.sum(F.when(pair, weight)).alias("weight_sum"),
    )
    return totals.select(
        "group",
        F.when(F.col("weight_sum") > 0,
               F.col("weighted_sum") / F.col("weight_sum")).alias("weighted_average"),
    )

result = {r["group"]: r["weighted_average"] for r in weighted_average(df).collect()}
assert result == {"a": 17.5, "b": None, "c": None, "d": -5.0}
spark.stop()

The validation action adds a scan; in a production pipeline, validate once at ingestion or cache deliberately if repeated computation warrants it. Double arithmetic can overflow for extreme magnitudes and differ slightly with reduction order; define numeric ranges or a suitable decimal precision when exact financial arithmetic is required.

For genuinely custom non-decomposable logic, a grouped aggregate pandas UDF is a supported alternative, but it loads each group in memory and lacks partial aggregation. Skew can exhaust worker memory. Prefer built-ins for this example and test repartitioned data, nulls, zero totals, empty input, and rejected weights.

Sources

GroupedData.agg; pandas_udf.

Related Questions