This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 644f24216aa2 [SPARK-58635][SS] Support for streaming aggregation in
Real-Time Mode (RTM)
644f24216aa2 is described below
commit 644f24216aa2c441d63112f5f11cbc06186f123a
Author: Boyang Jerry Peng <[email protected]>
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 <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
(cherry picked from commit 03a304104b9e0eb28f611ad2eb8f7d65da0c33ba)
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../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 c19a92401204..1f3987371e42 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
@@ -3770,6 +3770,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 f065f1de5cdc..92426b627cd2 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
@@ -930,6 +930,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: [email protected]
For additional commands, e-mail: [email protected]
