[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?

Before constructing a runtime Bloom filter, return the unchanged application-side plan when the extracted creation-side logical plan contains the PYTHON_UDF tree pattern. 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 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.