Question
Adaptive Query Execution (AQE) and Nested UDF Cost Estimation
Kavi River
0 reputation · 06 Apr 2020, 12:35 UTC
148.2K views1
Loading…
0 reputation · 06 Apr 2020, 12:35 UTC
Apache Spark's Adaptive Query Execution (AQE) optimizes query plans at runtime by utilizing statistics from completed shuffle stages. This allows the engine to dynamically coalesce shuffle partitions and switch join strategies, such as converting a SortMergeJoin to a BroadcastHashJoin when actual relation sizes fall below the threshold.
However, the integration of complex nested User Defined Functions (UDFs) introduces uncertainty into the optimizer's cost estimation. Because UDFs are treated as black boxes, the runtime statistics collected during the shuffle phase may not accurately reflect the computational overhead or the data distribution changes caused by nested logic.
In Spark 3.x, how does the AQE optimizer handle cost estimation for stages where nested UDFs precede the shuffle? Does the framework provide a mechanism to hint the actual cost of these UDFs to prevent suboptimal join strategy selection?
34775 reputation · 06 Apr 2020, 13:45 UTC
Spark 3.x Adaptive Query Execution (AQE) does not provide a mechanism to hint the computational cost or cardinality effects of nested UDFs. The optimizer treats UDFs as opaque functions and re-plans join strategies solely from runtime shuffle map-output statistics (bytes and records per partition) after the shuffle stage completes.
apachesparkspark.sql.adaptive.enabled=true, AQE can coalesce shuffle partitions, convert SortMergeJoin to BroadcastHashJoin when the observed relation size falls below spark.sql.adaptive.autoBroadcastJoinThreshold, and mitigate skew using actual partition sizes.BROADCAST, MERGE, SHUFFLE_HASH, SHUFFLE_REPLICATE_NL) force a strategy but do not teach the optimizer the UDF's cost; they bypass cost-based adaptation entirely.UDFs preceding a shuffle can alter row count, row width, and value distribution—factors that directly affect join strategy choice. Because Catalyst cannot analyze nested UDF bytecode or Python logic, it assumes the UDF is a simple projection with 1:1 cardinality and negligible CPU. AQE compensates by reacting to the actual shuffle output, but it cannot anticipate high CPU time hidden behind small output bytes (common with Python UDFs or complex nested Scala UDFs).
spark.conf.get("spark.sql.adaptive.enabled") and spark.conf.get("spark.sql.adaptive.autoBroadcastJoinThreshold").EXPLAIN FORMATTED with AQE on and off; confirm an AdaptiveSparkPlan node appears and note the join type chosen.df.hint("broadcast") or /*+ BROADCAST(t) */) after validating on representative data that the broadcast side fits in memory.df.persist()) so the shuffle statistics reflect a stable result.What is the observed shuffle write size (bytes/records per partition) for the stage immediately after the nested UDF? This determines whether AQE's broadcast conversion is triggering correctly or if the UDF is producing deceptively small output that hides high CPU cost.
Use comments to ask for clarification. Post a solution as an answer.
5,560 reputation · 06 Apr 2020, 13:07 UTC
AQE re-plans using post-shuffle map-output statistics such as rows, bytes, and partition sizes; it does not decompose map-side work into per-UDF cost components. In Spark 3.x, Catalyst treats nested UDFs as opaque expressions, and there is no configuration to inject an explicit UDF cost multiplier or cardinality adjustment before the shuffle. Consequently, a join strategy may still reflect pre-UDF estimates even when the UDF expands or filters rows. Verify by running EXPLAIN with spark.sql.adaptive.enabled toggled and comparing shuffle metrics against spark.sql.adaptive.autoBroadcastJoinThreshold. A BROADCAST join hint can override AQE, but confirm executor memory because the post-UDF relation may exceed the broadcast limit.