[SPARK-58310][SQL] Avoid runtime Bloom-filter subqueries containing Python UDFs

[sunchao]

Why are the changes needed?

With adaptive execution and runtime Bloom filters enabled, a nested broadcast join can cause InjectRuntimeFilter to copy a creation-side plan containing a Python UDF into a newly constructed Bloom-filter scalar subquery. Runtime filters are introduced after subquery optimization, so the copied Python UDF never passes through Python-UDF extraction. The physical plan subsequently attempts to generate JVM code for the unevaluable Python expression while executing BloomFilterAggregate, causing an otherwise valid query to fail through SubqueryExec and ObjectHashAggregateExec.

The problem affects ordinary Apache Spark and does not depend on an external table provider. A self-contained PySpark setup is:

from pyspark.sql.functions import pandas_udf, udf

@udf("string")
def normalize_runtime_bloom_url(value):
    return value

@pandas_udf("string")
def normalize_runtime_bloom_pandas_url(values):
    return values

spark.udf.register("normalize_runtime_bloom_url", normalize_runtime_bloom_url)
spark.udf.register(
    "normalize_runtime_bloom_pandas_url", normalize_runtime_bloom_pandas_url
)

spark.range(20000).selectExpr(
    "cast(id as int) as id",
    "concat('url-', cast(id % 20 as string)) as url",
).write.mode("overwrite").parquet("/tmp/spark-58310-fact")

spark.range(20).selectExpr(
    "concat('url-', cast(id as string)) as url",
    "case when id < 10 then 'NL' else 'US' end as country",
).write.mode("overwrite").parquet("/tmp/spark-58310-dimension")

spark.read.parquet("/tmp/spark-58310-fact").createOrReplaceTempView(
    "runtime_bloom_fact"
)
spark.read.parquet("/tmp/spark-58310-dimension").createOrReplaceTempView(
    "runtime_bloom_dimension"
)

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "false")
spark.conf.set("spark.sql.optimizer.runtime.bloomFilter.enabled", "true")
spark.conf.set(
    "spark.sql.optimizer.runtime.bloomFilter.applicationSideScanSizeThreshold", "0"
)
spark.conf.set("spark.sql.optimizer.runtime.bloomFilter.creationSideThreshold", "100MB")
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")

The following query reproduces the failure; substituting normalize_runtime_bloom_pandas_url reproduces the Arrow/Pandas variant:

WITH coverage AS (
  SELECT url, normalize_runtime_bloom_url(url) AS normalized_url
  FROM runtime_bloom_dimension
  WHERE country = 'NL'
),
normalized_keys AS (
  SELECT DISTINCT normalized_url
  FROM coverage
  WHERE normalized_url IS NOT NULL
),
matched_urls AS (
  SELECT /*+ BROADCAST(k) */ fact.url
  FROM runtime_bloom_fact fact
  JOIN normalized_keys k ON fact.url = k.normalized_url
),
snapshot AS (
  SELECT max(url) AS snapshot_url
  FROM runtime_bloom_fact
)
SELECT /*+ BROADCAST(m), BROADCAST(snapshot) */
  coverage.url,
  coverage.normalized_url,
  matched_urls.url AS matched_url,
  snapshot.snapshot_url
FROM coverage
LEFT JOIN matched_urls ON coverage.normalized_url = matched_urls.url
CROSS JOIN snapshot

What changes were proposed in this PR?

After constructing and column-pruning the runtime Bloom-filter aggregate, return the unchanged application-side plan when the pruned aggregate still contains the PYTHON_UDF tree pattern. Checking after column pruning preserves safe Bloom filters when a Python UDF is needed only by the outer query output and is absent from the actual Bloom-filter scalar subquery. This guard is deliberately limited to unsafe Bloom-filter creation. It does not disable runtime filtering, does not reject a Python UDF on the application side, and does not change Python-UDF extraction or adaptive execution.

Add SPARK-58310 coverage to InjectRuntimeFilterSuite that:

• Runs both standard scalar Python and Arrow/Pandas scalar UDFs when available.
• Reproduces nested CTEs, explicit broadcast joins, scalar aggregates, and adaptive execution with dynamic partition pruning disabled.
• Compares the complete result with the Bloom-disabled baseline and verifies a real Parquet write/read round trip.
• Proves that safe application-side and pruned creation-side runtime Bloom filters are still injected and produce the same results.

Does this PR introduce any user-facing change?

Yes. Queries that previously failed because a runtime Bloom-filter subquery attempted to code-generate a Python UDF now complete successfully. Safe runtime Bloom-filter optimizations remain enabled.

How was this PR tested?

First, the new SPARK-58310 regression was run against unmodified Apache Spark master and failed with the expected Python-UDF scalar-subquery/ObjectHashAggregateExec error. The same regression then passed after the creation-side guard was applied, with both pandas and pyarrow installed.

The following focused upstream regression suites and Scala-style checks were also run on the public Apache Spark branch:

build/sbt \
  'sql/testOnly org.apache.spark.sql.InjectRuntimeFilterSuite org.apache.spark.sql.DynamicPartitionPruningV1SuiteAEOn org.apache.spark.sql.SubquerySuite org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite org.apache.spark.sql.execution.python.ExtractPythonUDFsSuite' \
  'catalyst/scalastyle' \
  'sql/scalastyle' \
  'sql/Test/scalastyle'

How was this patch tested?

The upstream before/after reproduction, Python and Arrow regression, adjacent SQL suites, and Scala-style checks are described above.

Was this patch authored or co-authored using generative AI tooling?

Yes. Generated-by: OpenAI Codex (GPT-5).

[sunchao]

cc @viirya @cloud-fan @dongjoon-hyun @peter-toth

[viirya]

Focused, fail-safe fix, and the P2 refinement in the second commit is the right call — checking the pruned aggregate (post ColumnPruning/ConstantFolding) rather than the raw creation plan means a Python UDF that only feeds an output column, and is pruned out of the actual Bloom subquery, still gets its filter injected. The test pins both sides of that: the crash is gone, and getNumBloomFilters > 0 confirms the safe application-side and creation-side-output-only-UDF cases are still filtered. Coverage spans scalar Python and Arrow/Pandas, with a baseline comparison and a real Parquet round trip.

The guard is consistent with the existing isSimpleExpression treatment of PYTHON_UDF (which only guards the key/predicate expressions) — it closes the complementary gap where the key is simple but the pruned creation-side plan still carries a Python UDF. PYTHON_UDF is the right and sufficient pattern here: only Python UDFs need a separate eval operator that ExtractPythonUDFs inserts; Scala UDFs are directly codegen-able.

I looked at whether this could be fixed at the source (running extraction on the new subquery) instead of skipping the filter, and I don't think it can within a reasonable scope, for a concrete reason worth recording: ExtractPythonUDFs.apply uses plan.transformUpWithPruning, and per QueryPlan's contract the TreeNode transforms "do not traverse into query plans inside subqueries." So the ScalarSubquery(aggregate) that InjectRuntimeFilter builds is invisible to ExtractPythonUDFs no matter where it runs — reordering the batches wouldn't help (InjectRuntimeFilter already precedes the Extract Python UDFs batch), and re-running extraction wouldn't either (it's uncorrelated, so it never gets rewritten into a join and picked up later, unlike the SPARK-26293 case the rule comments describe). Making extraction descend into subqueries would mean changing a core rule's traversal, which is well out of scope for this fix.

One non-blocking thought: the alternative that would preserve the filter is to run ExtractPythonUDFs locally on the aggregate before wrapping it in the subquery, and only bail if a UDF survives. But that injects a BatchEvalPython/ArrowEvalPython on the Bloom creation side, which undercuts the premise that the creation side is cheap — so skipping the filter may well be the better trade anyway. Worth a sentence in the code comment (why the check sits after pruning, and that preserving the filter would require in-subquery extraction) so the next reader doesn't re-derive it; and if you think the source-side fix has value, a follow-up JIRA could capture it. Fine to leave as-is either way.

LGTM.

[uros-b]

LGTM, thank you @suncao and @viirya!

[peter-toth]

Thanks for the PR, @sunchao!

Fail-safe fix: InjectRuntimeFilter builds a Bloom-filter scalar subquery after ExtractPythonUDFs has run, so a Python UDF copied into the creation side never gets its eval operator and crashes codegen; the guard skips the filter in that case. I traced it independently and agree with @viirya -- the key subtlety is that the check sits after ColumnPruning/ConstantFolding (on the built aggregate, not the raw creation plan), so a Python UDF that only feeds an output column and is pruned out of the Bloom subquery still gets its filter (the creationSideQuery test pins this). PYTHON_UDF is the right and sufficient pattern (only Python UDFs need the separate eval operator; Scala UDFs codegen directly), and injectBloomFilter is the only runtime-filter construction path, so nothing else is exposed.

Non-blocking

• 1. Stale PR description: the "What changes" section says the guard returns the unchanged plan "when the extracted creation-side logical plan contains the PYTHON_UDF tree pattern," which describes the pre-refinement behavior -- the final code checks the pruned aggregate, so a prunable creation-side UDF (the creationSideQuery test) still gets a Bloom filter. Worth updating the description to match, alongside @viirya's suggestion to note the after-pruning placement in the code comment.

[dongjoon-hyun]

The test case name sounds a little misleading. Maybe,

Suggested change

	test("SPARK-58310: preserve Python UDFs in runtime bloom-filter scalar subqueries") {
	test("SPARK-58310: skip runtime bloom filter when pruned creation side contains Python UDF") {

[sunchao]

Updated the test name as suggested in 00ed2e05bf. The renamed focused regression passes for both standard Python and Arrow/Pandas UDFs. Thanks!

[dongjoon-hyun]

Please use import java.io.File and revise the following.

Suggested change

	val factPath = new java.io.File(dir, "fact").getCanonicalPath
	val dimensionPath = new java.io.File(dir, "dimension").getCanonicalPath
	val factPath = new File(dir, "fact").getCanonicalPath
	val dimensionPath = new File(dir, "dimension").getCanonicalPath

[sunchao]

Done in 00ed2e05bf: imported java.io.File and updated both paths to use File. Thanks!

[dongjoon-hyun]

+1, LGTM with two additional minor test case comments.

[cloud-fan]

0 blocking, 1 non-blocking, 1 nit.
The production fix and regression coverage are sound; two minor clarity issues already raised by reviewers remain open.

Already raised in existing discussion (2)

• Rename this test to describe that runtime Bloom-filter injection is skipped when the pruned creation side retains a Python UDF. The current name says the UDF is preserved in the scalar subquery, but the production guard prevents that scalar subquery from being created. -- existing discussion
• Update the PR description to say the guard checks the pruned Bloom aggregate, not the extracted creation-side plan. The distinction matters because output-only Python UDFs are pruned and still permit Bloom-filter injection. -- existing discussion

Verification

Traced InjectRuntimeFilter through post-pruning aggregate construction, scalar-subquery wrapping, and the optimizer/Python-UDF extraction ordering. Both selected mechanical scanners completed with no findings. Tests were not run as part of this review.

PR description suggestions

• Clarify that the guard checks the pruned Bloom aggregate, so output-only Python UDFs that column pruning removes do not disable the optimization.

[sunchao]

Merge Summary:

• merged into master 6f1c904

Posted by merge_spark_pr.py

[sunchao]

Merge Summary:

• merged into branch-4.x 5ef5af3

Posted by merge_spark_pr.py

[sunchao]

Thanks all for the review! Merged to master / branch-4.x