Skip to content

[SPARK-38204][SS] Use StatefulOpClusteredDistribution for stateful operators with respecting backward compatibility #35673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
HeartSaVioR wants to merge 13 commits into from
Closed

[SPARK-38204][SS] Use StatefulOpClusteredDistribution for stateful operators with respecting backward compatibility #35673

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
fa4cc46
WIP refining code changes...
HeartSaVioR Feb 25, 2022
411c372
Change the approach to inject streaming context in AggregateExec
HeartSaVioR Feb 28, 2022
11defa2
remove comment out of sync now
HeartSaVioR Feb 28, 2022
78a90a3
fix
HeartSaVioR Feb 28, 2022
448aa0a
unnecessary change
HeartSaVioR Feb 28, 2022
9afed1b
fix compliation
HeartSaVioR Feb 28, 2022
5d61b8a
fix failing tests
HeartSaVioR Feb 28, 2022
dc9d1fc
whitespace
HeartSaVioR Feb 28, 2022
0fa1fb4
Scala 2.13 fix
HeartSaVioR Mar 2, 2022
181490c
remove out of sync code comment
HeartSaVioR Mar 3, 2022
cc30d92
remove temporary workaround as the issue has been fixed in master branch
HeartSaVioR Mar 11, 2022
4dd373c
reflect review comments
HeartSaVioR Mar 15, 2022
6c7111b
reflect review comment
HeartSaVioR Mar 15, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1773,6 +1773,23 @@ object SQLConf {
.booleanConf
.createWithDefault(true)

val STATEFUL_OPERATOR_USE_STRICT_DISTRIBUTION =
buildConf("spark.sql.streaming.statefulOperator.useStrictDistribution")
.internal()
.doc("When true, the stateful operator for streaming query will use " +
"StatefulOpClusteredDistribution which guarantees stable state partitioning as long as " +
"the operator provides consistent grouping keys across the lifetime of query. " +
"When false, the stateful operator for streaming query will use ClusteredDistribution " +
"which is not sufficient to guarantee stable state partitioning despite the operator " +
"provides consistent grouping keys across the lifetime of query. " +
"This config will be set to true for new streaming queries to guarantee stable state " +
"partitioning, and set to false for existing streaming queries to not break queries " +
"which are restored from existing checkpoints. Please refer SPARK-38204 for details. " +
"The purpose of this config is only compatibility; DO NOT MANUALLY CHANGE THIS!!!")
Copy link
Member

@xuanyuanking xuanyuanking Mar 14, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's move the description of This config is for compatility only; DO NOT MANUALLY CHANGE THIS!!! as the first sentence in the doc.

.version("3.3.0")
.booleanConf
.createWithDefault(true)

val FILESTREAM_SINK_METADATA_IGNORED =
buildConf("spark.sql.streaming.fileStreamSink.ignoreMetadata")
.internal()
Expand Down
69 changes: 54 additions & 15 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,28 @@ object AggUtils {
}
}

private def createStreamingAggregate(
requiredChildDistributionExpressions: Option[Seq[Expression]] = None,
groupingExpressions: Seq[NamedExpression] = Nil,
aggregateExpressions: Seq[AggregateExpression] = Nil,
aggregateAttributes: Seq[Attribute] = Nil,
initialInputBufferOffset: Int = 0,
resultExpressions: Seq[NamedExpression] = Nil,
child: SparkPlan): SparkPlan = {
createAggregate(
requiredChildDistributionExpressions,
isStreaming = true,
groupingExpressions = groupingExpressions,
aggregateExpressions = aggregateExpressions,
aggregateAttributes = aggregateAttributes,
initialInputBufferOffset = initialInputBufferOffset,
resultExpressions = resultExpressions,
child = child)
}

private def createAggregate(
requiredChildDistributionExpressions: Option[Seq[Expression]] = None,
isStreaming: Boolean = false,
groupingExpressions: Seq[NamedExpression] = Nil,
aggregateExpressions: Seq[AggregateExpression] = Nil,
aggregateAttributes: Seq[Attribute] = Nil,
Expand All @@ -60,6 +80,8 @@ object AggUtils {
if (useHash && !forceSortAggregate) {
HashAggregateExec(
requiredChildDistributionExpressions = requiredChildDistributionExpressions,
isStreaming = isStreaming,
numShufflePartitions = None,
groupingExpressions = groupingExpressions,
aggregateExpressions = mayRemoveAggFilters(aggregateExpressions),
aggregateAttributes = aggregateAttributes,
Expand All @@ -73,6 +95,8 @@ object AggUtils {
if (objectHashEnabled && useObjectHash && !forceSortAggregate) {
ObjectHashAggregateExec(
requiredChildDistributionExpressions = requiredChildDistributionExpressions,
isStreaming = isStreaming,
numShufflePartitions = None,
groupingExpressions = groupingExpressions,
aggregateExpressions = mayRemoveAggFilters(aggregateExpressions),
aggregateAttributes = aggregateAttributes,
Expand All @@ -82,6 +106,8 @@ object AggUtils {
} else {
SortAggregateExec(
requiredChildDistributionExpressions = requiredChildDistributionExpressions,
isStreaming = isStreaming,
numShufflePartitions = None,
groupingExpressions = groupingExpressions,
aggregateExpressions = mayRemoveAggFilters(aggregateExpressions),
aggregateAttributes = aggregateAttributes,
Expand Down Expand Up @@ -121,7 +147,7 @@ object AggUtils {
// If we have session window expression in aggregation, we add MergingSessionExec to
// merge sessions with calculating aggregation values.
val interExec: SparkPlan = mayAppendMergingSessionExec(groupingExpressions,
aggregateExpressions, partialAggregate)
aggregateExpressions, partialAggregate, isStreaming = false)
Copy link
Member

@xuanyuanking xuanyuanking Mar 14, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: how about using default value in mayAppendMergingSessionExec and let here unchange?


// 2. Create an Aggregate Operator for final aggregations.
val finalAggregateExpressions = aggregateExpressions.map(_.copy(mode = Final))
Expand Down Expand Up @@ -153,7 +179,8 @@ object AggUtils {
// If we have session window expression in aggregation, we add UpdatingSessionsExec to
// calculate sessions for input rows and update rows' session column, so that further
// aggregations can aggregate input rows for the same session.
val maySessionChild = mayAppendUpdatingSessionExec(groupingExpressions, child)
val maySessionChild = mayAppendUpdatingSessionExec(groupingExpressions, child,
Copy link
Member

@xuanyuanking xuanyuanking Mar 14, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

isStreaming = false)

val distinctAttributes = normalizedNamedDistinctExpressions.map(_.toAttribute)
val groupingAttributes = groupingExpressions.map(_.toAttribute)
Expand Down Expand Up @@ -290,7 +317,7 @@ object AggUtils {
val partialAggregate: SparkPlan = {
val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = Partial))
val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
createAggregate(
createStreamingAggregate(
groupingExpressions = groupingExpressions,
aggregateExpressions = aggregateExpressions,
aggregateAttributes = aggregateAttributes,
Expand All @@ -302,7 +329,7 @@ object AggUtils {
val partialMerged1: SparkPlan = {
val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = PartialMerge))
val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
createAggregate(
createStreamingAggregate(
requiredChildDistributionExpressions =
Some(groupingAttributes),
groupingExpressions = groupingAttributes,
Expand All @@ -320,7 +347,7 @@ object AggUtils {
val partialMerged2: SparkPlan = {
val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = PartialMerge))
val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
createAggregate(
createStreamingAggregate(
requiredChildDistributionExpressions =
Some(groupingAttributes),
groupingExpressions = groupingAttributes,
Expand Down Expand Up @@ -348,7 +375,7 @@ object AggUtils {
// projection:
val finalAggregateAttributes = finalAggregateExpressions.map(_.resultAttribute)

createAggregate(
createStreamingAggregate(
requiredChildDistributionExpressions = Some(groupingAttributes),
groupingExpressions = groupingAttributes,
aggregateExpressions = finalAggregateExpressions,
Expand Down Expand Up @@ -407,7 +434,7 @@ object AggUtils {
val partialAggregate: SparkPlan = {
val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = Partial))
val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
createAggregate(
createStreamingAggregate(
groupingExpressions = groupingExpressions,
aggregateExpressions = aggregateExpressions,
aggregateAttributes = aggregateAttributes,
Expand All @@ -424,7 +451,8 @@ object AggUtils {
// this is to reduce amount of rows to shuffle
MergingSessionsExec(
requiredChildDistributionExpressions = None,
requiredChildDistributionOption = None,
isStreaming = true,
numShufflePartitions = None,
groupingExpressions = groupingAttributes,
sessionExpression = sessionExpression,
aggregateExpressions = aggregateExpressions,
Expand All @@ -447,8 +475,10 @@ object AggUtils {
val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = PartialMerge))
val aggregateAttributes = aggregateExpressions.map(_.resultAttribute)
MergingSessionsExec(
requiredChildDistributionExpressions = None,
requiredChildDistributionOption = Some(restored.requiredChildDistribution),
requiredChildDistributionExpressions = Some(groupingWithoutSessionAttributes),
isStreaming = true,
// This will be replaced with actual value in state rule.
Copy link
Member

@xuanyuanking xuanyuanking Mar 14, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's link the state rule class name here?

Copy link
Contributor Author

@HeartSaVioR HeartSaVioR Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't create a dedicate class for state rule. See

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala

Lines 105 to 209 in a30575e

/** Locates save/restore pairs surrounding aggregation. */
val state = new Rule[SparkPlan] {
/**
* Ensures that this plan DOES NOT have any stateful operation in it whose pipelined execution
* depends on this plan. In other words, this function returns true if this plan does
* have a narrow dependency on a stateful subplan.
*/
private def hasNoStatefulOp(plan: SparkPlan): Boolean = {
var statefulOpFound = false
def findStatefulOp(planToCheck: SparkPlan): Unit = {
planToCheck match {
case s: StatefulOperator =>
statefulOpFound = true
case e: ShuffleExchangeLike =>
// Don't search recursively any further as any child stateful operator as we
// are only looking for stateful subplans that this plan has narrow dependencies on.
case p: SparkPlan =>
p.children.foreach(findStatefulOp)
}
}
findStatefulOp(plan)
!statefulOpFound
}
override def apply(plan: SparkPlan): SparkPlan = plan transform {
case StateStoreSaveExec(keys, None, None, None, stateFormatVersion,
UnaryExecNode(agg,
StateStoreRestoreExec(_, None, _, child))) =>
val aggStateInfo = nextStatefulOperationStateInfo
StateStoreSaveExec(
keys,
Some(aggStateInfo),
Some(outputMode),
Some(offsetSeqMetadata.batchWatermarkMs),
stateFormatVersion,
agg.withNewChildren(
StateStoreRestoreExec(
keys,
Some(aggStateInfo),
stateFormatVersion,
child) :: Nil))
case SessionWindowStateStoreSaveExec(keys, session, None, None, None, stateFormatVersion,
UnaryExecNode(agg,
SessionWindowStateStoreRestoreExec(_, _, None, None, _, child))) =>
val aggStateInfo = nextStatefulOperationStateInfo
SessionWindowStateStoreSaveExec(
keys,
session,
Some(aggStateInfo),
Some(outputMode),
Some(offsetSeqMetadata.batchWatermarkMs),
stateFormatVersion,
agg.withNewChildren(
SessionWindowStateStoreRestoreExec(
keys,
session,
Some(aggStateInfo),
Some(offsetSeqMetadata.batchWatermarkMs),
stateFormatVersion,
child) :: Nil))
case StreamingDeduplicateExec(keys, child, None, None) =>
StreamingDeduplicateExec(
keys,
child,
Some(nextStatefulOperationStateInfo),
Some(offsetSeqMetadata.batchWatermarkMs))
case m: FlatMapGroupsWithStateExec =>
// We set this to true only for the first batch of the streaming query.
val hasInitialState = (currentBatchId == 0L && m.hasInitialState)
m.copy(
stateInfo = Some(nextStatefulOperationStateInfo),
batchTimestampMs = Some(offsetSeqMetadata.batchTimestampMs),
eventTimeWatermark = Some(offsetSeqMetadata.batchWatermarkMs),
hasInitialState = hasInitialState
)
case j: StreamingSymmetricHashJoinExec =>
j.copy(
stateInfo = Some(nextStatefulOperationStateInfo),
eventTimeWatermark = Some(offsetSeqMetadata.batchWatermarkMs),
stateWatermarkPredicates =
StreamingSymmetricHashJoinHelper.getStateWatermarkPredicates(
j.left.output, j.right.output, j.leftKeys, j.rightKeys, j.condition.full,
Some(offsetSeqMetadata.batchWatermarkMs)))
case l: StreamingGlobalLimitExec =>
l.copy(
stateInfo = Some(nextStatefulOperationStateInfo),
outputMode = Some(outputMode))
case StreamingLocalLimitExec(limit, child) if hasNoStatefulOp(child) =>
// Optimize limit execution by replacing StreamingLocalLimitExec (consumes the iterator
// completely) to LocalLimitExec (does not consume the iterator) when the child plan has
// no stateful operator (i.e., consuming the iterator is not needed).
LocalLimitExec(limit, child)
}
}

numShufflePartitions = None,
groupingExpressions = groupingAttributes,
sessionExpression = sessionExpression,
aggregateExpressions = aggregateExpressions,
Expand Down Expand Up @@ -476,8 +506,8 @@ object AggUtils {
// projection:
val finalAggregateAttributes = finalAggregateExpressions.map(_.resultAttribute)

createAggregate(
requiredChildDistributionExpressions = Some(groupingAttributes),
createStreamingAggregate(
requiredChildDistributionExpressions = Some(groupingWithoutSessionAttributes),
groupingExpressions = groupingAttributes,
aggregateExpressions = finalAggregateExpressions,
aggregateAttributes = finalAggregateAttributes,
Expand All @@ -491,10 +521,15 @@ object AggUtils {

private def mayAppendUpdatingSessionExec(
groupingExpressions: Seq[NamedExpression],
maybeChildPlan: SparkPlan): SparkPlan = {
maybeChildPlan: SparkPlan,
isStreaming: Boolean): SparkPlan = {
groupingExpressions.find(_.metadata.contains(SessionWindow.marker)) match {
case Some(sessionExpression) =>
UpdatingSessionsExec(
isStreaming = isStreaming,
// numShufflePartitions will be set to None, and replaced to the actual value in the
// state rule if the query is streaming.
Copy link
Member

@xuanyuanking xuanyuanking Mar 14, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

Copy link
Contributor Author

@HeartSaVioR HeartSaVioR Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here

numShufflePartitions = None,
groupingExpressions.map(_.toAttribute),
sessionExpression.toAttribute,
maybeChildPlan)
Expand All @@ -506,7 +541,8 @@ object AggUtils {
private def mayAppendMergingSessionExec(
groupingExpressions: Seq[NamedExpression],
aggregateExpressions: Seq[AggregateExpression],
partialAggregate: SparkPlan): SparkPlan = {
partialAggregate: SparkPlan,
isStreaming: Boolean): SparkPlan = {
groupingExpressions.find(_.metadata.contains(SessionWindow.marker)) match {
case Some(sessionExpression) =>
val aggExpressions = aggregateExpressions.map(_.copy(mode = PartialMerge))
Expand All @@ -519,7 +555,10 @@ object AggUtils {

MergingSessionsExec(
requiredChildDistributionExpressions = Some(groupingWithoutSessionsAttributes),
requiredChildDistributionOption = None,
isStreaming = isStreaming,
// numShufflePartitions will be set to None, and replaced to the actual value in the
// state rule if the query is streaming.
numShufflePartitions = None,
groupingExpressions = groupingAttributes,
sessionExpression = sessionExpression,
aggregateExpressions = aggExpressions,
Expand Down
23 changes: 20 additions & 3 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/BaseAggregateExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,15 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference,
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Final, PartialMerge}
import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistribution, Distribution, UnspecifiedDistribution}
import org.apache.spark.sql.execution.{AliasAwareOutputPartitioning, ExplainUtils, UnaryExecNode}
import org.apache.spark.sql.execution.streaming.StatefulOperatorPartitioning

/**
* Holds common logic for aggregate operators
*/
trait BaseAggregateExec extends UnaryExecNode with AliasAwareOutputPartitioning {
def requiredChildDistributionExpressions: Option[Seq[Expression]]
def isStreaming: Boolean
def numShufflePartitions: Option[Int]
def groupingExpressions: Seq[NamedExpression]
def aggregateExpressions: Seq[AggregateExpression]
def aggregateAttributes: Seq[Attribute]
Expand Down Expand Up @@ -92,7 +95,20 @@ trait BaseAggregateExec extends UnaryExecNode with AliasAwareOutputPartitioning
override def requiredChildDistribution: List[Distribution] = {
requiredChildDistributionExpressions match {
case Some(exprs) if exprs.isEmpty => AllTuples :: Nil
case Some(exprs) => ClusteredDistribution(exprs) :: Nil
case Some(exprs) =>
if (isStreaming) {
numShufflePartitions match {
case Some(parts) =>
StatefulOperatorPartitioning.getClusteredDistributionWithBackwardCompatibility(
exprs, parts, conf) :: Nil

case _ =>
throw new IllegalStateException("Expected to set the number of partitions before " +
Copy link
Contributor

@c21 c21 Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can also add a require assertion at class level, e.g. at Line 38

require(isStreaming == numShufflePartitions.isDefined, "Expected to set the number of partitions for streaming aggregate")

Or we can define only one variable numStatefulShufflePartitions: Option[Int] instead of two: isStreaming and numShufflePartitions?

Copy link
Contributor Author

@HeartSaVioR HeartSaVioR Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We create a node with numShufflePartitions = None and replace the value in state rule. That said, we can't check the condition before state rule has been performed.

Copy link
Contributor

@LuciferYang LuciferYang Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should use the new error framework to throw exception in newly added code.

Copy link
Contributor Author

@HeartSaVioR HeartSaVioR Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The error framework is for user facing errors. This is something like "this should not be called, internal error". I just made the error message be general to make our developer life be easier.

Copy link
Contributor

@LuciferYang LuciferYang Mar 15, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok

"constructing required child distribution!")
}
} else {
ClusteredDistribution(exprs) :: Nil
}
case None => UnspecifiedDistribution :: Nil
}
}
Expand All @@ -102,7 +118,8 @@ trait BaseAggregateExec extends UnaryExecNode with AliasAwareOutputPartitioning
*/
def toSortAggregate: SortAggregateExec = {
SortAggregateExec(
requiredChildDistributionExpressions, groupingExpressions, aggregateExpressions,
aggregateAttributes, initialInputBufferOffset, resultExpressions, child)
requiredChildDistributionExpressions, isStreaming, numShufflePartitions, groupingExpressions,
aggregateExpressions, aggregateAttributes, initialInputBufferOffset, resultExpressions,
child)
}
}
2 changes: 2 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ import org.apache.spark.util.Utils
*/
case class HashAggregateExec(
requiredChildDistributionExpressions: Option[Seq[Expression]],
isStreaming: Boolean,
numShufflePartitions: Option[Int],
groupingExpressions: Seq[NamedExpression],
aggregateExpressions: Seq[AggregateExpression],
aggregateAttributes: Seq[Attribute],
Expand Down
15 changes: 2 additions & 13 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/MergingSessionsExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Expression, MutableProjection, NamedExpression, SortOrder, UnsafeRow}
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.metric.SQLMetrics

Expand All @@ -41,7 +40,8 @@ import org.apache.spark.sql.execution.metric.SQLMetrics
*/
case class MergingSessionsExec(
requiredChildDistributionExpressions: Option[Seq[Expression]],
requiredChildDistributionOption: Option[Seq[Distribution]],
isStreaming: Boolean,
numShufflePartitions: Option[Int],
groupingExpressions: Seq[NamedExpression],
sessionExpression: NamedExpression,
aggregateExpressions: Seq[AggregateExpression],
Expand All @@ -59,17 +59,6 @@ case class MergingSessionsExec(

override def outputOrdering: Seq[SortOrder] = child.outputOrdering

override def requiredChildDistribution: List[Distribution] = {
requiredChildDistributionExpressions match {
case Some(exprs) if exprs.isEmpty => AllTuples :: Nil
case Some(exprs) => ClusteredDistribution(exprs) :: Nil
case None => requiredChildDistributionOption match {
case Some(distributions) => distributions.toList
case None => UnspecifiedDistribution :: Nil
}
}
}

override def requiredChildOrdering: Seq[Seq[SortOrder]] = {
Seq((keyWithoutSessionExpressions ++ Seq(sessionExpression)).map(SortOrder(_, Ascending)))
}
Expand Down
2 changes: 2 additions & 0 deletions ...ore/src/main/scala/org/apache/spark/sql/execution/aggregate/ObjectHashAggregateExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ import org.apache.spark.sql.execution.metric.SQLMetrics
*/
case class ObjectHashAggregateExec(
requiredChildDistributionExpressions: Option[Seq[Expression]],
isStreaming: Boolean,
numShufflePartitions: Option[Int],
groupingExpressions: Seq[NamedExpression],
aggregateExpressions: Seq[AggregateExpression],
aggregateAttributes: Seq[Attribute],
Expand Down
2 changes: 2 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ import org.apache.spark.sql.internal.SQLConf
*/
case class SortAggregateExec(
requiredChildDistributionExpressions: Option[Seq[Expression]],
isStreaming: Boolean,
numShufflePartitions: Option[Int],
groupingExpressions: Seq[NamedExpression],
aggregateExpressions: Seq[AggregateExpression],
aggregateAttributes: Seq[Attribute],
Expand Down
18 changes: 17 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/UpdatingSessionsExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, SortOrder}
import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistribution, Distribution, Partitioning}
import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
import org.apache.spark.sql.execution.streaming.StatefulOperatorPartitioning

/**
* This node updates the session window spec of each input rows via analyzing neighbor rows and
Expand All @@ -35,6 +36,8 @@ import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
* Refer [[UpdatingSessionsIterator]] for more details.
*/
case class UpdatingSessionsExec(
isStreaming: Boolean,
numShufflePartitions: Option[Int],
groupingExpression: Seq[Attribute],
sessionExpression: Attribute,
child: SparkPlan) extends UnaryExecNode {
Expand Down Expand Up @@ -63,7 +66,20 @@ case class UpdatingSessionsExec(
if (groupingWithoutSessionExpression.isEmpty) {
AllTuples :: Nil
} else {
ClusteredDistribution(groupingWithoutSessionExpression) :: Nil
if (isStreaming) {
numShufflePartitions match {
case Some(parts) =>
StatefulOperatorPartitioning.getClusteredDistributionWithBackwardCompatibility(
groupingWithoutSessionExpression, parts, conf) :: Nil

case _ =>
throw new IllegalStateException("Expected to set the number of partitions before " +
"constructing required child distribution!")
}

} else {
ClusteredDistribution(groupingWithoutSessionExpression) :: Nil
}
}
}

Expand Down
13 changes: 5 additions & 8 deletions .../src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Expression, SortOrder, UnsafeRow}
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution}
import org.apache.spark.sql.catalyst.plans.physical.Distribution
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.streaming.StreamingSymmetricHashJoinHelper._
import org.apache.spark.sql.execution.streaming.state._
Expand Down Expand Up @@ -93,13 +93,10 @@ case class FlatMapGroupsWithStateExec(
* to have the same grouping so that the data are co-lacated on the same task.
*/
override def requiredChildDistribution: Seq[Distribution] = {
// NOTE: Please read through the NOTE on the classdoc of StatefulOpClusteredDistribution
// before making any changes.
// TODO(SPARK-38204)
ClusteredDistribution(
groupingAttributes, requiredNumPartitions = stateInfo.map(_.numPartitions)) ::
ClusteredDistribution(
initialStateGroupAttrs, requiredNumPartitions = stateInfo.map(_.numPartitions)) ::
StatefulOperatorPartitioning.getClusteredDistributionWithBackwardCompatibility(
groupingAttributes, getStateInfo, conf) ::
StatefulOperatorPartitioning.getClusteredDistributionWithBackwardCompatibility(
initialStateGroupAttrs, getStateInfo, conf) ::
Nil
}

Expand Down
Loading