[SPARK-53941][SS] Support AQE in stateless streaming workloads #52642
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,61 @@ | ||
| /* | ||
| * 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.internal.{Logging, LogKeys} | ||
| import org.apache.spark.sql.execution.SparkPlan | ||
| import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec | ||
| import org.apache.spark.sql.execution.adaptive.QueryStageExec | ||
| import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec | ||
|
|
||
| /** | ||
| * This is an utility object placing methods to traverse the query plan for streaming query. | ||
| * This is used for patterns of traversal which are repeated in multiple places. | ||
| */ | ||
| object StreamingQueryPlanTraverseHelper extends Logging { | ||
| def collectFromUnfoldedPlan[B]( | ||
| executedPlan: SparkPlan)( | ||
| pf: PartialFunction[SparkPlan, B]): Seq[B] = { | ||
| executedPlan.flatMap { | ||
| // InMemoryTableScanExec is a node to represent a cached plan. The node has underlying | ||
| // actual executed plan, which we should traverse to collect the required information. | ||
| case s: InMemoryTableScanExec => collectFromUnfoldedPlan(s.relation.cachedPlan)(pf) | ||
|
|
||
| // AQE physical node contains the executed plan, pick the plan. | ||
| // In most cases, AQE physical node is expected to contain the final plan, which is | ||
| // appropriate for the caller. | ||
| // Even it does not contain the final plan (in whatever reason), we just provide the | ||
| // plan as best effort, as there is no better way around. | ||
| case a: AdaptiveSparkPlanExec => | ||
| if (!a.isFinalPlan) { | ||
| logWarning(log"AQE plan is captured, but the executed plan in AQE plan is not" + | ||
| log"the final one. Providing incomplete executed plan. AQE plan: ${MDC( | ||
| LogKeys.AQE_PLAN, a)}") | ||
| } | ||
| collectFromUnfoldedPlan(a.executedPlan)(pf) | ||
|
|
||
| // There are several AQE-specific leaf nodes which covers shuffle. We should pick the | ||
| // underlying plan of these nodes, since the underlying plan has the actual executed | ||
| // nodes which we want to collect metrics. | ||
| case e: QueryStageExec => collectFromUnfoldedPlan(e.plan)(pf) | ||
|
|
||
| case p if pf.isDefinedAt(p) => Seq(pf(p)) | ||
| case _ => Seq.empty[B] | ||
| } | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -27,7 +27,7 @@ import org.apache.spark.internal.LogKeys | |
| import org.apache.spark.internal.LogKeys.BATCH_ID | ||
| import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder | ||
| import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, CurrentBatchTimestamp, CurrentDate, CurrentTimestamp, FileSourceMetadataAttribute, LocalTimestamp} | ||
| import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Deduplicate, DeduplicateWithinWatermark, Distinct, FlatMapGroupsInPandasWithState, FlatMapGroupsWithState, GlobalLimit, Join, LeafNode, LocalRelation, LogicalPlan, Project, StreamSourceAwareLogicalPlan, TransformWithState, TransformWithStateInPySpark} | ||
| import org.apache.spark.sql.catalyst.streaming.{StreamingRelationV2, WriteToStream} | ||
| import org.apache.spark.sql.catalyst.trees.TreePattern.CURRENT_LIKE | ||
| import org.apache.spark.sql.catalyst.util.truncatedString | ||
|
|
@@ -344,9 +344,40 @@ class MicroBatchExecution( | |
| setLatestExecutionContext(execCtx) | ||
|
|
||
| populateStartOffsets(execCtx, sparkSessionForStream) | ||
|
|
||
| // SPARK-53941: This code path is executed for the first batch, regardless of whether it's a | ||
| // fresh new run or restart. | ||
| disableAQESupportInStatelessIfUnappropriated(sparkSessionForStream) | ||
|
Contributor
There was a problem hiding this comment. Do you have any plans to try to either remove this or have a way to opt-out since other checks should theoretically be enough when choosing wither to insert an adaptive plan? We (sadly) only recently discovered the performance benefits of manually re-enabling AQE inside a FEB and it would be nice to have that work by default, and I think this is still preventing that from happening?
Contributor
Author
There was a problem hiding this comment. It doesn't prevent you to do that unless your "streaming part" (before FEB) of the query is stateful. Do I miss something?
Contributor
There was a problem hiding this comment. No that's the case I'm talking about. A common use case for us is a stateful stream with a FEB to merge into a delta table, and enabling AQE inside the batch helps with the merge performance
Contributor
Author
There was a problem hiding this comment. That is very dependent to the DSv1 sink implementation - FEB sink is somehow safe because we finalize RDD plan for streaming part. If not, it should have been problematic likewise I mentioned in the PR description. I'd still leave this to user's risk - you can still override AQE config in FEB sink after this fix to get the same behavior (which I really hate that config is overridable in FEB sink but...)
Contributor
There was a problem hiding this comment. I guess my question is why isn't https://github.com/apache/spark/pull/52642/files#diff-1473ce928d51533500dcd01054219c3b1fc8e4c69c763a68cb01a22e145ce1dcR59-R67 "enough". Why does the AQE setting also need to be disabled? |
||
|
|
||
| logInfo(log"Stream started from ${MDC(LogKeys.STREAMING_OFFSETS_START, execCtx.startOffsets)}") | ||
| execCtx | ||
| } | ||
|
|
||
| private def disableAQESupportInStatelessIfUnappropriated( | ||
| sparkSessionToRunBatches: SparkSession): Unit = { | ||
| def containsStatefulOperator(p: LogicalPlan): Boolean = { | ||
| p.exists { | ||
| case node: Aggregate if node.isStreaming => true | ||
| case node: Deduplicate if node.isStreaming => true | ||
| case node: DeduplicateWithinWatermark if node.isStreaming => true | ||
| case node: Distinct if node.isStreaming => true | ||
| case node: Join if node.left.isStreaming && node.right.isStreaming => true | ||
| case node: FlatMapGroupsWithState if node.isStreaming => true | ||
|
Contributor
There was a problem hiding this comment. What about AIPWS and TWSInPandas ?
Contributor
Author
There was a problem hiding this comment.
|
||
| case node: FlatMapGroupsInPandasWithState if node.isStreaming => true | ||
| case node: TransformWithState if node.isStreaming => true | ||
| case node: TransformWithStateInPySpark if node.isStreaming => true | ||
| case node: GlobalLimit if node.isStreaming => true | ||
| case _ => false | ||
| } | ||
| } | ||
|
|
||
| if (containsStatefulOperator(analyzedPlan)) { | ||
| // SPARK-53941: We disable AQE for stateful workloads as of now. | ||
| logWarning(log"Disabling AQE since AQE is not supported in stateful workloads.") | ||
| sparkSessionToRunBatches.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") | ||
|
Contributor
There was a problem hiding this comment. Do you think it makes sense to make this restriction stronger? In addition to set this conf there, we also set this AQE disabled as a property of the query, and perform a conf check every batch. This can prevent users override this config during foreachBatch and save potential support burden. I can add a followup if you think it makes sense
Contributor
Author
There was a problem hiding this comment. While I think this direction is great in general, it does not actually prevent the issue completely, because we will check the config change after the change is applied for a single batch. If we want to prevent the issue holistically, we should actually disallow (either ignore with logging, or even fail the query) setting the configs under streaming engine's control in FEB sink. This may need more thought since when user function is executed it's not under streaming engine's control. For example, spark session from given DataFrame in user function should not allow changing a list of configs under streaming engine's control. But even with the above, we can't prevent the case of referencing external spark session in user function. That's more of a fundamental issue, so it'd probably be fine if the solution does not address this. |
||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Repeatedly attempts to run batches as data arrives. | ||
| */ | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
shall we have a dedicated config for streaming AQE? It seems overkill to turn off AQE entirely.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm OK with introducing a config to control over streaming AQE - should we do this as layered configs? e.g. If AQE is turned off via spark.sql.adaptive.enabled, we turn off both batch and streaming. The config for streaming AQE takes effect only when the AQE config spark.sql.adaptive.enabled is turned on.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
layered config SGTM.