Two-aggregate pivot

In September 2024's T-SQL Tuesday, Deepthi Goguri wants to hear about a recent technical issue you resolved. I recently ran into an issue with a client where a Spark Structured Streaming query I was developing just refused to run, and I'll talk about that here – thanks for hosting, Deepthi! 😊

The client's requirement was to split incoming event data into time buckets, then pivot each bucket to aggregate different event types in different columns. I won't present real client data here, but you can imagine taking an incoming source of events like this:

and transforming it into a streaming output that looks like this:

For the purposes of this example, I can produce the incoming source stream using Spark's rate source:

from pyspark.sql.functions import col, when, lit, window, first
 
df_src = (
    spark.readStream
    .format("rate")
    .option("rowsPerSecond", 3)
    .load()
    .withColumn("text_value", col("value").cast("string"))
    .withColumn("label", when(col("value") % 3 == 0, lit("x")).otherwise(lit("y")))
)

This snippet of code produces a pivoted output like the screenshot above:

df_pivot = (
    df_src
    .withWatermark("timestamp", "10 seconds")
    .groupBy(window("timestamp", "10 seconds"))
    .pivot("label", ["x", "y"])
    .agg(first("value"))
)

I'm using a stream watermark to handle late arriving data – basically¹⁾ my watermark enables the stream to accept data arriving up to 10 seconds late …and that's where the problem shows up.

When I run this streaming query – in Azure Databricks I can do this simply with display(df_pivot) – I receive the error:

A solution is provided in the error message – set the configuration property spark.sql.streaming.statefulOperator.checkCorrectness.enabled to false (you can stop reading now if you like 😉) – but caveats this advice with “if you understand the possible risk of correctness issue…”.

I thought I'd better find out a bit more.

What Spark was trying to protect me from was a side-effect of its “global” watermark. The stream watermark – created by .withWatermark(“timestamp”, … – keeps track of the latest timestamp encountered by the stream, but a streaming query tracks only one watermark value for the whole query. This can cause problems if a query contains more than one aggregate operation – when the first aggregate moves the watermark on, that might cause subsequent aggregates to discard data that they haven't actually processed. This article contains a nice explanation of the phenomenon.

Thanks for saving me from falling in the well, Sparky! As Spark is open source, I can even search the error message and locate the code where this is enforced 😊:

def checkStreamingQueryGlobalWatermarkLimit(plan: LogicalPlan, outputMode: OutputMode): Unit = {
  val failWhenDetected = SQLConf.get.statefulOperatorCorrectnessCheckEnabled
  try {
    plan.foreach { subPlan =>
    ...

But wait… my query only has one aggregate operation, doesn't it? Being able to look at the source code told me one thing I didn't already know – that the check for multiple stateful operations takes place in the Query Analyzer's LogicalPlan.

I can use df_pivot.explain(True) to take a look at the analyzed logical query plan:

window: struct<start:timestamp,end:timestamp>, x: bigint, y: bigint
~Project [window#5251-T10000ms, __pivot_first(value) AS `first(value)`#5258[0] AS x#5259L, __pivo
+- ~Aggregate [window#5251-T10000ms], [window#5251-T10000ms, pivotfirst(label#5214, first(value)#
   +- ~Aggregate [window#5251-T10000ms, label#5214], [window#5251-T10000ms, label#5214, first(val
      +- ~Project [named_struct(start, knownnullable(precisetimestampconversion(((precisetimestam
         +- ~Filter isnotnull(timestamp#5206-T10000ms)
            +- ~EventTimeWatermark timestamp#5206: timestamp, 10 seconds
               +- ~Project [timestamp#5206, value#5207L, text_value#5210, CASE WHEN...

..and as the highlighted lines show, the Analyzer has implemented my single aggregate using two logical Aggregate operators. I can find this in the source too:

if (aggregates.forall(a => PivotFirst.supportsDataType(a.dataType))) {
  // Since evaluating |pivotValues| if statements for each input row can get slow this is an
  // alternate plan that instead uses two steps of aggregation.
  val namedAggExps: Seq[NamedExpression] = aggregates.map(a => Alias(a, a.sql)())
  val namedPivotCol = pivotColumn match {
  ...

As the first line of code indicates, this is only an option available to the Analyzer for supported types. Right now, the Analyzer won't choose the two-aggregate implementation to pivot string values. If you don't need an arithmetic aggregate – as here where I'm using first – then a (fairly unattractive) workaround might be to:

    ...
    .pivot("label", ["x", "y"])
    .agg(first("text_value"))
)

It looks like maybe it was Spark that pushed me into the well after all 😕, but at least I have some confidence now that it's safe to turn off the correctness check – I'm not actually performing two aggregations, it's the Analyzer's internal implementation of pivot. If I've understood this correctly, the Analyzer's correctness check is over-cautious here and it's safe to turn the check off:

spark.conf.set("spark.sql.streaming.statefulOperator.checkCorrectness.enabled", "false")

But what if I aggregate my pivoted results later in the stream? I'll fall right back in the well! If this feels like a bug to you, feel free to vote for it 😉.

When I ran into this issue it felt pretty intractable, even with Spark offering me a solution, because I wasn't confident it would be safe. The availability of Spark's source code online makes it both open and searchable, which allowed me to look ito the engine and come to a conclusion.