Ilya created SPARK-59204:
----------------------------
Summary: [SQL][AQE][CACHE] Severe planning-time regression in
iterative cached DataFrame workload after Spark 3.4.3 -> 4.1.3
Key: SPARK-59204
URL: https://issues.apache.org/jira/browse/SPARK-59204
Project: Spark
Issue Type: Bug
Components: SQL
Affects Versions: 4.1.3
Environment: Regression workload:
* Apache Spark: 4.1.3
* Scala: 2.13
* Java: 17
* CatBoost Spark: {{ai.catboost:catboost-spark_4.1_2.13:1.2.10}}
Regression baseline:
* Apache Spark: 3.4.3
* Scala: 2.13
The standalone reproducer itself uses only Spark DataFrame APIs and synthetic
data.
h2.
Reporter: Ilya
Attachments: spark_aqe_cache_plan_growth_reproducer.py
We observe a severe performance regression after migrating an iterative PySpark
DataFrame workload from Apache Spark 3.4.3 to Spark 4.1.3.
The workload repeatedly:
# derives features from the current cached state;
# derives a prediction from those features;
# appends the prediction back to the previous state;
# repartitions the new state;
# caches and materializes it with count().
The previous cached state is therefore referenced both directly by the next
state and indirectly through the feature/prediction branch.
Spark 3.4.3 shows gradual degradation as this recursive dependency becomes more
complex.
Spark 4.1.3 with AQE enabled during cache materialization shows dramatically
stronger nonlinear degradation.
Disabling AQE only while materializing the newly cached iterative state reduces
the degradation by more than an order of magnitude.
A standalone PySpark reproducer without CatBoost, Iceberg, S3, Airflow, or
production data is attached.
h2. Original workload regression
The production workload was first tested to isolate the dependency upgrade from
the Spark upgrade.
Results:
{code:java}
Spark 3.4.3 + CatBoost Spark 1.2.8
~48 minutes
Spark 3.4.3 + CatBoost Spark 1.2.10
~48 minutes
{code}
Therefore, upgrading CatBoost from 1.2.8 to 1.2.10 does not reproduce the
regression on Spark 3.4.3.
After upgrading Spark:
{code:java}
Spark 4.1.3 + CatBoost Spark 1.2.10
AQE enabled during iterative cache materialization
{code}
the workload exhibits severe degradation and, with the default
spark.sql.maxPlanStringLength, eventually fails with:
{code:java}
java.lang.OutOfMemoryError: Java heap space
{code}
h2. Effect of spark.sql.maxPlanStringLength
Changing only:
{code:java}
spark.sql.maxPlanStringLength=8192
{code}
prevents the driver OOM.
However, it does not fix the underlying performance degradation.
The application continues running, but late iterations become extremely
expensive and the full workload takes approximately 3-4 hours instead of
approximately 48 minutes on Spark 3.4.3.
This OOM behavior appears related to SPARK-50992 and is being reported there
separately.
h2. Original workload pattern
The production workload is conceptually similar to:
{code:python}
extended = state.union(dummy_rows)
features = (
feature_pipeline
.transform(extended)
.cache()
)
features.count()
prediction = model.transform(features)
state = (
state
.union(prediction)
.repartition(200, "key1", "key2")
)
state = state.cache()
state.count()
{code}
The previous state is therefore referenced:
# directly in the next outer Union;
# indirectly again through features -> prediction.
h2. Production planning-time degradation
One timing measurement surrounds only the lazy model transform() call.
There is no Spark action inside this timing window, so this is not
executor-side model inference time.
With Spark 4.1.3 and AQE enabled during cache materialization:
{code:java}
step 11: 0.322 s
step 12: 0.703 s
step 13: 0.989 s
step 14: 1.570 s
step 15: 2.956 s
step 16: 5.797 s
step 17: 11.336 s
step 18: 22.606 s
{code}
In another run the growth continued:
{code:java}
step 13: 0.852 s
step 14: 1.593 s
step 15: 2.982 s
step 16: 5.828 s
step 17: 11.794 s
step 18: 22.561 s
step 19: 45.476 s
step 20: 94.096 s
step 21: 179.974 s
step 22: 367.955 s
{code}
The later sequence is close to doubling on each iteration.
Execution time also starts degrading.
For example:
{code:java}
feature execution:
step 16: 38.8 s
step 18: 59.1 s
step 19: 84.0 s
step 20: 134.5 s
step 21: 240.9 s
step 22: 426.4 s
{code}
and iterative-state materialization:
{code:java}
step 16: 93.7 s
step 18: 122.6 s
step 20: 206.2 s
step 21: 296.2 s
step 22: 534.4 s
{code}
h2. AQE workaround in the original workload
The following scoped workaround dramatically improves the workload:
{code:python}
previous_aqe = spark.conf.get(
"spark.sql.adaptive.enabled"
)
try:
spark.conf.set(
"spark.sql.adaptive.enabled",
"false",
)
state = state.cache()
state.count()
finally:
spark.conf.set(
"spark.sql.adaptive.enabled",
previous_aqe,
)
{code}
AQE remains enabled during the rest of the application.
With this workaround on Spark 4.1.3:
* all 26 iterations complete;
* feature calculation remains approximately stable;
* state execution/materialization remains approximately 74-83 seconds across
the iterations;
* sum of measured iteration times is approximately 48 minutes;
* full application runtime is approximately 60 minutes.
We also observe substantially fewer nested Union branches in cached-plan
representations compared with Spark 4.1.3 when AQE remains enabled during
materialization.
h2. Standalone PySpark reproducer
The attached file:
{code:java}
spark_aqe_cache_plan_growth_reproducer.py
{code}
uses only synthetic data and Spark DataFrame APIs.
It intentionally reproduces this dependency structure:
{code:java}
previous state -------------------------------+
| |
+--> derived features --> prediction ----+--> union
|
repartition
|
cache
|
count
{code}
The reproducer can be run with AQE enabled:
{code:bash}
spark-submit spark_aqe_cache_plan_growth_reproducer.py --mode aqe_on
--retention isolated
{code}
or with AQE disabled only around state cache materialization:
{code:bash}
spark-submit spark_aqe_cache_plan_growth_reproducer.py --mode
scoped_aqe_off --retention isolated
{code}
retention=isolated unpersists the previous state and features only after the
new state has been fully materialized.
This was done to reduce accumulated cache-memory usage as a confounding factor.
h2. Standalone reproduction: Spark 4.1.3
With AQE enabled:
{code:java}
state materialization:
step 1: ~2.0 s
step 4: ~2.9 s
step 5: ~5.6 s
step 6: ~21.0 s
step 7: ~110-118 s
{code}
The result was reproduced in multiple Spark 4.1.3 runs.
For example, one run produced:
{code:java}
step 1: 1.92 s
step 4: 2.66 s
step 5: 5.62 s
step 6: 21.05 s
step 7: 115.67 s
{code}
Another produced:
{code:java}
step 1: 3.49 s
step 4: 7.72 s
step 5: 7.64 s
step 6: 25.65 s
step 7: 117.87 s
{code}
Iteration 8 becomes extremely slow.
h2. Standalone reproduction: Spark 4.1.3 with scoped AQE disable
On the same Spark 4.1.3 environment, disabling AQE only while materializing the
newly cached state dramatically reduces the degradation.
For example:
{code:java}
step 1: 2.38 s
step 4: 2.30 s
step 7: 5.39 s
{code}
The equivalent AQE-enabled iteration 7 requires approximately:
{code:java}
109-118 s
{code}
versus:
{code:java}
5.4 s
{code}
with AQE disabled only during materialization.
This is approximately a 20x difference at iteration 7.
The workaround does not completely eliminate the inherent cost of the
recursively growing dependency structure, but it dramatically reduces the Spark
4.1.3-specific degradation.
h2. Standalone comparison with Spark 3.4.3
The same synthetic workload on Spark 3.4.3 with AQE enabled also becomes
gradually more expensive, but the degradation is substantially less aggressive.
Example:
{code:java}
state materialization:
step 3: 6.62 s
step 5: 10.32 s
step 9: 20.46 s
{code}
The shape of the degradation differs significantly from Spark 4.1.3:
{code:java}
Spark 3.4.3:
gradual growth
Spark 4.1.3 + AQE:
~2 s -> ~5 s -> ~21 s -> ~115 s
{code}
Therefore, the issue is not that recursive iterative plans have zero cost on
Spark 3.4.3.
Rather, Spark 4.1.3 with AQE enabled during cached-state materialization
amplifies this cost dramatically.
h2. Physical-plan diagnostics
Initially, inspecting only the outer optimized logical plan appeared to show no
growth.
At iterations 2, 4, and 6 it remained approximately:
{code:java}
optimized nodes: 5
Union: 1
InMemoryRelation: 2
{code}
However, InMemoryRelation hides the physical cached plan.
We therefore recursively traversed:
* AdaptiveSparkPlanExec.executedPlan
* QueryStageExec.plan
* InMemoryTableScanExec.relation.cachedPlan
without calling treeString() or explain().
h3. Recursive physical-plan references
On Spark 4.1.3 with AQE enabled, this exposes rapid recursive growth.
{code:java}
step 1:
all node references: 51
AdaptiveSparkPlanExec: 5
TableCacheQueryStageExec: 3
InMemoryTableScanExec: 3
ShuffleQueryStageExec: 3
Union: 2
step 2:
all node references: 147
AdaptiveSparkPlanExec: 15
TableCacheQueryStageExec: 13
InMemoryTableScanExec: 13
ShuffleQueryStageExec: 7
Union: 6
step 4:
all node references: 723
AdaptiveSparkPlanExec: 75
TableCacheQueryStageExec: 73
InMemoryTableScanExec: 73
ShuffleQueryStageExec: 31
Union: 30
step 5:
all node references: 1491
AdaptiveSparkPlanExec: 155
TableCacheQueryStageExec: 153
InMemoryTableScanExec: 153
ShuffleQueryStageExec: 63
Union: 62
step 6:
all node references: 3027
AdaptiveSparkPlanExec: 315
TableCacheQueryStageExec: 313
InMemoryTableScanExec: 313
ShuffleQueryStageExec: 127
Union: 126
{code}
For several node types this is very close to a recurrence of:
{code:java}
N(step + 1) ~= 2 * N(step) + constant
{code}
For example:
{code:java}
TableCacheQueryStageExec:
73
153
313
{code}
and:
{code:java}
Union:
30
62
126
{code}
The traversal time itself also grows:
{code:java}
step 1: 0.065 s
step 2: 0.089 s
step 4: 0.357 s
step 5: 0.626 s
step 6: 1.247 s
{code}
These counts represent recursive physical-plan references reached during
traversal.
They should not be interpreted as proof that every reference corresponds to a
distinct JVM object.
However, they demonstrate that the small outer InMemoryRelation plan hides a
rapidly expanding recursive physical-plan structure involving
AdaptiveSparkPlanExec, TableCacheQueryStageExec, and InMemoryTableScanExec.
The rapid growth correlates strongly with the cache-materialization slowdown.
h2. Other experiments
h3. spark.sql.optimizer.canChangeCachedPlanOutputPartitioning
We tested:
{code:java}
spark.sql.optimizer.canChangeCachedPlanOutputPartitioning=false
{code}
while leaving AQE enabled during state materialization.
This does not resolve the issue.
Planning/execution time still grows rapidly and the production workload takes
more than two hours.
h3. Different JVM GC
We tested ParallelGC using:
{code:java}
-XX:+UseParallelGC
{code}
for driver and executors.
This did not resolve the regression.
h3. localCheckpoint
Periodically calling localCheckpoint() truncates the recursive dependency
structure and resets the growth after the checkpoint boundary.
However, checkpoint materialization adds significant runtime overhead in the
production workload compared with disabling AQE only around cache
materialization.
h2. Expected behavior
We expect Spark 4.1.3 not to exhibit an orders-of-magnitude stronger
cache-materialization degradation than Spark 3.4.3 for the same iterative
cached-DataFrame structure.
In particular, enabling AQE during cache materialization should not change an
iteration from approximately:
{code:java}
5 s
{code}
to:
{code:java}
110+ s
{code}
when disabling AQE only around the same cache() + count() restores the lower
runtime.
h2. Actual behavior
With Spark 4.1.3 and AQE enabled during iterative cache materialization:
* cache materialization becomes rapidly and nonlinearly slower;
* recursive physical-plan traversal exposes rapidly growing references to
AdaptiveSparkPlanExec, TableCacheQueryStageExec, InMemoryTableScanExec, and
nested Union nodes;
* the same workload is dramatically faster when AQE is disabled only during
cached-state materialization;
* the degradation is substantially more aggressive than on Spark 3.4.3;
* in the original workload, lazy transformation planning time also grows from
sub-second timings to hundreds of seconds;
* with the default spark.sql.maxPlanStringLength, the original workload
eventually fails with driver OOM.
h2. Suspected area
The results suggest that the regression is related to the interaction between
AQE and recursively referenced cached InMemoryRelation plans.
In particular, TableCacheQueryStageExec is heavily represented in the
recursively expanded cached physical plans.
We are not claiming that TableCacheQueryStageExec itself is proven to be the
root cause.
However, disabling AQE only while materializing the cached state removes most
of the regression, making the AQE cached-plan path a strong candidate for
investigation.
h2. Related issues
* SPARK-42101 - Wrap InMemoryTableScanExec with QueryStage
* SPARK-45443 - Revisit TableCacheQueryStage to avoid replicated
InMemoryRelation materialization
* SPARK-46995 - Allow AQE coalesce final stage in SQL cached plan
* SPARK-50992 - OOMs and performance issues with AQE in large plans
* apache/spark#55094 - AQE does not apply optimizations for queries involving
TableCacheQueryStageExec
h2. Attachments
* spark_aqe_cache_plan_growth_reproducer.py
Additional sanitized diagnostics can be provided if useful:
* per-iteration timings;
* driver OOM stack trace;
* physical-plan node/reference statistics;
* AQE-on vs scoped-AQE-off comparisons.
--
This message was sent by Atlassian Jira
(v8.20.10#820010)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
