You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ka...@apache.org on 2026/08/10 06:12:19 UTC
(spark) branch master updated: [SPARK-58635][SS] Support for streaming aggregation in Real-Time Mode (RTM)
This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 03a304104b9e [SPARK-58635][SS] Support for streaming aggregation in Real-Time Mode (RTM)
03a304104b9e is described below
commit 03a304104b9e0eb28f611ad2eb8f7d65da0c33ba
Author: Boyang Jerry Peng <je...@databricks.com>
AuthorDate: Mon Aug 10 15:12:04 2026 +0900
[SPARK-58635][SS] Support for streaming aggregation in Real-Time Mode (RTM)
### What changes were proposed in this pull request?
This PR adds support for **streaming aggregation in Real-Time Mode (RTM)**, building on the pipelined
shuffle and RTM stateful work from SPARK-58508.
RTM cannot use the ordinary streaming aggregation plan. The standard plan materializes a partial
aggregate before the shuffle and holds results until the batch ends, which is incompatible with RTM's
requirement to emit rows continuously within a long-running batch. This PR introduces a *streamline*
aggregation plan that reads, merges, and emits per row instead.
The plan is three operators:
1. **`ProjectAggregationBufferExec`** (Partial) - projects input rows into aggregation buffers.
2. **`StatefulStreamlineAggregateExec`** (PartialMerge) - the stateful core. For each row it reads the
current buffer from the state store, merges the new value, and emits the updated buffer
immediately. Buffers written during the batch are held in a bounded cache (Guava `maximumSize`)
and flushed to the state store at batch end.
3. **`ProjectAggregationBufferExec`** (Final) - projects the merged buffers to the output schema.
Supporting pieces:
- `GenericBufferAggregationIterator` - shared buffer-merge iteration used by the operators above.
- `EvictionIterator` - incremental state-store cleanup, so eviction cost is spread across the batch
rather than paid in one pass at batch end.
- `AggUtils.planStreamlineStreamingAggregation` - the planner entry point, selected by
`StatefulAggregationStrategy` when the plan is an RTM plan.
- `IncrementalExecution` wires the operators into the RTM path and admits the
`stateStoreSave` <-> `StatefulStreamlineAggregate` operator transition in the state-metadata check
(the two share the same state manager and format, so their on-disk state is mutually readable);
`RealTimeModeAllowlist` admits the new operators.
- Two new internal configs: `spark.sql.streaming.useStreamlineAggregator` (default `false`) and
`spark.sql.streaming.statefulOperator.incrementalCleanupFactor` (default `0`, meaning cleanup at
batch end). RTM uses the streamline operator regardless of the first config; it exists so the
operator can also be exercised under a microbatch trigger.
The change is inert outside RTM: the streamline plan is only selected for a plan containing a
`RealTimeStreamScanExec` leaf (or when the config above is set), so the ordinary microbatch
aggregation path is untouched.
### Why are the changes needed?
Before this PR, a streaming aggregation was rejected in RTM: it planned into `HashAggregateExec`,
which is not in `RealTimeModeAllowlist`, so the query failed at start with
`STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST`. Aggregation is one of the most common
streaming workloads, so this was a significant gap in RTM's coverage. SPARK-58508 enabled stateful
queries in RTM (deduplication); this extends that to aggregation.
### Does this PR introduce _any_ user-facing change?
Yes, additive. A streaming aggregation query can now run under `Trigger.RealTime(...)` where it
previously failed at query start. Behavior outside RTM is unchanged: the two new configs are
`internal()` and default off, and the streamline plan is only selected for an RTM plan.
### How was this patch tested?
New and updated tests:
- `StreamlineStreamingAggregationRealTimeSuite` - end-to-end RTM aggregation coverage, including
incremental cleanup.
- `StreamingAggregationSuite` - the shared suite now also exercises the streamline operator (via the
`useStreamlineAggregator` config), so the new path is covered by the same assertions as the
existing micro-batch path. Adds a test that a query keeps its state when switched between the
micro-batch and streamline operators in either direction, guarding the operator-transition
allowance above.
- `StreamRealTimeModeAllowlistSuite` - updated for the newly admitted operators; the obsolete
aggregation-rejection test is removed.
Results on this branch: `StreamingAggregationSuite` 55/55,
`StreamlineStreamingAggregationRealTimeSuite` + `StreamRealTimeModeAllowlistSuite` 8/8, and the
streamline cases in `RocksDBStateStoreStreamingAggregationSuite` 26/26.
Several defects were caught by these tests during development and fixed, each in its own commit:
counting every output row in `ProjectAggregationBufferExec`; grouping the final aggregate by
attributes rather than expressions; copying the grouping key before caching it; allowing the
micro-batch <-> streamline operator transition on restart; and failing the task (rather than silently
dropping the write and still committing) when a state-store write fails during the batch-end flush.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code
Closes #57812 from jerrypeng/stack/pipelined-shuffle-pr9-rtm-aggregation.
Authored-by: Boyang Jerry Peng <je...@databricks.com>
Signed-off-by: Jungtaek Lim <ka...@gmail.com>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 30 ++
.../spark/sql/execution/SparkStrategies.scala | 42 +-
.../spark/sql/execution/aggregate/AggUtils.scala | 90 ++++
.../GenericBufferAggregationIterator.scala | 90 ++++
.../stateful/ProjectAggregationBufferExec.scala | 121 +++++
.../stateful/StatefulStreamlineAggregateExec.scala | 494 +++++++++++++++++++++
.../StreamingAggregationStateManager.scala | 23 +-
.../operators/stateful/statefulOperators.scala | 1 +
.../streaming/runtime/IncrementalExecution.scala | 39 +-
.../streaming/runtime/RealTimeModeAllowlist.scala | 7 +
.../streaming/state/EvictionIterator.scala | 105 +++++
.../StreamRealTimeModeAllowlistSuite.scala | 31 +-
.../sql/streaming/StreamingAggregationSuite.scala | 201 +++++++++
...reamlineStreamingAggregationRealTimeSuite.scala | 134 ++++++
14 files changed, 1373 insertions(+), 35 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 435aeb6e085f..5f80aaf9623a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -3783,6 +3783,36 @@ object SQLConf {
.checkValue(v => Set(1, 2).contains(v), "Valid versions are 1 and 2")
.createWithDefault(2)
+ val STREAMING_USE_STREAMLINE_AGGREGATOR =
+ buildConf("spark.sql.streaming.useStreamlineAggregator")
+ .internal()
+ .doc("Test/development only, not intended for production use. When true, plan a streaming " +
+ "aggregation with the streamline aggregation operator, which merges each input row " +
+ "against state and emits immediately, instead of the microbatch operators that only emit " +
+ "once the batch ends. Real-Time Mode queries use the streamline operator regardless of " +
+ "this config; this flag exists only so the operator can be exercised under an ordinary " +
+ "microbatch trigger in tests, and changes an aggregation's output timing when set.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
+ val STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR =
+ buildConf("spark.sql.streaming.statefulOperator.incrementalCleanupFactor")
+ .internal()
+ .doc("For a stateful operator that evicts by watermark, the number of eviction-eligible " +
+ "records to remove per input record processed, spreading eviction cost across the batch " +
+ "instead of paying it all at batch end. When 0, incremental cleanup is disabled and all " +
+ "eviction happens at batch end. When k, up to k eligible records are removed per input; " +
+ "any still-eligible records left over are removed at batch end. Only applies to modes " +
+ "that evict (e.g. Append/Update); has no effect in Complete mode, which never evicts. " +
+ "Currently read only by the streamline aggregation operator.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .longConf
+ .checkValue(_ >= 0, "The incremental cleanup factor must not be negative.")
+ .createWithDefault(0L)
+
val STREAMING_STOP_ACTIVE_RUN_ON_RESTART =
buildConf("spark.sql.streaming.stopActiveRunOnRestart")
.doc("Running multiple runs of the same streaming query concurrently is not supported. " +
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 996ecac61b72..577f09d9a685 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -39,7 +39,7 @@ import org.apache.spark.sql.execution.aggregate.AggUtils
import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec}
import org.apache.spark.sql.execution.command._
import org.apache.spark.sql.execution.datasources.{LogicalRelation, WriteFiles, WriteFilesExec}
-import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, StreamingDataSourceV2ScanRelation}
import org.apache.spark.sql.execution.exchange.{REBALANCE_PARTITIONS_BY_COL, REBALANCE_PARTITIONS_BY_NONE, REPARTITION_BY_COL, REPARTITION_BY_NUM, ShuffleExchangeExec}
import org.apache.spark.sql.execution.python._
import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec}
@@ -560,6 +560,20 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
* [[org.apache.spark.sql.execution.streaming.StreamExecution]]
*/
object StatefulAggregationStrategy extends Strategy {
+
+ /**
+ * Whether this plan reads a streaming source in Real-Time Mode, which is the case when the
+ * relation carries a real-time mode duration -- the same signal that decides whether to plan
+ * a [[org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec]] for it.
+ */
+ private def isRealTimeMode(plan: LogicalPlan): Boolean = {
+ plan.collectLeaves().exists {
+ case s: StreamingDataSourceV2ScanRelation =>
+ s.relation.realTimeModeDuration.isDefined
+ case _ => false
+ }
+ }
+
override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case _ if !plan.isStreaming => Nil
@@ -601,12 +615,26 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
case None =>
val stateVersion = conf.getConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION)
- AggUtils.planStreamingAggregation(
- normalizedGroupingExpressions,
- aggregateExpressions,
- rewrittenResultExpressions,
- stateVersion,
- planLater(child))
+ // A Real-Time Mode batch runs until its duration elapses rather than until its input
+ // is exhausted, so an aggregation that only emits once the batch ends would hold every
+ // result back for the whole batch. Plan the streamline operator instead, which merges
+ // each input row against state and emits immediately.
+ if (isRealTimeMode(child) ||
+ conf.getConf(SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR)) {
+ AggUtils.planStreamlineStreamingAggregation(
+ normalizedGroupingExpressions,
+ aggregateExpressions,
+ rewrittenResultExpressions,
+ stateVersion,
+ planLater(child))
+ } else {
+ AggUtils.planStreamingAggregation(
+ normalizedGroupingExpressions,
+ aggregateExpressions,
+ rewrittenResultExpressions,
+ stateVersion,
+ planLater(child))
+ }
}
case _ => Nil
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala
index 7fda56086380..16c34b404de6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala
@@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.catalyst.plans.logical.Aggregate
import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.streaming.{ProjectAggregationBufferExec, StatefulStreamlineAggregateExec}
import org.apache.spark.sql.execution.streaming.operators.stateful._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.util.Utils
@@ -425,6 +426,95 @@ object AggUtils {
finalAndCompleteAggregate :: Nil
}
+ /**
+ * Plans a "streamlined" streaming aggregation using the following progression:
+ * (Here the term "streamline" represents the loop of "read-process-output" for each input.)
+ *
+ * - Initialize Aggregation Buffer (Passthrough aggregation)
+ * - Shuffle
+ * - Streamlined Stateful Aggregation
+ * - For each input, do the following
+ * - Read the previous value for grouping key in state store
+ * - Merge the input and previous value (if any)
+ * - Store the new value to the state store
+ * - Complete (output the current result of the aggregation)
+ *
+ * The concept is to aggregate only between input and the state store, enabling an output to
+ * be available just from processing a single input. In update mode, each input will produce
+ * an output, which won't have any blocking operation in real time mode, leading to the
+ * lowest output latency.
+ */
+ def planStreamlineStreamingAggregation(
+ groupingExpressions: Seq[NamedExpression],
+ functionsWithoutDistinct: Seq[AggregateExpression],
+ resultExpressions: Seq[NamedExpression],
+ stateFormatVersion: Int,
+ child: SparkPlan): Seq[SparkPlan] = {
+ val groupingAttributes = groupingExpressions.map(_.toAttribute)
+
+ val initAggBuffer: SparkPlan = {
+ val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = Partial))
+ val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
+
+ ProjectAggregationBufferExec(
+ numShufflePartitions = None,
+ groupingExpressions = groupingExpressions,
+ aggregateExpressions = aggregateExpressions,
+ aggregateAttributes = aggregateAttributes,
+ resultExpressions = groupingAttributes ++
+ aggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes),
+ isFinalAggregate = false,
+ child = child)
+ }
+
+ val aggregate: SparkPlan = {
+ val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = PartialMerge))
+ val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
+
+ StatefulStreamlineAggregateExec(
+ requiredChildDistributionExpressions =
+ Some(groupingAttributes),
+ groupingExpressions = groupingAttributes,
+ aggregateExpressions = mayRemoveAggFilters(aggregateExpressions),
+ aggregateAttributes = aggregateAttributes,
+ initialInputBufferOffset = groupingAttributes.length,
+ resultExpressions = groupingAttributes ++
+ aggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes),
+ isFinalAggregate = false,
+ numShufflePartitions = None,
+ outputMode = None,
+ stateFormatVersion = stateFormatVersion,
+ child = initAggBuffer,
+ stateInfo = None,
+ eventTimeWatermarkForLateEvents = None,
+ eventTimeWatermarkForEviction = None)
+ }
+
+ val finalAndCompleteAggregate: SparkPlan = {
+ val finalAggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = Final))
+ // The attributes of the final aggregation buffer, which is presented as input to the result
+ // projection:
+ val finalAggregateAttributes = finalAggregateExpressions.map(_.resultAttribute)
+
+ ProjectAggregationBufferExec(
+ requiredChildDistributionExpressions = Some(groupingAttributes),
+ numShufflePartitions = None,
+ // The child here is the post-shuffle output of the stateful aggregate, whose grouping
+ // columns are already resolved attributes, so group by those rather than by the original
+ // expressions. This matches planStreamingAggregation's final stage and the stateful
+ // aggregate above, both of which group by the attributes.
+ groupingExpressions = groupingAttributes,
+ aggregateExpressions = finalAggregateExpressions,
+ aggregateAttributes = finalAggregateAttributes,
+ initialInputBufferOffset = groupingAttributes.length,
+ resultExpressions = resultExpressions,
+ isFinalAggregate = true,
+ child = aggregate)
+ }
+
+ finalAndCompleteAggregate :: Nil
+ }
+
/**
* Plans a streaming session aggregation using the following progression:
*
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/GenericBufferAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/GenericBufferAggregationIterator.scala
new file mode 100644
index 000000000000..839e52bde8fb
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/GenericBufferAggregationIterator.scala
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.streaming
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, MutableProjection, NamedExpression, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.execution.aggregate.AggregationIterator
+
+/**
+ * This base class extends [[AggregationIterator]] and produces a new aggregation buffer on demand.
+ * The instance of aggregation buffer is optimal for given aggregate functions.
+ *
+ * This base class is useful for cases where aggregation buffer needs to be initialized frequently,
+ * e.g. more than the cardinality of grouping keys.
+ */
+abstract class GenericBufferAggregationIterator(
+ partIndex: Int,
+ groupingExpressions: Seq[NamedExpression],
+ originalInputAttributes: Seq[Attribute],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ initialInputBufferOffset: Int,
+ resultExpressions: Seq[NamedExpression],
+ newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection)
+ extends AggregationIterator(
+ partIndex,
+ groupingExpressions,
+ originalInputAttributes,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ resultExpressions,
+ newMutableProjection) {
+
+ protected val useUnsafeBuffer = aggregateFunctions.flatMap(_.aggBufferAttributes)
+ .map(_.dataType).forall(UnsafeRow.isMutable)
+
+ /**
+ * Returns an aggregation buffer containing initial buffer values. Each call will produce the
+ * different buffer instance.
+ */
+ protected def newAggregationBuffer(): InternalRow = {
+ val buffer = initialAggregationBuffer.copy()
+ // if we are using a GenericInternalRow which
+ // is just a wrapper for an underlying data structured
+ // we need to re-initialize the buffer since
+ // copy does not actually create a new copy
+ // of the underlying data structure
+ if (!useUnsafeBuffer) {
+ initializeBuffer(buffer)
+ }
+ buffer
+ }
+
+ // An aggregation buffer containing initial buffer values. It is used to
+ // initialize other aggregation buffers.
+ private val initialAggregationBuffer: InternalRow = createNewAggregationBuffer()
+
+ private def createNewAggregationBuffer(): InternalRow = {
+ val bufferSchema = aggregateFunctions.flatMap(_.aggBufferAttributes)
+ val bufferRowSize: Int = bufferSchema.length
+ val genericMutableBuffer = new GenericInternalRow(bufferRowSize)
+
+ val buffer = if (useUnsafeBuffer) {
+ val unsafeProjection =
+ UnsafeProjection.create(bufferSchema.map(_.dataType))
+ val buf = unsafeProjection.apply(genericMutableBuffer)
+ initializeBuffer(buf)
+ buf
+ } else {
+ genericMutableBuffer
+ }
+ buffer
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/ProjectAggregationBufferExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/ProjectAggregationBufferExec.scala
new file mode 100644
index 000000000000..7046e4236aa0
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/ProjectAggregationBufferExec.scala
@@ -0,0 +1,121 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.streaming
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, MutableProjection, NamedExpression, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.execution.metric.SQLMetrics
+
+/**
+ * This class handles the part of aggregation functions in the input rows, based on the function's
+ * mode. This class intends to either initialize the aggregation buffer or complete the aggregate
+ * buffer and produce the result, so it is expected to be used for two aggregate modes:
+ * 1) partial merge 2) final. This class is pass-through and does not perform the actual
+ * aggregation.
+ */
+case class ProjectAggregationBufferExec(
+ requiredChildDistributionExpressions: Option[Seq[Expression]] = None,
+ numShufflePartitions: Option[Int],
+ groupingExpressions: Seq[NamedExpression] = Nil,
+ aggregateExpressions: Seq[AggregateExpression] = Nil,
+ aggregateAttributes: Seq[Attribute] = Nil,
+ initialInputBufferOffset: Int = 0,
+ resultExpressions: Seq[NamedExpression] = Nil,
+ isFinalAggregate: Boolean,
+ child: SparkPlan)
+ extends BaseAggregateExec {
+
+ override val isStreaming: Boolean = true
+
+ override lazy val metrics = Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ metrics // force lazy initialization at driver
+
+ val numOutputRows = longMetric("numOutputRows")
+
+ child.execute().mapPartitionsWithIndex { case (partIdx, iter) =>
+ val aggProcessor = new ProjectAggregationBufferProcessor(
+ partIdx,
+ groupingExpressions,
+ inputAttributes,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ resultExpressions,
+ (expressions, inputSchema) =>
+ MutableProjection.create(expressions, inputSchema))
+
+ iter.map { row =>
+ numOutputRows += 1
+ aggProcessor.process(row)
+ }
+ }
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ copy(child = newChild)
+
+ override def toSortAggregate: SortAggregateExec = {
+ throw new IllegalStateException("This class cannot be replaced with SortAggregate!")
+ }
+}
+
+/**
+ * This class is an implementation of GenericBufferAggregationIterator which only handles the
+ * aggregation buffer of input, depending on the aggregate mode. This class is pass-through
+ * and does not perform the actual aggregation.
+ */
+class ProjectAggregationBufferProcessor(
+ partIndex: Int,
+ groupingExpressions: Seq[NamedExpression],
+ originalInputAttributes: Seq[Attribute],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ initialInputBufferOffset: Int,
+ resultExpressions: Seq[NamedExpression],
+ newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection)
+ extends GenericBufferAggregationIterator(
+ partIndex,
+ groupingExpressions,
+ originalInputAttributes,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ resultExpressions,
+ newMutableProjection) {
+
+ def hasNext: Boolean =
+ throw new UnsupportedOperationException(
+ "hasNext is not supported in ProjectAggregationBufferProcessor")
+
+ def next(): UnsafeRow =
+ throw new UnsupportedOperationException(
+ "next is not supported in ProjectAggregationBufferProcessor")
+
+ def process(newInput: InternalRow): UnsafeRow = {
+ val groupingKey = groupingProjection.apply(newInput)
+ val buffer = newAggregationBuffer()
+ processRow(buffer, newInput)
+ generateOutput(groupingKey, buffer)
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala
new file mode 100644
index 000000000000..18d7d1a68b5f
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala
@@ -0,0 +1,494 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.streaming
+
+import java.util.concurrent.TimeUnit.NANOSECONDS
+
+import scala.util.control.NonFatal
+
+import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache, RemovalNotification}
+import org.apache.hadoop.conf.Configuration
+
+import org.apache.spark.TaskContext
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.WidenStatefulOpNullability
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, MutableProjection, NamedExpression, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.streaming.InternalOutputModes.{Append, Complete, Update}
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorCustomMetric, StatefulOperatorCustomSumMetric, StatefulOperatorStateInfo, StatefulOperatorsUtils, StateStoreWriter, StreamingAggregationStateManager, WatermarkSupport}
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateSchemaCompatibilityChecker, StateSchemaValidationResult, StateStore, StateStoreColFamilySchema, StateStoreOps}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.CompletionIterator
+
+/**
+ * The physical plan of streaming aggregation which "streamlines" the process of aggregation.
+ * (Here the term "streamline" represents the loop of "read-process-output" for each input.)
+ *
+ * Refer to the classdoc of [[StatefulStreamlineAggregationProcessor]] for more details.
+ *
+ * For producing result table as output, it follows the semantic of output mode.
+ *
+ * - Append mode: accumulated result for the grouping key will be produced once the watermark
+ * passes and there will be no further update against the grouping key.
+ * - Update mode: each input will produce intermediate accumulated result as an output.
+ * Note that this is different from streaming aggregation in microbatch mode
+ * ([[StateStoreSaveExec]]) which produces the final intermediate accumulated result for
+ * each grouping key only in this microbatch.
+ * - Complete mode: the entire result table will be produced per each microbatch.
+ */
+case class StatefulStreamlineAggregateExec(
+ requiredChildDistributionExpressions: Option[Seq[Expression]],
+ numShufflePartitions: Option[Int],
+ groupingExpressions: Seq[NamedExpression],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ initialInputBufferOffset: Int,
+ resultExpressions: Seq[NamedExpression],
+ isFinalAggregate: Boolean,
+ outputMode: Option[OutputMode] = None,
+ stateFormatVersion: Int,
+ child: SparkPlan,
+ stateInfo: Option[StatefulOperatorStateInfo] = None,
+ eventTimeWatermarkForLateEvents: Option[Long] = None,
+ eventTimeWatermarkForEviction: Option[Long] = None)
+ extends BaseAggregateExec with StateStoreWriter with WatermarkSupport {
+
+ override val isStreaming: Boolean = true
+
+ override def shortName: String =
+ StatefulOperatorsUtils.STATEFUL_STREAMLINE_AGGREGATE_EXEC_OP_NAME
+
+ override def keyExpressions: Seq[Attribute] = groupingExpressions.map(_.toAttribute)
+
+ // SPARK-57003 component (b): widen StatefulStreamlineAggregateExec output.
+ override def output: Seq[Attribute] =
+ WidenStatefulOpNullability.widenOutputForStatefulOp(
+ resultExpressions.map(_.toAttribute))
+
+ override def customStatefulOperatorMetrics: Seq[StatefulOperatorCustomMetric] = {
+ Seq(
+ StatefulOperatorCustomSumMetric(
+ "numRowsReadDuringEviction", "number of state rows read during state eviction"
+ ),
+ StatefulOperatorCustomSumMetric(
+ "numRowsIncrementallyRemoved", "number of state rows removed during incremental eviction"
+ )
+ )
+ }
+
+ private[sql] val stateManager = StreamingAggregationStateManager.createStateManager(
+ keyExpressions, child.output, stateFormatVersion)
+
+ // SPARK-57003 component (a): widen state schemas to nullable at construction. Both the
+ // schema-check site (`validateAndMaybeEvolveStateSchema`) and the runtime
+ // `mapPartitionsWithStateStore` site read these.
+ private val stateKeySchema: StructType =
+ WidenStatefulOpNullability.widenStateSchema(keyExpressions.toStructType)
+ private val stateValueSchema: StructType =
+ WidenStatefulOpNullability.widenStateSchema(stateManager.getStateValueSchema)
+
+ private val incrementalCleanupFactor = session.sessionState.conf.getConf(
+ SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR)
+
+ private def doIncrementalCleanup = incrementalCleanupFactor > 0
+
+ override def validateAndMaybeEvolveStateSchema(
+ hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
+ List[StateSchemaValidationResult] = {
+ val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME,
+ 0, stateKeySchema, 0, stateValueSchema))
+ List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo,
+ hadoopConf, newStateSchema, session.sessionState, stateSchemaVersion))
+ }
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ metrics // force lazy init at driver
+
+ val numOutputRows = longMetric("numOutputRows")
+ val numUpdatedStateRows = longMetric("numUpdatedStateRows")
+ val allUpdatesTimeMs = longMetric("allUpdatesTimeMs")
+ val numRowsReadDuringEviction = longMetric("numRowsReadDuringEviction")
+ val numRemovedStateRows = longMetric("numRemovedStateRows")
+ val allRemovalsTimeMs = longMetric("allRemovalsTimeMs")
+ val commitTimeMs = longMetric("commitTimeMs")
+ val numRowsIncrementallyRemoved = longMetric("numRowsIncrementallyRemoved")
+
+ assert(outputMode.nonEmpty,
+ "Incorrect planning in IncrementalExecution, outputMode has not been set")
+
+ child.execute().mapPartitionsWithStateStore(
+ getStateInfo,
+ stateKeySchema,
+ stateValueSchema,
+ NoPrefixKeyStateEncoderSpec(stateKeySchema),
+ session.sessionState,
+ Some(session.streams.stateStoreCoordinator)) { (store, iter) =>
+
+ // It's feasible to overload the method to provide a partition index, but now it's too
+ // many...
+ val partIdx = TaskContext.get().partitionId()
+
+ // Filter late date using watermark if specified
+ val baseIterator = watermarkPredicateForDataForLateEvents match {
+ case Some(predicate) => applyRemovingRowsOlderThanWatermark(iter, predicate)
+ case None => iter
+ }
+
+ val aggProcessor = new StatefulStreamlineAggregationProcessor(
+ partIdx,
+ groupingExpressions,
+ inputAttributes,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ resultExpressions,
+ (expressions, inputSchema) =>
+ MutableProjection.create(expressions, inputSchema),
+ stateManager,
+ store,
+ numUpdatedStateRows)
+
+ // Each output row from aggIter references the same underlying row.
+ // It is the caller's responsibility to ensure that each row is consumed before the next
+ // one is produced.
+ //
+ // flushDirtyWrites is aggIter's completion action, so it runs when aggIter is exhausted and
+ // BEFORE store.commit() in every mode below: Complete drains aggIter (line ~191) ahead of the
+ // commit iterator; Append drains it (~222) ahead of its own; Update chains it inside the
+ // outer CompletionIterator whose completion commits, and draining the outer drains aggIter
+ // first. This ordering is what lets flushDirtyWrites surface a state-write failure (see its
+ // rethrow) in time to fail the task before the batch commits -- a partial write can never be
+ // committed.
+ var tmpRow: UnsafeRow = null
+ val aggIter = CompletionIterator[UnsafeRow, Iterator[UnsafeRow]](
+ baseIterator.map { row =>
+ allUpdatesTimeMs += timeTakenMs {
+ tmpRow = aggProcessor.process(row)
+ }
+ tmpRow
+ },
+ aggProcessor.flushDirtyWrites() // Lazily evaluated
+ )
+
+ // Remaining logic performs the same thing with StateStoreSaveExec. It's mostly duplicated,
+ // with slight modification to cover the usage of aggregation iterator.
+
+ outputMode match {
+ // Update and output all rows in the StateStore.
+ case Some(Complete) =>
+ // consume iterator fully to process all inputs and save the result into state store.
+ aggIter.foreach(_ => ())
+
+ // SPARK-45582 - Ensure that store instance is not used after commit is called
+ // to invoke the iterator.
+ val rangeIter = stateManager.values(store)
+
+ CompletionIterator[UnsafeRow, Iterator[UnsafeRow]](
+ rangeIter.map { valueRow =>
+ numOutputRows += 1
+ valueRow
+ }, {
+ allRemovalsTimeMs += 0
+ commitTimeMs += timeTakenMs {
+ store.commit()
+ }
+ setStoreMetrics(store)
+ setOperatorMetrics()
+ }
+ )
+
+ // Update and output only rows being evicted from the StateStore
+ // Assumption: watermark predicates must be non-empty if append mode is allowed
+ case Some(Append) =>
+ assert(watermarkPredicateForDataForLateEvents.isDefined,
+ "Watermark needs to be defined for streaming aggregation query in append mode")
+
+ assert(watermarkPredicateForKeysForEviction.isDefined,
+ "Watermark needs to be defined for streaming aggregation query in append mode")
+
+ allUpdatesTimeMs += timeTakenMs {
+ // consume iterator fully to process all inputs and save the result into state store.
+ while (aggIter.hasNext) {
+ aggIter.next()
+ }
+ }
+
+ val removalStartTimeNs = System.nanoTime
+ val evictionIterator =
+ stateManager.evictionIterator(store, eventTimeWatermarkForEviction)
+
+ CompletionIterator[UnsafeRow, Iterator[UnsafeRow]](
+ evictionIterator.map(_.value), {
+ numRowsReadDuringEviction += evictionIterator.numRowsReadDuringEvictionSoFar
+ numRemovedStateRows += evictionIterator.numRowsRemovedSoFar
+ numOutputRows += evictionIterator.numRowsRemovedSoFar
+
+ // Note: Due to the iterator lazy exec, this metric also captures the time taken
+ // by the consumer operators in addition to the processing in this operator.
+ allRemovalsTimeMs += NANOSECONDS.toMillis(System.nanoTime - removalStartTimeNs)
+ commitTimeMs += timeTakenMs {
+ store.commit()
+ }
+ setStoreMetrics(store)
+ setOperatorMetrics()
+ })
+
+ // Update and output modified rows from the StateStore.
+ case Some(Update) =>
+ /**
+ * When doing incremental cleanup, we have to be careful what watermark to use. Because
+ * the late events timestamp is less than the eviction timestamp, within a batch it is
+ * possible for us to receive events whose timestamps is less than the eviction
+ * timestamp. Thus, it is possible to evict a record at timestamp t, such that
+ * t < evictionTimestamp, and, within the same batch, receive a record at timestamp t.
+ *
+ * Thus, when using incremental eviction, we have to make sure to clean up records
+ * up to the timestamp before which we will _never_ receive new records. This would be
+ * the event time watermark for late events.
+ */
+ val incrementalAwareEvictionWatermark = if (doIncrementalCleanup) {
+ eventTimeWatermarkForLateEvents
+ } else {
+ eventTimeWatermarkForEviction
+ }
+
+ // Only create it if we are doing incremental cleanup. If we instantiate it and are
+ // not doing incremental cleanup, then the iterator will not iterate through records
+ // inserted into the store during the batch.
+ val incrementalEvictionIter = if (doIncrementalCleanup) {
+ Some(stateManager.evictionIterator(store, incrementalAwareEvictionWatermark))
+ } else {
+ None
+ }
+
+ val updateIter = aggIter.map { row =>
+ incrementalEvictionIter.foreach { evictionIter =>
+ allRemovalsTimeMs += timeTakenMs {
+ var numRemovalsCurrRecord = 0
+ // NOTE: EvictionIterator removes (and counts) a row inside hasNext, not next(), so
+ // a bare hasNext already deletes and bumps the metrics. To stop at exactly
+ // incrementalCleanupFactor removals per input we must re-check the count before
+ // each hasNext rather than draining the iterator.
+ while (numRemovalsCurrRecord < incrementalCleanupFactor
+ && evictionIter.hasNext) {
+ // The removal happens inside of the iterator; in Update mode,
+ // we don't need the result.
+ evictionIter.next()
+ numRemovalsCurrRecord += 1
+ }
+ }
+ }
+ numOutputRows += 1
+ row
+ }
+
+ CompletionIterator[UnsafeRow, Iterator[UnsafeRow]](updateIter, {
+ // Anything removed so far must have been part of incremental eviction
+ incrementalEvictionIter.foreach { iter =>
+ numRowsIncrementallyRemoved += iter.numRowsRemovedSoFar
+ }
+
+ // If the incremental eviction iterator is defined, we'll finish eviction here
+ // if any records remain. If it's not, we'll construct an eviction iterator
+ // and also use it up here.
+ val evictionIter = incrementalEvictionIter.getOrElse(
+ stateManager.evictionIterator(store, eventTimeWatermarkForEviction))
+
+ allRemovalsTimeMs += timeTakenMs {
+ while (evictionIter.hasNext) {
+ evictionIter.next()
+ }
+ }
+
+ numRowsReadDuringEviction += evictionIter.numRowsReadDuringEvictionSoFar
+ numRemovedStateRows += evictionIter.numRowsRemovedSoFar
+
+ commitTimeMs += timeTakenMs {
+ store.commit()
+ }
+ setStoreMetrics(store)
+ setOperatorMetrics()
+ })
+
+ case _ => throw QueryExecutionErrors.unsupportedOutputModeForStreamingOperationError(
+ outputMode.get, "streaming aggregations")
+ }
+ }
+ }
+
+ override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = {
+ (outputMode.contains(Append) || outputMode.contains(Update)) &&
+ eventTimeWatermarkForEviction.isDefined &&
+ newInputWatermark > eventTimeWatermarkForEviction.get
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = {
+ copy(child = newChild)
+ }
+
+ // FIXME: How we can prevent this to be called?
+ override def toSortAggregate: SortAggregateExec = {
+ throw new IllegalStateException("This class cannot be replaced with SortAggregate!")
+ }
+}
+
+/**
+ * This class is an implementation of GenericBufferAggregationIterator which performs the
+ * aggregation against state store instead of maintaining aggregation hash table.
+ *
+ * For each input, do the following
+ * - Read the previous value for grouping key in state store
+ * - Merge the input and previous value (if any)
+ * - Produce the merged result
+ * - Store the new value to the dirty writes.
+ *
+ * This class maintains dirty writes which represent the cache of state store, performing both
+ * caching and deferred writes. Note that there is a size limit of dirty writes which should be
+ * considered carefully in both 1) memory usage and 2) latency spike on flushing writes.
+ *
+ * TODO: dirty writes can be implemented via LRU, which will optimize the ability of cache and
+ * also less spike on flushing writes (as evicted entry would be small portion of LRU).
+ */
+class StatefulStreamlineAggregationProcessor(
+ partIndex: Int,
+ groupingExpressions: Seq[NamedExpression],
+ originalInputAttributes: Seq[Attribute],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ initialInputBufferOffset: Int,
+ resultExpressions: Seq[NamedExpression],
+ newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection,
+ stateManager: StreamingAggregationStateManager,
+ stateStore: StateStore,
+ numUpdatedStateRows: SQLMetric)
+ extends GenericBufferAggregationIterator(
+ partIndex,
+ groupingExpressions,
+ originalInputAttributes,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ resultExpressions,
+ newMutableProjection) {
+
+ // The value for unsafe buffer is just an conservative arbitrary value - it should be probably
+ // safer to increase the value.
+ // The value for safe buffer is picked from the default value of fallback threshold of object
+ // hash aggregate, 128.
+ // That said, it's conservative, but should we still make this be configurable?
+ // NOTE: We need to consider the latency spike on flushing, so the number should not be too high
+ // even though there is more available memory to cache more entries.
+ protected val flushThresholdNumKeys: Int = if (useUnsafeBuffer) 1000 else 100
+
+ // Holds the first failure thrown while a removal listener writes an entry to the state store.
+ // Guava logs and SWALLOWS any exception a removal listener throws (see the
+ // CacheBuilder.removalListener scaladoc), so the put below cannot fail the task on its own --
+ // without this capture a failed write would be silently dropped and the batch would still commit,
+ // losing that key's update. flushDirtyWrites rethrows it so the task fails instead of committing
+ // partial state.
+ private var writeFailure: Option[Throwable] = None
+
+ // The writes which did not go into state store yet. We also leverage this dirty writes to the
+ // cache of state store.
+ private val dirtyWrites: LoadingCache[UnsafeRow, UnsafeRowReference] = CacheBuilder.newBuilder()
+ .maximumSize(flushThresholdNumKeys)
+ .removalListener((notification: RemovalNotification[UnsafeRow, UnsafeRowReference]) => {
+ // Skip once a write has already failed: the task is going to fail anyway, and further puts
+ // into a broken store are pointless.
+ if (writeFailure.isEmpty) {
+ try {
+ val value = notification.getValue
+ assert(value.getRow != null, "dirty writes should contain the valid row to update, " +
+ "but found null.")
+ stateManager.put(stateStore, value.getRow)
+ numUpdatedStateRows += 1
+ } catch {
+ case NonFatal(e) => writeFailure = Some(e)
+ }
+ }
+ })
+ .build(new CacheLoader[UnsafeRow, UnsafeRowReference] {
+ override def load(key: UnsafeRow): UnsafeRowReference = {
+ val newRef = new UnsafeRowReference
+ val existingValueInState = stateManager.get(stateStore, key)
+ // NOTE: newRef.getRow could be still null after this line
+ newRef.putRow(existingValueInState)
+ newRef
+ }
+ })
+
+ def hasNext: Boolean =
+ throw new UnsupportedOperationException(
+ "hasNext is not supported for StatefulStreamlineAggregationProcessor")
+
+ def next(): UnsafeRow =
+ throw new UnsupportedOperationException(
+ "next is not supported for StatefulStreamlineAggregationProcessor")
+
+ /**
+ * Flush the dirty writes to state store. This should be called at the end of each microbatch.
+ *
+ * Rethrows the first state-store write failure a removal listener captured, so that a failed put
+ * fails the task rather than committing state that is missing an update.
+ */
+ def flushDirtyWrites(): Unit = {
+ dirtyWrites.invalidateAll()
+ writeFailure.foreach { e =>
+ throw new IllegalStateException(
+ "Failed to write an aggregation update to the state store", e)
+ }
+ }
+
+ def process(newInput: InternalRow): UnsafeRow = {
+ // groupingProjection hands back the same UnsafeRow on every call, overwriting its bytes, so the
+ // key has to be copied before the cache can store it. Without the copy, one shared row becomes
+ // the key of every entry: each entry is then only findable through the hash captured when it
+ // was inserted, and the equality check that should discriminate between keys degenerates to a
+ // reference match against that one row.
+ val groupingKey = groupingProjection.apply(newInput).copy()
+ val buffer = newAggregationBuffer()
+
+ val existingValueRef = dirtyWrites.get(groupingKey)
+
+ if (existingValueRef.getRow != null) {
+ processRow(buffer, existingValueRef.getRow)
+ }
+
+ processRow(buffer, newInput)
+
+ val output = generateOutput(groupingKey, buffer)
+ existingValueRef.putRow(output.copy())
+ output
+ }
+}
+
+class UnsafeRowReference {
+ private var row: UnsafeRow = _
+
+ def getRow: UnsafeRow = row
+
+ def putRow(r: UnsafeRow): Unit = {
+ row = r
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala
index 8357053cdc46..63dab2c392c0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala
@@ -21,7 +21,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeRow}
import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateUnsafeProjection, GenerateUnsafeRowJoiner}
import org.apache.spark.sql.catalyst.types.DataTypeUtils
-import org.apache.spark.sql.execution.streaming.state.{NoopStatePartitionKeyExtractor, ReadStateStore, StateStore, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{EvictionIterator, NoopStatePartitionKeyExtractor, ReadStateStore, StateStore, UnsafeRowPair}
import org.apache.spark.sql.types.StructType
/**
@@ -61,6 +61,15 @@ sealed trait StreamingAggregationStateManager extends Serializable {
/** Return an iterator containing all the values in target state store. */
def values(store: ReadStateStore): Iterator[UnsafeRow]
+
+ /**
+ * Return an iterator over the rows of the target state store whose event time is older than
+ * `evictionTimestamp`, removing each row as it is returned. See [[EvictionIterator]].
+ */
+ def evictionIterator(
+ store: StateStore,
+ evictionTimestamp: Option[Long],
+ allowMultipleEventTimeColumns: Boolean = false): EvictionIterator
}
object StreamingAggregationStateManager extends Logging {
@@ -96,6 +105,18 @@ abstract class StreamingAggregationStateManagerBaseImpl(
// discard and don't convert values to avoid computation
store.iterator().map(_.key)
}
+
+ override def evictionIterator(
+ store: StateStore,
+ evictionTimestamp: Option[Long],
+ allowMultipleEventTimeColumns: Boolean = false): EvictionIterator = {
+ EvictionIterator(
+ store,
+ iterator(store),
+ keyExpressions,
+ allowMultipleEventTimeColumns,
+ evictionTimestamp)
+ }
}
/**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala
index 022fa3469eea..48d08b719492 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala
@@ -1677,6 +1677,7 @@ object StatefulOperatorsUtils {
)
val SYMMETRIC_HASH_JOIN_EXEC_OP_NAME = "symmetricHashJoin"
val STATE_STORE_SAVE_EXEC_OP_NAME = "stateStoreSave"
+ val STATEFUL_STREAMLINE_AGGREGATE_EXEC_OP_NAME = "StatefulStreamlineAggregate"
val DEDUPLICATE_EXEC_OP_NAME = "dedupe"
val DEDUPLICATE_WITHIN_WATERMARK_EXEC_OP_NAME = "dedupeWithinWatermark"
val SESSION_WINDOW_STATE_STORE_SAVE_EXEC_OP_NAME = "sessionWindowStateStoreSaveExec"
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
index f7b64af45a31..e46a4bf28690 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
@@ -40,7 +40,7 @@ import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec
import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader
import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec}
-import org.apache.spark.sql.execution.streaming.{StreamingErrors, StreamingQueryPlanTraverseHelper}
+import org.apache.spark.sql.execution.streaming.{ProjectAggregationBufferExec, StatefulStreamlineAggregateExec, StreamingErrors, StreamingQueryPlanTraverseHelper}
import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata, OffsetSeqMetadataBase}
import org.apache.spark.sql.execution.streaming.operators.stateful.{SessionWindowStateStoreRestoreExec, SessionWindowStateStoreSaveExec, StatefulOperator, StatefulOperatorStateInfo, StateStoreRestoreExec, StateStoreSaveExec, StateStoreWriter, StreamingDeduplicateExec, StreamingDeduplicateWithinWatermarkExec, StreamingGlobalLimitExec, StreamingLocalLimitExec, UpdateEventTimeColumnExec}
import org.apache.spark.sql.execution.streaming.operators.stateful.flatmapgroupswithstate.FlatMapGroupsWithStateExec
@@ -195,6 +195,9 @@ class IncrementalExecution(
case a: UpdatingSessionsExec if a.isStreaming =>
a.copy(numShufflePartitions = Some(numStateStores))
+
+ case a: ProjectAggregationBufferExec if a.isStreaming =>
+ a.copy(numShufflePartitions = Some(numStateStores))
}
}
@@ -317,6 +320,15 @@ class IncrementalExecution(
object StateOpIdRule extends SparkPlanPartialRule {
override val rule: PartialFunction[SparkPlan, SparkPlan] = {
+ case a: StatefulStreamlineAggregateExec =>
+ val aggStateInfo = nextStatefulOperationStateInfo()
+ a.copy(
+ numShufflePartitions = Some(aggStateInfo.numPartitions),
+ stateInfo = Some(aggStateInfo),
+ outputMode = Some(outputMode),
+ eventTimeWatermarkForLateEvents = None,
+ eventTimeWatermarkForEviction = None)
+
case StateStoreSaveExec(keys, None, None, None, None, stateFormatVersion,
UnaryExecNode(agg,
StateStoreRestoreExec(_, None, _, child))) =>
@@ -441,6 +453,12 @@ class IncrementalExecution(
}
override val rule: PartialFunction[SparkPlan, SparkPlan] = {
+ case a: StatefulStreamlineAggregateExec if a.stateInfo.isDefined =>
+ a.copy(
+ eventTimeWatermarkForLateEvents = inputWatermarkForLateEvents(a.stateInfo.get),
+ eventTimeWatermarkForEviction = inputWatermarkForEviction(a.stateInfo.get)
+ )
+
case s: StateStoreSaveExec if s.stateInfo.isDefined =>
s.copy(
eventTimeWatermarkForLateEvents = inputWatermarkForLateEvents(s.stateInfo.get),
@@ -585,6 +603,23 @@ class IncrementalExecution(
rulesToCompose.reduceLeft { (ruleA, ruleB) => ruleA orElse ruleB }
}
+ /**
+ * Returns true if a checkpoint whose metadata records operator `oldOpName` may be reopened by a
+ * plan whose operator is `newOpName`, without tripping the operator-mismatch guard.
+ *
+ * A streaming aggregation may move between the micro-batch operator ([[StateStoreSaveExec]],
+ * "stateStoreSave") and the streamline operator ([[StatefulStreamlineAggregateExec]]) in either
+ * direction: the two share [[StreamingAggregationStateManager]] and the same state format
+ * version, so the on-disk state written by one is readable by the other.
+ */
+ private def isOperatorMetadataConvertible(oldOpName: String, newOpName: String): Boolean = {
+ oldOpName == newOpName || ((oldOpName, newOpName) match {
+ case ("stateStoreSave", "StatefulStreamlineAggregate") => true
+ case ("StatefulStreamlineAggregate", "stateStoreSave") => true
+ case _ => false
+ })
+ }
+
private def checkOperatorValidWithMetadata(
planWithStateOpId: SparkPlan,
batchId: Long): Unit = {
@@ -640,7 +675,7 @@ class IncrementalExecution(
(opMapInMetadata.keySet ++ opMapInPhysicalPlan.keySet).foreach { opId =>
val opInMetadata = opMapInMetadata.getOrElse(opId, "not found")
val opInCurBatch = opMapInPhysicalPlan.getOrElse(opId, "not found")
- if (opInMetadata != opInCurBatch) {
+ if (!isOperatorMetadataConvertible(opInMetadata, opInCurBatch)) {
throw QueryExecutionErrors.statefulOperatorNotMatchInStateMetadataError(
opMapInMetadata,
opMapInPhysicalPlan)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
index 2cae2f7ee2ad..47cc286bb617 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
@@ -63,6 +63,13 @@ object RealTimeModeAllowlist extends Logging {
// exchange is a supported member of a pipelined group rather than a materialization barrier.
"org.apache.spark.sql.execution.exchange.ShuffleExchangeExec",
"org.apache.spark.sql.execution.joins.BroadcastHashJoinExec",
+ // Streaming aggregation. A Real-Time Mode batch does not end when its input is exhausted, so
+ // an aggregation is planned as the streamline operator, which merges each input row against
+ // state and emits immediately, wrapped by the two buffer-projection stages that initialize
+ // the aggregation buffer and produce the result columns (see
+ // AggUtils.planStreamlineStreamingAggregation).
+ "org.apache.spark.sql.execution.streaming.ProjectAggregationBufferExec",
+ "org.apache.spark.sql.execution.streaming.StatefulStreamlineAggregateExec",
// Streaming deduplication and the state-store access operators it plans into. These run in the
// pipelined-shuffle consumer stage, keyed by the same columns the shuffle repartitions on.
"org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreRestoreExec",
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala
new file mode 100644
index 000000000000..e8b71adeb746
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.streaming.state
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Predicate}
+import org.apache.spark.sql.execution.streaming.operators.stateful.WatermarkSupport
+
+/**
+ * An iterator over the evictable rows of a [[StateStore]], which removes each row it returns and
+ * reports how far it has progressed.
+ *
+ * Only the rows that are actually evicted are returned, so every row this iterator yields has been
+ * removed from the store. The removal happens in `hasNext` rather than `next` (see the note on
+ * `pending` below) precisely so that a caller which stops early still leaves the store consistent
+ * with what it observed. That lets a caller both emit the evicted rows -- streaming aggregation in
+ * append mode outputs a grouping key once the watermark passes it -- and count real removals rather
+ * than state rows scanned.
+ */
+trait EvictionIterator extends Iterator[UnsafeRowPair] {
+ /** Number of state rows examined so far, whether or not they were evicted. */
+ def numRowsReadDuringEvictionSoFar: Long
+
+ /** Number of state rows removed so far. */
+ def numRowsRemovedSoFar: Long
+}
+
+object EvictionIterator {
+
+ /**
+ * Returns an [[EvictionIterator]] over the rows of `store` whose event time is older than
+ * `evictionTimestamp`, removing each row as it is returned.
+ *
+ * The event time is read from the state store key, using the watermark metadata on
+ * `keyExpressions`. If those attributes carry no event time column, or `evictionTimestamp` is
+ * empty, nothing can be evicted and the iterator is empty.
+ *
+ * Note `evictionTimestamp` is not necessarily the current watermark: a caller doing incremental
+ * cleanup may pass an earlier timestamp, before which no further input can arrive.
+ */
+ def apply(
+ store: StateStore,
+ storeIterator: Iterator[UnsafeRowPair],
+ keyExpressions: Seq[Attribute],
+ allowMultipleEventTimeColumns: Boolean,
+ evictionTimestamp: Option[Long]): EvictionIterator = {
+
+ val evictionPredicate = WatermarkSupport.watermarkExpression(
+ WatermarkSupport.findEventTimeColumn(keyExpressions, allowMultipleEventTimeColumns),
+ evictionTimestamp).map { expr =>
+ Predicate.create(expr, keyExpressions)
+ }
+
+ new EvictionIterator {
+ private var rowsRead = 0L
+ private var rowsRemoved = 0L
+
+ override def numRowsReadDuringEvictionSoFar: Long = rowsRead
+ override def numRowsRemovedSoFar: Long = rowsRemoved
+
+ // The row hasNext has advanced to and already removed from the store, held so next() can
+ // return it. Removal happens in hasNext (not next()) so that a caller which stops iterating
+ // after hasNext -- without the matching next() -- still leaves the store consistent with the
+ // rows it was told are evicted; every row this iterator surfaces has already been removed.
+ private var pending: Option[UnsafeRowPair] = None
+
+ override def hasNext: Boolean = evictionPredicate match {
+ case Some(predicate) =>
+ while (pending.isEmpty && storeIterator.hasNext) {
+ val rowPair = storeIterator.next()
+ rowsRead += 1
+ if (predicate.eval(rowPair.key)) {
+ store.remove(rowPair.key)
+ rowsRemoved += 1
+ pending = Some(rowPair)
+ }
+ }
+ pending.isDefined
+ case None => false
+ }
+
+ override def next(): UnsafeRowPair = {
+ if (!hasNext) {
+ throw new NoSuchElementException("End of the iterator")
+ }
+ val rowPair = pending.get
+ pending = None
+ rowPair
+ }
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
index de4af97656a6..ee101cbede9d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
@@ -131,29 +131,10 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase {
}
}
- // A stateful aggregation is still rejected under RTM because HashAggregateExec is not in the
- // allowlist, even though the pipelined shuffle and the StateStore operators now are. When RTM
- // supports aggregation (SPARK-54236), remove this test.
- test("stateful queries not allowed") {
- val inputData = LowLatencyMemoryStream[Int](2)
-
- val df = inputData.toDF()
- .select(col("value").as("key"))
- .groupBy(col("key"))
- .count()
- .select(concat(col("key"), lit("-"), col("count")))
-
- val query = runStreamingQuery("stateful_allowlist", df)
-
- eventually(timeout(60.seconds)) {
- checkError(
- exception = query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException],
- condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST",
- parameters = Map(
- "errorType" -> "operator",
- "message" -> "org.apache.spark.sql.execution.aggregate.HashAggregateExec is"
- )
- )
- }
- }
+ // The "stateful queries not allowed" test that used to live here asserted that a streaming
+ // aggregation was rejected, because it planned into the micro-batch aggregation operators and
+ // HashAggregateExec is not allowlisted. A Real-Time Mode aggregation is now planned as the
+ // streamline aggregate operator, which is allowlisted, so there is no rejection left to assert.
+ // StreamlineStreamingAggregationRealTimeSuite covers the aggregation itself end to end, and the
+ // generic operator-allowlist test above still guards the operators that remain unsupported.
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
index 09dffb4f65cc..0055b1f91a7c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
@@ -962,6 +962,207 @@ class StreamingAggregationSuite extends StateStoreMetricsTest with Assertions {
(inputData2, aggregated2)
}
+ // Streaming aggregation planned as StatefulStreamlineAggregateExec, which merges each input row
+ // against the state store and emits as it goes, rather than the micro-batch operators that emit
+ // once the batch ends. Real-Time Mode always plans it; here the config selects it so the operator
+ // can be covered under a micro-batch trigger too. See StreamlineStreamingAggregationRealTimeSuite
+ // for the Real-Time Mode coverage.
+ private val streamlineEnabled =
+ SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR.key -> "true"
+
+ testWithAllStateVersions("streamline aggregation: update mode emits per input row",
+ streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+
+ val aggregated = inputData.toDF()
+ .groupBy($"value")
+ .agg(count("*"))
+ .as[(Int, Long)]
+
+ testStream(aggregated, Update)(
+ AddData(inputData, 3),
+ CheckLastBatch((3, 1)),
+ AddData(inputData, 3, 2),
+ CheckLastBatch((3, 2), (2, 1)),
+ StopStream,
+ StartStream(),
+ AddData(inputData, 3, 2, 1),
+ CheckLastBatch((3, 3), (2, 2), (1, 1)),
+ // The distinguishing behaviour: four rows for the same key produce four outputs, one per
+ // input. Micro-batch aggregation would emit only the final (4, 4) for this batch.
+ AddData(inputData, 4, 4, 4, 4),
+ CheckLastBatch((4, 1), (4, 2), (4, 3), (4, 4))
+ )
+ }
+
+ testWithAllStateVersions("streamline aggregation: complete mode outputs the whole result table",
+ streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+
+ val aggregated = inputData.toDF()
+ .groupBy($"value")
+ .agg(count("*"))
+ .as[(Int, Long)]
+
+ testStream(aggregated, Complete)(
+ AddData(inputData, 3),
+ CheckLastBatch((3, 1)),
+ AddData(inputData, 2),
+ CheckLastBatch((3, 1), (2, 1)),
+ StopStream,
+ StartStream(),
+ AddData(inputData, 3, 2),
+ CheckLastBatch((3, 2), (2, 2))
+ )
+ }
+
+ testWithAllStateVersions("streamline aggregation: sum, min, max and avg", streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+
+ val aggregated = inputData.toDF()
+ .selectExpr("value % 2 AS key", "value")
+ .groupBy($"key")
+ .agg(sum("value"), min("value"), max("value"), avg("value"))
+ .as[(Int, Long, Int, Int, Double)]
+
+ testStream(aggregated, Complete)(
+ AddData(inputData, 1, 2, 3, 4),
+ // key 1: values 1, 3 -> sum 4, min 1, max 3, avg 2.0
+ // key 0: values 2, 4 -> sum 6, min 2, max 4, avg 3.0
+ CheckLastBatch((1, 4L, 1, 3, 2.0), (0, 6L, 2, 4, 3.0)),
+ AddData(inputData, 5, 6),
+ // key 1 gains 5 -> sum 9, max 5, avg 3.0; key 0 gains 6 -> sum 12, max 6, avg 4.0
+ CheckLastBatch((1, 9L, 1, 5, 3.0), (0, 12L, 2, 6, 4.0))
+ )
+ }
+
+ testWithAllStateVersions("streamline aggregation: multiple grouping keys", streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+
+ val aggregated = inputData.toDF()
+ .selectExpr("value", "value % 2 AS k1", "value % 3 AS k2")
+ .groupBy($"k1", $"k2")
+ .agg(count("*"))
+ .as[(Int, Int, Long)]
+
+ testStream(aggregated, Complete)(
+ AddData(inputData, 1, 2, 3, 4, 5, 6),
+ CheckLastBatch(
+ (1, 1, 1), // 1
+ (0, 2, 1), // 2
+ (1, 0, 1), // 3
+ (0, 1, 1), // 4
+ (1, 2, 1), // 5
+ (0, 0, 1)) // 6
+ )
+ }
+
+ testWithAllStateVersions("streamline aggregation: recovery from a restart keeps state",
+ streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+
+ val aggregated = inputData.toDF()
+ .groupBy($"value")
+ .agg(count("*"))
+ .as[(Int, Long)]
+
+ testStream(aggregated, Complete)(
+ AddData(inputData, 1, 1, 2),
+ CheckLastBatch((1, 2), (2, 1)),
+ StopStream,
+ StartStream(),
+ // The counts continue from the committed state rather than restarting at 1.
+ AddData(inputData, 1, 2),
+ CheckLastBatch((1, 3), (2, 2))
+ )
+ }
+
+ testWithAllStateVersions("streamline aggregation: computed grouping key", streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+ // The grouping key is a COMPUTED expression, not a bare column reference. The final stage is
+ // planned with the original groupingExpressions rather than the post-shuffle attributes, so if
+ // that is wrong this is where it shows: the final stage would try to re-evaluate `value % 3`
+ // against its child's output, which only carries the already-grouped attribute.
+ val aggregated = inputData.toDF()
+ .groupBy(($"value" % 3).as("k"))
+ .agg(count("*"), sum("value"))
+ .as[(Int, Long, Long)]
+
+ testStream(aggregated, Complete)(
+ AddData(inputData, 1, 2, 3, 4, 5, 6),
+ // k=1: 1,4 -> count 2 sum 5 ; k=2: 2,5 -> count 2 sum 7 ; k=0: 3,6 -> count 2 sum 9
+ CheckLastBatch((1, 2L, 5L), (2, 2L, 7L), (0, 2L, 9L))
+ )
+ }
+
+ // Append mode is the one mode that drives the operator's eviction path: a windowed grouping key
+ // is emitted only once the watermark passes it, via EvictionIterator (which removes in hasNext).
+ // Mirrors the stateStoreSave Append test above so the streamline operator is held to the same
+ // watermark/eviction behaviour.
+ testWithAllStateVersions("streamline aggregation: append mode emits windows past the watermark",
+ streamlineEnabled) {
+ val inputData = MemoryStream[Int]
+
+ val aggWithWatermark = inputData.toDF()
+ .withColumn("eventTime", timestamp_seconds($"value"))
+ .withWatermark("eventTime", "10 seconds")
+ .groupBy(window($"eventTime", "5 seconds") as Symbol("window"))
+ .agg(count("*") as Symbol("count"))
+ .select($"window".getField("start").cast("long").as[Long], $"count".as[Long])
+
+ testStream(aggWithWatermark, Append)(
+ StartStream(additionalConfs = Map(SQLConf.SHUFFLE_PARTITIONS.key -> "3")),
+ AddData(inputData, 3, 2, 1, 9),
+ // Nothing is emitted yet: no window has fallen fully below the watermark.
+ CheckLastBatch(),
+ AddData(inputData, 25), // Advance watermark to 15s; windows ending <= 15s are now evictable.
+ // Both closed windows are emitted once, past the watermark, and then evicted:
+ // [0,5) has 1,2,3 -> count 3 ; [5,10) has 9 -> count 1.
+ CheckLastBatch((0, 3), (5, 1))
+ )
+ // Note: state-row *metrics* (e.g. updated rows) are not asserted here because the streamline
+ // operator writes state per input row rather than once per key per batch, so its counts
+ // legitimately differ from the micro-batch stateStoreSave operator. Output correctness -- what
+ // Append actually guarantees -- is covered by the CheckLastBatch assertions above.
+ }
+
+ // The micro-batch aggregation operator (stateStoreSave) and the streamline operator share the
+ // same StreamingAggregationStateManager and state format, so a checkpoint written by one can be
+ // read by the other. The operator-name check in IncrementalExecution must therefore treat the
+ // switch as a supported transition rather than a changed stateful operator (which would abort the
+ // query on restart). Both directions are exercised.
+ test("switching between the micro-batch and streamline aggregation operators keeps state") {
+ Seq(false -> true, true -> false).foreach { case (firstStreamline, secondStreamline) =>
+ withTempDir { dir =>
+ val inputData = MemoryStream[Int]
+ val aggregated = inputData.toDF().groupBy($"value").agg(count("*")).as[(Int, Long)]
+ spark.conf.set(
+ SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR.key, firstStreamline.toString)
+ testStream(aggregated, Complete)(
+ StartStream(checkpointLocation = dir.getAbsolutePath),
+ AddData(inputData, 1, 1, 2),
+ CheckLastBatch((1, 2), (2, 1)),
+ StopStream,
+ Execute { _ =>
+ spark.conf.set(
+ SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR.key, secondStreamline.toString)
+ },
+ StartStream(checkpointLocation = dir.getAbsolutePath),
+ // State survives the operator switch: the counts continue rather than resetting.
+ AddData(inputData, 1, 2),
+ CheckLastBatch((1, 3), (2, 2)),
+ StopStream
+ )
+ }
+ }
+ }
+
+ // Note: the streamline operator's Append branch asserts a watermark is present, but Append
+ // aggregation WITHOUT a watermark is already rejected at analysis time by
+ // UnsupportedOperationChecker (STREAMING_OUTPUT_MODE.UNSUPPORTED_OPERATION), so that assert is a
+ // defensive internal invariant no normal query reaches. The Append-WITH-watermark path is
+ // covered by the test above.
+
@tailrec
private def findStateSchemaNotCompatible(exc: Throwable):
Option[SparkUnsupportedOperationException] = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamlineStreamingAggregationRealTimeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamlineStreamingAggregationRealTimeSuite.scala
new file mode 100644
index 000000000000..5c5409f7a7d5
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamlineStreamingAggregationRealTimeSuite.scala
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.streaming
+
+import org.apache.spark.sql.execution.streaming.StatefulStreamlineAggregateExec
+import org.apache.spark.sql.execution.streaming.operators.stateful.StreamingAggregationStateManager
+import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream}
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Tests for streaming aggregation in Real-Time Mode, which is planned as
+ * [[StatefulStreamlineAggregateExec]] rather than the micro-batch aggregation operators.
+ */
+class StreamlineStreamingAggregationRealTimeSuite extends StreamRealTimeModeSuiteBase
+ with StateStoreMetricsTest {
+
+ import testImplicits._
+
+ def stateFormatVersions: Seq[Int] = StreamingAggregationStateManager.supportedVersions
+
+ def executeFuncWithStateVersionSQLConf(
+ stateVersion: Int,
+ confPairs: Seq[(String, String)],
+ func: => Any): Unit = {
+ withSQLConf(confPairs ++
+ Seq(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> stateVersion.toString): _*) {
+ func
+ }
+ }
+
+ def testWithAllStateVersions(name: String, confPairs: (String, String)*)
+ (func: => Any): Unit = {
+ for (version <- stateFormatVersions) {
+ test(s"$name - state format version $version") {
+ executeFuncWithStateVersionSQLConf(version, confPairs, func)
+ }
+ }
+ }
+
+ testWithAllStateVersions("aggregation runs in Real-Time Mode") {
+ val inputData = LowLatencyMemoryStream[(String, Int)](2)
+
+ val agg = inputData.toDF().select($"_1".as("key"), $"_2".as("value"))
+ .groupBy($"key")
+ .agg(sum("value").as("total"))
+
+ testStream(agg, OutputMode.Update, sink = new ContinuousMemorySink())(
+ StartStream(),
+ AddData(inputData, ("a", 1), ("b", 2), ("a", 3)),
+ // Update mode emits an intermediate result per input row, so "a" is seen twice: once with
+ // its own value and once merged with the earlier one. Micro-batch aggregation would instead
+ // emit only the final value per key per batch.
+ CheckAnswerWithTimeout(60000, ("a", 1L), ("b", 2L), ("a", 4L)),
+ Execute { q =>
+ val aggregates = q.lastExecution.executedPlan.collect {
+ case a: StatefulStreamlineAggregateExec => a
+ }
+ assert(aggregates.size == 1,
+ s"expected the streamline aggregate operator, got:\n${q.lastExecution.executedPlan}")
+ },
+ StopStream
+ )
+ }
+
+ // The eviction counts below are per state store, so the state has to live in a single partition
+ // for them to be predictable. Testing incremental cleanup without pinning the partition count
+ // would need a partition-aware data generator.
+ testWithAllStateVersions("update mode aggregation with incremental cleanup evicts records " +
+ "not removed during incremental cleanup",
+ SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR.key -> "2",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1"
+ ) {
+ val inputData = LowLatencyMemoryStream[Int]
+ val aggWithWatermark = inputData.toDF()
+ .withColumn("eventTime", timestamp_seconds($"value"))
+ .withWatermark("eventTime", "50 seconds")
+ .groupBy(window($"eventTime", "10 seconds") as Symbol("window"))
+ .agg(count("*") as Symbol("count"))
+ .select($"window".getField("end").cast("long").as[Long], $"count".as[Long])
+
+ // With incremental eviction, we evict incrementalCleanupFactor * numInputRows from the
+ // state store.
+ testStream(aggWithWatermark, OutputMode.Update, sink = new ContinuousMemorySink())(
+ StartStream(),
+ AddData(inputData, 9, 19, 29, 39, 49),
+ WaitUntilBatchProcessed(0),
+ // Batch 0: watermark is 0.
+ CheckAnswerWithTimeout(60000, (10, 1), (20, 1), (30, 1), (40, 1), (50, 1)),
+
+ // Batch 1: watermark starts at 0 and moves past the [40, 50) window end.
+ AddData(inputData, 101),
+ // Wait until batch 1 and the no data batch complete
+ WaitUntilBatchProcessed(2),
+ CheckAnswerWithTimeout(60000, (110, 1), (10, 1), (20, 1), (30, 1), (40, 1), (50, 1)),
+
+ Execute { q =>
+ val batch1Metrics = q.recentProgress.filter(_.batchId == 1).head
+ assert(
+ batch1Metrics.stateOperators.head.customMetrics.get("numRowsIncrementallyRemoved") == 0)
+ assert(batch1Metrics.stateOperators.head.numRowsRemoved === 0)
+ },
+
+ AddData(inputData, 100),
+ WaitUntilBatchProcessed(3),
+ CheckAnswerWithTimeout(60000,
+ (110, 1), (110, 2), (10, 1), (20, 1), (30, 1), (40, 1), (50, 1)),
+
+ Execute { q =>
+ val batch3Metrics = q.recentProgress.filter(_.batchId == 3).head
+ // 1 record is in the batch, and the incremental cleanup factor is 2. Thus, 2 rows should be
+ // incrementally cleaned up, but we should still remove 5 total.
+ assert(
+ batch3Metrics.stateOperators.head.customMetrics.get("numRowsIncrementallyRemoved") == 2)
+ assert(batch3Metrics.stateOperators.head.numRowsRemoved === 5)
+ }
+ )
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org