Skip to content

fix: support struct-typed scalar subquery results - #5889

Draft
LinSimon-901101 wants to merge 1 commit into
apache:mainfrom
LinSimon-901101:fix/5834-struct-scalar-subquery
Draft

LinSimon-901101 wants to merge 1 commit into
apache:mainfrom
LinSimon-901101:fix/5834-struct-scalar-subquery

Conversation

@LinSimon-901101

@LinSimon-901101 LinSimon-901101 commented Sep 12, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5834.
Closes #3336.

Rationale for this change

Spark merges related scalar subqueries into struct results. Rejecting those results forces consuming operators to fall back to Spark. Spark 4.2's MergeSubplans also introduces this shape in aggregate and UNION queries without explicit subqueries.

This change allows supported struct results to stay native while preserving batch sizes, null semantics, and Spark fallback for consumers whose calendar or timezone behavior cannot cover the full input range.

What changes are included in this PR?

  • Support non-empty struct scalar-subquery results with distinct field names and supported scalar or nested-struct fields. Transfer materialized results through Arrow IPC and resolve them into owned literals per physical expression occurrence during native planning, preserving field order, planned nested metadata, and the distinction between a null struct and a struct containing null fields.
  • Make GetStructField over a scalar struct return an array of batch.num_rows(), including empty batches. Downstream consumers receive the same shape as a column, including when every row selects the same IF branch.
  • Register subqueries for each per-partition Top-K iterator before native physical planning and remove registrations at task completion. Retain regression coverage for the existing codegen-dispatch rejection of embedded subqueries and use a regular Spark sort for the reference results affected by Spark 3.4's Top-K limitation.
  • Handle scalar timestamp arguments in hour, minute, second, and unix_timestamp. Re-enable the four Native engine panics on all-scalar inputs for hour, minute, second, unix_timestamp #3336 SQL cases as native expectations with constant folding disabled. Fixed-offset clock extraction handles full-range timestamp scalars, arrays, and dictionaries; timestamp-without-time-zone clock fields remain independent of the session timezone.
  • Replace the calendar-consumer blacklist with an explicit allow-list on paths enabled by struct subqueries. Raw transport, comparisons, supported epoch conversions, fixed-offset/NTZ clock extraction, DayOfWeek, and WeekDay remain native. Unknown calendar consumers fall back, including TruncDate and consumers returning non-calendar types. Allowed wrappers still check their children, and the guard follows aliases, generators, and sibling/downstream consumers, including after a projection drops the struct.
  • Extend the lpad/rpad string-argument fallback from literals to expressions without column references (str.references.isEmpty), including direct scalar subqueries. Conditional inputs with column references can remain native. A small adapter reuses the existing padding kernel for runtime scalar strings from conditional branches or serializer rewrites, while preserving NULL lengths, the literal-pad restriction, and Spark's NULL short-circuit behavior.
  • Add SQL, Scala, and Rust regressions for plain and merged subqueries, batch sizes 0/1/many, parent and nested nulls, metadata, conditional padding, ANSI errors, full-range timestamps, calendar fallback, aliases, generators, and floating-point special values. Update the q9/q77a approved plans and narrow the Spark 4.2 merged-UNION limitation to duplicate-field-name structs.

Unsupported struct shapes, duplicate field names, and unsupported consumers retain fallback. The existing malformed-UTF-8 normalization limitation also applies to newly native struct results and remains documented.

CI execution settings match integrated upstream 80b518173d0631d31028b732b385c68f69ace326. The only workflow differences are the two required CometScalarSubquerySuite registrations in the Linux and macOS test lists. There are no PR-specific heap, timeout, watchdog, retry, matrix, or failure-rule changes.

How are these changes tested?

[Full fork CI] passed at 6d0d35f:

  • 150 jobs: 137 successful, 13 conditional skips, zero failed or cancelled jobs.
  • All 45 Spark SQL test shards passed across Spark 3.4, 3.5, 4.0, 4.1, and 4.2.
  • Spark 3.4 sql_core-3: 1,468 ScalaTest cases passed, zero failed, and zero aborted suites.
  • Spark 3.5 Catalyst: 7,210 ScalaTest cases passed, zero failed, and zero aborted suites, including the previously failing HLL case.
  • Spark 4.2 sql_core-1: 10,298 ScalaTest cases passed, zero failed, and zero aborted suites.
  • Preflight and Required Checks passed. All 19 Required Checks dependencies reported success.

@github-actions github-actions Bot added bug Something isn't working area:expressions Expression evaluation area:ffi Arrow FFI / JNI boundary labels Sep 12, 2026
@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

Local validation at 918202206 on macOS arm64:

  • Spark 4.2.0 / JDK 21: 3 focused tests passed, 0 failed.
  • Rust scalar-subquery tests: 5 passed, 0 failed.
  • Native build, Scalastyle and Spotless passed.

The Spark 4.1.3 logs show the pre-merge red/green comparison: the same two regressions failed against the original base due to unsupported struct results/projection fallback, then passed with the fix. The baseline failures are expected.

Logs are attached below.
spark-4.2-focused.txt
rust-struct-subquery.txt
spark-4.1-baseline-expected failures.txt
spark-4.1-fixed-before upstream merge.txt

@LinSimon-901101
LinSimon-901101 marked this pull request as ready for review September 12, 2026 20:08

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 91820220 against de1eb4f8. No verified P1/P2 findings.

Correctness

Previously, the struct returned by Spark's merged scalar subqueries caused the consuming projection to fall back. The recursive support check now admits non-empty structs with distinct field names and supported leaves. Array/map/interval fields, duplicate names, non-default collations and unsupported decimal scales retain fallback.

The one-row, one-column IPC envelope preserves the distinction between a null struct and a present struct containing null fields. Native decoding validates shape, field order, names, nullability and leaf types before restoring planned metadata. Timestamp fields retain UTC microseconds, while timestamp-without-time-zone fields remain unzoned. Spark still enforces scalar cardinality before native evaluation. Each task iterator builds a fresh native plan, so the owned cache stays within that execution; cache contents do not affect expression equality or hashing.

The SQL fixture covers supported leaves, nested nulls, empty results, separate subqueries and recursive fallback. The optimizer tests verify actual struct subqueries and a native consuming projection, including the Spark 4.2 union/aggregate case. I inspected the author's attached logs: they report three Spark 4.2 tests and five Rust tests passing at the stated head, plus the pre-merge Spark 4.1 red/green comparison. These are author-provided results, not independent CI. CI remains action_required with zero jobs; the sole successful check is labeling. I ran no local JVM/native tests. Canonical source checks covered the maintained Spark 3.5/4.0 branches; the required 3.4/4.1 branches were unavailable.

Performance

The first evaluation pays for row normalization, Arrow allocation, IPC serialization and decoding. Subsequent evaluations reuse the owned scalar, avoiding repeated IPC work per batch for that expression. Repeated references can still own separate caches, so this does not establish one transfer per logical subquery. Primitive transfer paths remain unchanged apart from the empty-cache lookup. No benchmark was run and no measured speedup is established. Could you add a focused comparison against the previous Spark fallback, varying struct width and input-batch count, to quantify first-transfer cost and how well the cache amortizes it?

Design

Reusing the existing Arrow writer keeps value transfer aligned with the project's supported primitive representations. The Scala support gate and native wire validation serve separate purposes: the former decides fallback before execution, while the latter rejects a malformed or mismatched result. Metadata restoration checks the value shape rather than introducing casts. The change addresses the struct-transfer seam and explicitly retains the existing malformed-UTF-8 limitation and separate scan/shuffle limitations.

Abstraction & complexity

The implementation adds a focused serializer, decoder and execution-local cache without a new plan protocol or global cache. The explicit equality/hash implementations preserve the original three identity fields while excluding mutable state. Recursive normalization and metadata alignment are limited to the supported struct shape; the added tests exercise those responsibilities directly. I found no additional abstraction or simplification that warrants a pre-merge change.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up at unchanged 91820220 / base de1eb4f8:

[P2] Regenerate q9’s approved plan for the newly native struct-subquery projection. All five exec jobs fail only check simplified (tpcds-v1.4/q9): 3.4, 3.5, 4.0, 4.1, 4.2. The unchanged golden expects Project [COMET: Unsupported data type: StructType(...)]; each actual plan has CometColumnarToRow → CometProject and reports 38 accelerated operators instead of 37. Please regenerate and commit the affected approved plan using the documented golden-file workflow, then rerun the supported-profile plan-stability checks.

These are actual plan-comparison failures at the exact tested H/B merge. The native producer and JVM consumers share artifact 10307852860 and its verified digest. The merged-struct result/operator test and scalar_subquery_struct.sql pass on all five profiles; the five Rust struct tests also pass. This supports an expected-plan omission, without demonstrating a query-result regression. I replayed the captured plan comparison locally; I did not run Spark/Rust tests or benchmarks locally. Other CI jobs remain pending.

@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

Thanks for the review; I've regenerated q9's shared golden file to reflect the native CometProject, with no production code or assertion changes.
Both complete TPC-DS plan-stability suites passed locally on all five Spark profiles (3.4–4.2) with golden generation disabled: 129 tests per version, 0 failures.
Spark 3.4 used JDK 17 locally rather than CI's JDK 11, and CI validation of the update remains pending.

Comment on lines 37 to 54
private def supportedStructField(dt: DataType): Boolean = dt match {
case s: StructType =>
s.nonEmpty && s.fieldNames.distinct.length == s.length &&
s.fields.forall(f => supportedStructField(f.dataType))
case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
StringType | BinaryType | DateType | TimestampType | TimestampNTZType | NullType =>
true
case d: DecimalType => d.scale >= 0 && d.scale <= d.precision
case _ => false
}

override def getSupportLevel(expr: ScalarSubquery): SupportLevel = {
val supported = expr.dataType match {
case s: StructType => supportedStructField(s)
case dt => supportedDataType(dt)
}
if (supported) {
Compatible()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Related to #5025, which centralizes these recursive type checks. Could we coordinate this with that helper instead of adding another independent predicate? It already covers duplicate names, string types, and interval restrictions. This path would still need its struct-only and decimal-scale restrictions preserved.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. I've adopted the shared supportedDataType API and implementation proposed in #5025, while retaining the scalar-subquery bridge's struct-only and decimal-scale restrictions.

Code: QueryPlanSerde.scala:579–614, with bridge-specific restrictions in CometScalarSubquery.scala:35–55.

Comment on lines +283 to +285
// Concurrent first evaluations may both initialize the same immutable result.
// Failed evaluations are never cached.
let _ = self.struct_value.set(value.clone());

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can the same native Subquery instance be evaluated concurrently here? Separate Spark tasks have separate instances. Understanding the sharing would help decide whether lazy initialization is needed.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each Spark task creates its own native execution context, and struct results are now resolved into owned, immutable Literals during native physical planning. I've removed OnceLock and evaluation-time initialization, so the implementation no longer depends on assumptions about concurrent evaluation.

Code: task-local iterator creation and subquery registration in CometExecRDD.scala:125–139, and immutable literal construction in planner.rs:786–798.

Comment on lines +160 to +162
if let Some(value) = self.struct_value.get() {
return Ok(ColumnarValue::Scalar(value.clone()));
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By native physical planning, Spark has materialized and registered the subquery result. Could we resolve it into an immutable scalar there? What requires deferring initialization until evaluate?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since Spark completes the materialization and registration of results before native physical planning, there's no reason to delay until evaluate. It's now resolved to owned immutable Literals during planning, and OnceLock has been removed.

Code: planning-time resolution in planner.rs:786–798, using Subquery::resolve_struct in subquery.rs:60–86.

if reader.next().transpose()?.is_some() {
return internal_err!("Scalar subquery IPC result contains more than one batch");
}
let value = align_struct_metadata(batch.column(0), data_type)?;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Rust test covers restoring metadata. Could we also exercise a struct with Parquet field IDs through the JVM serializer, showing the schema mismatch that occurs without this alignment?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. I've added a JVM serializer test to prove that the Arrow IPC wire schema loses nested Parquet field IDs, while the planned protobuf type retains metadata. A native integration test also verifies that the output after alignment restores nested IDs and covers NULL children and NULL structs.

Tests: JVM serializer coverage in CometArrowStreamSuite.scala:65–140, and native-output metadata/NULL coverage in CometExecSuite.scala:2375–2480.


/// The JVM bridge emits one row with one struct column. Validate the wire shape and type before
/// creating the scalar; Arrow IPC validation also keeps malformed strings out of native arrays.
fn decode_struct_result(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The JVM writer always emits one row, one column, and one batch. Can these shape checks fail through a supported production path? If not, could we simplify them?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These shape checks cannot fail on the current supported production path, since the JVM serializer always emits one row, one column, and one batch. I've removed the extra-batch probe and retained the no-batch and row/column guards to prevent panics if an internal bug violates that contract.

Code: the fixed producer shape in CometArrowConverters.scala:62–76, and the simplified decoder with retained defensive guards in subquery.rs:89–105.

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think GetStructField over the new scalar struct only works when the consuming batch has exactly one row.

Subquery::evaluate returns ColumnarValue::Scalar(ScalarValue::Struct(..)) holding a one-row StructArray. The scalar arm of GetStructField::evaluate handles that by returning ColumnarValue::Array(child_with_parent_nulls(&struct_array, self.ordinal)?), and child_with_parent_nulls hands back the child of that one-row struct, so the result is a length-1 Array rather than a scalar. DataFusion then runs it through ColumnarValue::into_array_of_size(batch.num_rows()), which rejects an Array whose length is not the batch row count with Array length 1 does not match expected length N.

So SELECT c1, (SELECT max(x) FROM t2), (SELECT min(x) FROM t2) FROM t1 should fail as soon as t1 produces a batch with more than one row, and that is the main shape MergeScalarSubqueries creates. Would it make sense for the scalar arm of GetStructField::evaluate to return ColumnarValue::Scalar(ScalarValue::try_from_array(&child, 0)?) so the projection broadcasts it instead?

I think the reason the tests do not catch this is that all of them land on one-row batches. CometTestBase runs local[5], so the 4-row INSERT ... VALUES in scalar_subquery_struct.sql becomes four one-row Parquet files and withParquetTable((0 until 5)...) becomes five one-row files. The Spark 4.2 test's projection sits directly over one-row aggregates, and q9 filters reason down to r_reason_sk = 1. Could you add a fixture where the table is written as a single file so the consuming projection sees a batch with many rows, and check that it fails without the GetStructField change?

@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

Thanks for pointing out the issue; the scalar branch of GetStructField now returns ColumnarValue::Scalar, allowing DataFusion to broadcast the value to the consuming batch's row count. The regression test now uses a single Parquet file with 1,024 rows and fails without the fix, with additional Rust coverage for batch sizes 0/1/4, NULLs, and nested structs.

Code: get_struct_field.rs:102–108. Tests: the 1,024-row regression in CometExecSuite.scala:2300, and batch-size/NULL/nested-struct coverage in get_struct_field.rs:148–244.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed fb658a69 against 0690d38d.

The previous [P2] q9 golden mismatch is fixed. The updated golden exactly matches the actual plan captured from each of the five previously failing exec jobs. Replaying the unchanged text comparison fails with the old golden and passes with the new one. This reuses captured CI plans, rather than generating a new plan at this head.

The scalar GetStructField branch now returns a scalar for broadcasting. The regression uses one 1,024-row Parquet file and asserts multi-row batches with AQE off/on. I also checked the planning-time immutable literals, registration order, shared type-gate defaults and nested field-metadata coverage. These address the discussed implementation issues without a new verified P1/P2 finding.

Current-head CI is still action_required with zero jobs. The sole successful check is labeling. I ran the captured-plan replay and source checks, but no local Spark/Rust tests or benchmarks. The new runtime tests remain unverified independently, and the earlier benchmark request remains unmeasured.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed 6c5ed5ff against fb658a69 and base 58ab5f61. The merge retains the same 16-file authored contribution. The conflict resolution keeps both the struct metadata regression and the inherited dictionary-stream test. Scalar-field broadcasting, nested null handling, the 1,024-row regression and subquery registration before native planning remain intact. No new or remaining verified P1/P2 findings.

The q9 golden and comparator are unchanged. Replaying the five historical captured plans still fails with the original golden and passes with the current one. This is a text-comparison replay, with no fresh plan generation. Source checks passed against maintained Spark 3.5/4.0. No local product tests or benchmarks were rerun, and the earlier benchmark request remains unmeasured. Current CI and CodeQL require approval and have zero jobs. The successful label workflow checked out base/main.

@andygrove

Copy link
Copy Markdown
Member

I would like to do a deeper review of this PR before it is merged.

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I built this merged with current main and ran it against about 35 query shapes that read the struct natively, including filters, aggregates, windows, native shuffle partitioning, coalesce/IF/CASE against struct columns, struct equality, hashing, casts, -0.0 and NaN fields, timestamps under a non-UTC session, and decimal(38). They all match Spark on 3.5, 4.1 and 4.2, and plan stability only changes q9. Two shapes that work today fail with this change, though.

The first is a merged subquery feeding a Scala UDF. Spark wraps a UDF that takes a primitive in a null check, so SELECT plus1((SELECT max(a) FROM t)), plus1((SELECT min(a) FROM t)) FROM t becomes if(isnull(m.max), null, plus1(m.max)) once MergeScalarSubqueries runs. On main and in 1.1 the isnull side can't convert the struct subquery, so the projection falls back and returns Spark's answer. With this PR the isnull side converts, the UDF goes to the codegen dispatcher with the subquery inside it, and the task fails with requirement failed: Subquery ... has not finished. The dispatcher closure-serializes its tree at planning time, before the subquery has run, so a plain subquery inside a UDF already fails on main. I filed #6674 for that. Adding case _: ExecSubqueryExpression => true to the rejection list in CometBatchKernelCodegen.canHandle makes these fall back correctly, and CometCodegenSuite still passes.

The second is a merged or struct subquery in the sort keys of a top-K over more than one partition, for example ORDER BY a + (SELECT max(a) FROM t), b + (SELECT min(b) FROM t) LIMIT 7. CometTakeOrderedAndProjectExec builds its per-partition top-K plan without registering subqueries. Now that struct values are resolved while that plan is built, it fails with Subquery N not found for plan M where main falls back. A plain subquery in the sort keys already hits this on main, and I filed #6676 for it. Registering the subqueries on that iterator the way the final plan already does fixes it in my local run.

Could we include both fixes here, or land them first, and add the two queries as tests? Neither root cause is yours, but this PR is what moves these shapes from a working fallback onto the broken paths.

The branch also conflicts with main now. subquery.rs moved to native/spark-expr/src/subquery.rs in #6450, so the new code goes there, with CometError coming from datafusion_comet_jni_bridge::errors. With just those changes it builds and the tests pass for me. I also added run-spark-4.1-tests and run-spark-4.2-tests, since Spark's SubquerySuite has five merged-subquery tests that inspect the plan. They will run on your next push.


This list is illustrative, not exhaustive: the per-function tables are not the complete set of expressions Comet can accelerate.

Scalar subqueries can return structs, including those created when Spark merges multiple scalar subqueries. Struct results are transferred from Spark through Arrow IPC during native physical planning and retained as owned, immutable literals for that plan's execution. Supported fields include booleans, numeric types, default-collation strings, binary, dates, timestamps, nulls, and nested structs. Decimal fields require a non-negative scale no greater than their precision. Structs must be non-empty and have distinct field names at each level; arrays, maps, intervals, and other unsupported field types still cause fallback to Spark. Existing non-struct scalar-subquery paths are unchanged.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Most of this paragraph describes the implementation rather than what a user sees. The Arrow IPC and owned-literal details don't give a user anything to act on, and "Existing non-struct scalar-subquery paths are unchanged" only makes sense next to this diff. One thing a user does need is missing. A string field that holds invalid UTF-8 now comes back with U+FFFD replacement where these queries used to fall back, so hex(cast(s.v AS binary)) returns EFBFBD28 where Spark returns C328. Could we trim this to what is supported and what falls back, and point to the "Strings with non-UTF-8 bytes" section of the compatibility guide? The list of boundaries in that section could name scalar subquery results too.

DecimalType(38, 38))

// Freeze the pre-refactor shared predicate: existing callers must retain their accepted types.
private def legacySupported(dt: DataType, allowComplex: Boolean): Boolean = dt match {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

legacySupported and the first two tests check QueryPlanSerde.supportedDataType itself against a copy of the predicate from before #5025. Now that #5025 has merged with its own boundary test in QueryPlanSerdeSuite, these duplicate that coverage. The frozen copy also means a legitimate later change to the shared helper would fail a suite named for scalar subqueries. Could we drop them and keep the scalar-subquery checks, with the non-struct case comparing supported(dt) to supportedDataType(dt) directly?

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

  • Prior state and problem: Spark’s merged scalar subqueries return structs, forcing consuming projections to fall back. Spark 4.2’s MergeSubplans exposes this in additional query shapes.
  • Design approach: Admit supported struct shapes recursively, transfer their materialized values through Arrow IPC, and resolve them into owned native literals during physical planning.
  • Correctness / compatibility analysis: Null structs remain distinct from structs containing null fields. Field order, planned metadata, timestamp representations, and scalar broadcasting are preserved. Relevant Spark implementations were checked across 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Two already-reported P2 regressions remain unresolved: merged subqueries inside Scala UDFs fail with Subquery ... has not finished (#6674), and merged subqueries in multi-partition top-K sort keys fail with Subquery ... not found (#6676). The reported reproductions are corroborated by the unchanged faulty paths in this checkout. These shapes previously fell back successfully.
  • Key design decisions: Immutable literals avoid mutable initialization state and repeated JNI lookup during batch evaluation. IPC serialization and decoding still occur per physical expression occurrence. No measured performance improvement is established.
  • Implementation sketch: The serde gate selects supported types, the JVM helper serializes one struct value, native decoding restores planned field metadata, and scalar GetStructField results broadcast to the consuming batch. This uses existing Arrow infrastructure without adding a global cache.
  • Behavioral changes worth calling out: Supported struct consumers now execute natively, including q9’s projection. Unsupported nested types retain fallback. The previously reported malformed-UTF-8 normalization caveat also remains.
  • Suggested improvements: Address the existing blockers by rejecting embedded ExecSubqueryExpression trees in codegen dispatch and registering and cleaning subqueries for the per-partition top-K iterator, with regression coverage for both reported queries.

Reviewed the entire 15-file diff at fd94f0a6d95b939914afd766b203cad9d89127c1 against 76493618769ff85e7a43e6f6c79f8b58d9e0b1ad. Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, and review-comet-memory-pr.

Exact-head CI: Comet CI and CodeQL require approval and have zero jobs. Only labeling passed. There is no exact-head Spark SQL suite verdict.

Validation: Six struct IPC/literal Rust tests and seven field-extraction Rust tests passed locally. Core tests used --no-default-features, excluding HDFS. No local JVM integration tests, full Spark SQL suites, or benchmarks were run. The checkout remains unchanged.

No additional introduced P1/P2 issues found within this review. Existing blockers remain, so this is not a clean verdict.

@LinSimon-901101

LinSimon-901101 commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor Author

@andygrove ,Thanks for pointing out these issues. I've addressed the review feedback, and the targeted local tests have passed, including the Spark 4.1 and 4.2 subquery tests. The full CI matrix is now running on my fork. I'll push the changes to this PR once CI passes.

@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

A quick update: I've pushed the review fixes and resolved the merge conflict with main. The test adjustment for Spark 3.4 remains on my fork pending CI validation.

The previous fork CI run exposed six failures on Spark 3.4.3 while collecting the Spark reference results with Comet disabled. I reproduced this locally; the failure matches [SPARK-45584], where Spark's TopK does not wait for subquery execution. This is separate from the Comet subquery-registration issue in #6676.

The test adjustment uses a regular Spark sort for the reference results, while retaining Comet TopK execution and its plan assertions. Removing the Comet registration fix still makes all six cases fail, confirming that the regression coverage remains effective.

After merging the latest main, all six affected Spark 3.4 cases pass locally. The Spark 4.1 codegen, native-UDF nesting, and TopK checks also passed (116 tests).

I've triggered a [new full CI run on my fork], including both the merge resolution and the test adjustment. I'll add the test adjustment to this PR once CI passes.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

  • Prior state and problem: Spark’s merged scalar subqueries return structs, forcing consuming operators to fall back. Spark 4.2’s MergeSubplans exposes this in additional query shapes.
  • Design approach: Admit supported struct shapes, transfer materialized values through Arrow IPC, and resolve them into owned native literals during physical planning.
  • Correctness / compatibility analysis: Reviewed null-struct semantics, scalar broadcasting, field order, metadata restoration, cardinality enforcement, and timestamp representations against Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0 sources. The previously reported UDF and top-K registration defects are addressed in the implementation. One already-discussed P2 test blocker remains, described below.
  • Key design decisions: Immutable literals avoid mutable initialization and repeated JNI lookup during batch evaluation. Serialization and decoding still occur per physical expression occurrence. JVM Arrow allocations and native literal storage sit outside pool reservations. No benchmark or measured performance regression is established.
  • Implementation sketch: The shared recursive type gate retains bridge-specific restrictions. The JVM serializes one struct value, Rust validates and restores planned metadata, and GetStructField returns a scalar that broadcasts correctly. Top-K registration now mirrors the existing final-stage lifecycle.
  • Behavioral changes worth calling out: Compared with branch-1.1, supported struct consumers intentionally become native, including q9’s projection. Embedded subqueries now prevent codegen dispatch. Unsupported struct shapes retain fallback. The previously discussed malformed-UTF-8 normalization difference is documented.
  • Suggested improvements: Include the already-prepared Spark 3.4 reference-answer correction. At CometTopKSuite.scala:174, checkSparkAnswerAndOperator collects the Spark reference through its affected top-K implementation. All six new cases failed with Subquery ... has not finished in the predecessor CI job. These tests and their comparison helper are unchanged at this head, and Spark 3.4.3 source confirms the missing subquery wait. The author’s acknowledged correction uses regular sorting for the Spark reference while retaining Comet top-K assertions. It remains outside this PR head.

Reviewed the entire 20-file diff at 8e4c6ca77eed0415ca7756b687a2dcb115cb9785 against b80bf4e08198d992450c50ec85cd804c72095246. Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, and review-comet-memory-pr.

Exact-head CI: Comet CI and CodeQL require approval and have zero jobs. Only labeling passed. The successful fork run uses 324ce0370136a6461b6bf7803c0dfc37d94af498, which includes the test correction, so it is not an exact-head verdict.

Validation: All 13 focused Rust tests passed locally: six struct IPC/literal tests and seven field-extraction tests. No local JVM integration build, full Spark SQL suites, or benchmarks were run. The checkout is unchanged.

No additional introduced P1/P2 issues found within this review. The substantiated existing Spark 3.4 test blocker remains.

@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

The [fork CI run] completed successfully, including the Spark 3.4 test adjustment. The [subsequent full CI run] , which also includes the merge-conflict resolution, passed as well.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

  • Prior state and problem: Spark’s merged scalar subqueries return structs, forcing consuming operators to fall back. Spark 4.2’s MergeSubplans exposes this in additional query shapes.
  • Design approach: Admit supported struct shapes, transfer materialized results through Arrow IPC, and resolve them into owned native literals during physical planning.
  • Correctness / compatibility analysis: Checked null structs versus null fields, scalar broadcasting, metadata restoration, timestamps, cardinality enforcement, and registration lifetime against Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0 sources. The previously reported broadcasting, codegen, top-K registration, q9 golden, and Spark 3.4 reference-test blockers are addressed.
  • Key design decisions: Immutable literals avoid mutable initialization and repeated JNI lookup during evaluation. Serialization and decoding occur per physical expression occurrence. JVM Arrow buffers and native literals remain outside pool reservations. No benchmark establishes a performance improvement or regression.
  • Implementation sketch: The shared type gate retains bridge-specific restrictions. The JVM serializes one struct value, Rust validates its shape and restores planned metadata, and GetStructField returns a broadcastable scalar. Top-K registration follows the existing final-stage lifecycle without introducing another abstraction.
  • Behavioral changes worth calling out: Compared with branch-1.1, supported struct consumers intentionally become native. Unsupported shapes retain fallback, and embedded subqueries now prevent codegen dispatch. The already-discussed malformed-UTF-8 normalization difference is documented for the new path.
  • Suggested improvements: No additional P1/P2 changes requested.

Reviewed the entire 20-file diff at afd99bd5e2a1aa149308d73d4d23a74de1d651ab against 9dc8c3ca89962506078831f7064a25c75da83883, including all prerequisite commits. The PR remains non-draft. Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, and review-comet-memory-pr.

Exact-head CI: Comet CI and CodeQL require approval and have zero jobs. Only labeling passed. The successful fork run tested ancestor 324ce0370136a6461b6bf7803c0dfc37d94af498. Its logs confirm all six Spark 3.4 top-K regressions, q9, the struct SQL fixture, and both UDF fallback cases passed. This is supporting evidence, not an exact-head CI verdict.

Validation: All 13 focused Rust tests passed locally: six struct IPC/literal tests and seven field-extraction tests. No local JVM integration build, full Spark SQL suites, or benchmarks were run. The checkout is unchanged.

No introduced P1/P2 issues found within this review. No substantiated existing P1/P2 blockers remain.

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for working through the last round. I merged this with current main and checked both fixes. The per-partition top-K now registers its subqueries, and the six new CometTopKSuite tests fail without that change. The UDF case falls back through the guard that landed in #6666, and the new CometCodegenSuite tests pin it. The PR's tests and TPC-DS plan stability pass on 3.4 through 4.2, and about a hundred query shapes that read the struct natively match Spark.

I found one more shape that works today and fails with this change, and it follows the same pattern as last time. Native hour, minute and second (which date_part and extract use), unix_timestamp, and lpad/rpad reject a scalar argument. SELECT a, hour((SELECT max(ts) FROM t)), hour((SELECT min(ts) FROM t)) FROM t fails with hour(scalar) should be fold in Spark JVM side. On main the two subqueries merge into a struct, the projection falls back, and the answer is right. With this PR the field read returns a scalar, so it reaches the kernel. lpad((SELECT max(s) FROM t), 8, '0') next to a second subquery on t fails the same way. A single plain subquery already hits this on main and in 1.1, which is #3336. Could we land a fix that lets these kernels take a scalar first, or include it here, and add one of these queries as a test? Wrapping the scalar in a one-row array and converting the result back fixed all of them for me, merged and plain.

compatibility/spark-versions.md still lists "Merged subqueries in UNION branches" as a Spark 4.2 limitation that sends q77a's projection, union and aggregates back to Spark. With this PR, q77a runs 87 of its 89 operators in Comet. Could we drop that bullet, or narrow it to the case that still falls back? That is a merged struct whose field names repeat, such as SELECT sum(s) FROM (SELECT max(a) AS s FROM t UNION ALL SELECT min(b) AS s FROM t).

-- specific language governing permissions and limitations
-- under the License.

-- Config: spark.sql.session.timeZone=UTC

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since timestamps in these structs now cross into native code, could this run under -- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles, with a query that casts or compares s.instant rather than only projecting it? Both pass with the current code when I try them.

val schema = StructType(Seq(StructField("value", dataType, nullable = true)))
val output = new ByteArrayOutputStream()
Using.resource(
VectorSchemaRoot.create(Utils.toArrowSchema(schema, "UTC"), CometArrowAllocator)) { root =>

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this use CometArrowStream.NATIVE_TIMEZONE rather than the literal "UTC"? timezones.md says every JVM producer uses that constant.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

  • Prior state and problem: Spark’s merged scalar subqueries return structs that previously forced consuming operators to fall back. Spark 4.2’s MergeSubplans exposes this in additional aggregate and union queries.
  • Design approach: Admit supported struct shapes recursively, serialize materialized results through Arrow IPC, and resolve them into owned native literals during physical planning.
  • Correctness: Checked cardinality enforcement, null structs versus null fields, field extraction, metadata restoration, and registration order against Spark sources. The existing P2 scalar-consumer blocker remains. For example, SELECT a, hour((SELECT max(ts) FROM t)), hour((SELECT min(ts) FROM t)) FROM t previously fell back successfully but now reaches a kernel that rejects scalar arguments. Local reproduction confirmed rejection by hour, minute, second, unix_timestamp, lpad, and rpad after struct-field extraction. Equivalent one-row array inputs succeed.
  • Compatibility analysis: Compared relevant implementations across Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Timestamp fields retain UTC microseconds, while TimestampNTZType remains unzoned. Unsupported nested types, duplicate names, collations, and decimal scales retain fallback. Compared with branch-1.1 at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1, native struct consumption is intended, but exposing the scalar-consumer failures is an unintended regression.
  • Key design decisions: Representing the result as one struct column preserves parent validity. Native decoding validates structural compatibility before restoring planned metadata. Returning extracted fields as scalars permits correct broadcasting to empty and multi-row batches.
  • Implementation sketch: The serde gate controls admission, getStruct invokes the JVM serializer, and resolve_struct constructs the native literal. Per-partition top-K iterators now register subqueries and remove them at task completion. Regression tests cover the existing codegen rejection of embedded subqueries.
  • Performance: Immutable literals avoid repeated JNI lookup and IPC decoding during batch evaluation. Transfer still occurs per physical expression occurrence, including normalization and copying. JVM Arrow buffers and native literal storage are not pool-reserved. No benchmark establishes a net speedup or a separate performance regression.
  • Design: Reusing the Arrow writer and DataFusion literals gives the bridge a straightforward ownership model. The remaining correctness gap is admission into existing scalar-intolerant consumers, rather than struct transport itself.
  • Abstraction & complexity: The shared type gate handles general capabilities, while the small bridge-specific predicate handles arrays/maps and decimal scales. Metadata alignment has a separate, necessary responsibility. No additional abstraction or simplification meeting the P1/P2 bar was identified.
  • Behavioral changes worth calling out: The approved plans now show q9 accelerating 38 of 43 eligible operators and Spark 4.2 q77a accelerating 87 of 89. Non-struct transfer logic remains unchanged. The documented malformed-UTF-8 normalization difference also applies to newly native struct fields. The earlier broadcasting, top-K registration, UDF-dispatch, and reference-test concerns are addressed.
  • Suggested improvements: Resolve the already-reported scalar-consumer blocker by supporting scalar inputs in the affected kernels or retaining fallback for those consumers, with regression coverage for merged subqueries. No additional P1/P2 changes are requested.

Reviewed the entire 21-file diff, including prerequisite commits, at 2f1e93fd06fd05e09a52979b74a718ba0a8da164 against 7d294535e87367d819c4e2a23422cd1158ac31b8. The PR remains non-draft. Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, and review-comet-memory-pr.

Exact-head CI: Comet CI and CodeQL both require approval and have zero jobs. Only labeling passed. The applied Spark 4.1/4.2 labels have not produced an exact-head suite verdict.

Validation: All 13 focused Rust tests passed: six struct IPC/literal tests and seven field-extraction tests. Two disposable reproduction tests confirmed the six affected kernels reject extracted scalars while their array equivalents succeed. No local JVM integration tests, full Spark SQL suites, or benchmarks were run. The reproduction did not exercise the JNI transfer end to end. Disposable tests were removed, and the checkout is unchanged.

No additional introduced P1/P2 issues found within this review. The substantiated existing P2 blocker remains unresolved.

@andygrove
andygrove enabled auto-merge October 23, 2026 08:40

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed commit bb1ce9023ad727e229e371010782fc6957426c83. New P1/P2 findings are detailed inline.

};
return Some(ConsumerProperties {
constant: true,
has_field: false,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Include whole-struct sources in bounded consumer evaluation. A resolved Literal(Struct) contributes its payload size but sets has_field: false, so consumers without GetStructField never receive the wrapper. With a Parquet src containing one 524,288-byte string s, Parquet outer_rows containing IDs 0–8191, AQE disabled, and the default 8,192-row Comet batch, run SELECT id, octet_length(CAST(IF(id > 0, (SELECT max(named_struct('s', s, 'd', 0D)) FROM src), named_struct('s', 'small', 'd', 0D)) AS STRING)) FROM outer_rows. Spark returns 12 for ID 0 and 524295 for every other row. The native expression instead panics while expanding the intermediate string array, before octet_length can reduce it to integers. This newly admitted shape previously fell back successfully. Please include whole-struct sources in bounded evaluation, or preserve fallback until these consumer trees are handled safely.

Evidence: Reproduced at bb1ce90 with the exact optimized expression shape: OctetLength(If(id > 0, Cast(Literal(Struct{s: 524288-byte string, d: 0.0}), Utf8), Literal("{small, 0.0}"))), applying ScalarSubqueryConsumer::wrap at each construction point as the planner does. An 8,192-row RecordBatch panicked with overflow in arrow-buffer 59.3.0 builder/offset.rs:63. Spark 3.5.9 retained the whole-struct scalar subquery, pushed the cast into the IF branches, and returned all 8,192 rows with the expected counts. Source inspection confirms native admission without an incompatible-expression opt-in, while both the base and release struct-subquery gates reject this shape. The disposable reproduction is preserved at /tmp/pr5889-root-bb1-bounds-repro.rs; its output is /tmp/pr5889-root-bb1-conditional.log, and the Spark reference output is /tmp/pr5889-root-bb1-reference.log.

@andygrove
andygrove added this pull request to the merge queue Oct 23, 2026
@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

Thanks for the reviews, @andygrove and @sunchao! Since this PR is already in the merge queue, should we pause the merge and address the new P2 here?

@andygrove
andygrove removed this pull request from the merge queue due to a manual request Oct 23, 2026
@andygrove

Copy link
Copy Markdown
Member

Thanks for the reviews, @andygrove and @sunchao! Since this PR is already in the merge queue, should we pause the merge and address the new P2 here?

Thanks for the ping, I removed from the merge queue

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed commit 30b4fe5f5fce0437296be1e10a452a19a988daeb. New P1/P2 findings are detailed inline.


override def getSupportLevel(expr: ScalarSubquery): SupportLevel = {
val supported = expr.dataType match {
case s: StructType =>

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] [P2] Preserve a supported planning path for merged Bloom-filter subqueries

This new struct admission also enables the shape Spark creates when an injected runtime Bloom-filter aggregate merges with another scalar aggregate. CometBloomFilterMightContain serializes the resulting GetStructField(ScalarSubquery(...)), but native BloomFilterMightContain::try_new evaluates it on an empty batch and accepts only a scalar binary value. GetStructField returns an empty binary array, so planning fails with Bloom filter expression should be evaluated as a scalar binary value. This affects ordinary joins with runtime Bloom filtering and a mergeable scalar aggregate, which previously retained Spark fallback. Please adapt Bloom’s planning-time evaluator to extract one binary value, or retain fallback for this argument shape, while preserving the batch-sized field contract for normal consumers.

Evidence: SQL reproduction on Spark 3.5.9: create Parquet dim with IDs 0–99 and v = 'keep' for IDs 0 and 1, otherwise 'drop'. Create Parquet fact with IDs 0–999. Set spark.sql.adaptive.enabled=false, spark.sql.execution.reuseSubquery=true, spark.sql.autoBroadcastJoinThreshold=-1, spark.sql.optimizer.runtime.bloomFilter.enabled=true, spark.sql.optimizer.runtime.bloomFilter.applicationSideScanSizeThreshold=0, spark.sql.optimizer.runtime.bloomFilter.creationSideThreshold=10485760, and spark.sql.shuffle.partitions=1. Run SELECT fact.id, (SELECT percentile(id, 0.5) FROM dim WHERE v = 'keep' AND id IS NOT NULL) AS p FROM fact JOIN dim ON fact.id = dim.id WHERE dim.v = 'keep'. Spark produces a merged named_struct(bloomFilter, ..., percentile(id, 0.5, 1), ...) and a filter containing might_contain(scalar-subquery[].bloomFilter, xxhash64(id, 42)). Expected rows are (0, 0.5) and (1, 0.5), which Spark returned. At the reviewed head, construct the resolved Literal(ScalarValue::Struct(...)), select its binary field with GetStructField, pass it through ScalarSubqueryConsumer::wrap, then call BloomFilterMightContain::try_new. Both null and valid non-null filter cases reproduce the stated planning error. Equivalent direct ScalarValue::Binary controls succeed. The production planner calls this same constructor at native/core/src/execution/planner.rs:977, and the scalar-only match is at native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs:56. Reproduction source and logs are preserved under /tmp/pr5889-30b4-review-bloom-*. A full Comet JVM SQL run was not performed.

@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

Thanks for pointing this out. I've fixed the Bloom filter’s planning-time evaluation while preserving the runtime array contract.
[Full fork CI] passed at 0209366.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed commit 02093665167aeecd9b862c97f29c810071a9250f. New P1/P2 findings are detailed inline.

if hasStructSubqueryContext(in) &&
(in.value.dataType.isInstanceOf[StructType] ||
in.value.dataType.isInstanceOf[ArrayType]) &&
containsVariableWidth(in.value.dataType) &&

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Apply the dynamic IN evaluation guard independently of field width. Primitive and fixed-width struct values bypass it. A later candidate then runs on a row that already matched and can abort the query. Constant fixed-width membership can also skip a shared candidate that Spark evaluates. Preserve Spark’s per-row and shared-expression rules, or retain fallback for these shapes.

Evidence: With ANSI and subquery reuse enabled, AQE disabled, and single-file Parquet tables items(id,txt)=[(1,'bad'),(2,'2')] and nums(k)=[0,1], run SELECT id, id IN ((SELECT max(k) FROM nums), CAST(txt AS INT)), (SELECT min(k) FROM nums) FROM items. Spark 4.1.3 and the base returned (1,true,0),(2,true,0). Head CometProject raised CAST_INVALID_INPUT on row 1. A fixed-width struct membership query reproduced the same failure. In the shared-candidate control, Spark and base raised the cast error with subexpression elimination enabled, while head returned (true,true) for both rows.

)
.map_err(Into::into)
} else {
Ok(comparison)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve per-row short-circuiting for AND and OR on projections enabled by merged subqueries. This branch returns an ordinary DataFusion BinaryExpr. Its Boolean preselection is heuristic, so a mixed batch can evaluate the right operand where Spark already knows the answer. Apply the appropriate row mask when evaluation is observable, or keep the projection on Spark.

Evidence: With ANSI and subquery reuse on, AQE off, items(id,txt)=[(1,'bad'),(2,'2')] in one batch, and nums(k)=[0,1], run SELECT id, id=(SELECT max(k) FROM nums) OR CAST(txt AS INT)=2, (SELECT min(k) FROM nums) FROM items. Spark and base returned (1,true,0),(2,true,0). Head raised CAST_INVALID_INPUT. Replacing the predicate with id<>(SELECT max(k) FROM nums) AND CAST(txt AS INT)=2 returned (1,false,0),(2,true,0) on both controls and failed on head. The 50% split falls outside DataFusion 55.2.0’s 20% preselection threshold.

case hour: Hour => clockFallbackReason(hour.child, hour.timeZoneId)
case minute: Minute => clockFallbackReason(minute.child, minute.timeZoneId)
case second: Second => clockFallbackReason(second.child, second.timeZoneId)
case _ if !hasStructSubqueryContext(expr) || !touchesCalendar => None

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve nullable operand order when admitting non-calendar consumers. Addition evaluates a failing right operand after a NULL left operand, and DIV evaluates its numerator after a NULL divisor. I also reproduced eager later arguments in startswith and substring. These paths can abort a successful query or advance a stateful expression and change a later row. Apply each operator’s Spark evaluation order or retain fallback.

Evidence: Use one-file Parquet items(id,txt)=[(1,'bad'),(2,'2')], nums(k)=[0,1], ANSI on, and AQE off. Substitute z=(SELECT max(named_struct('f',k)) FROM nums WHERE k<0).f. Both SELECT z+CAST(txt AS INT) FROM items and SELECT CAST(txt AS INT) DIV z FROM items returned two NULLs on Spark and base. Both raised CAST_INVALID_INPUT on head. On one partition, SELECT id, IF(id=1,z,id)+monotonically_increasing_id() FROM items returned (1,NULL),(2,2) on the controls and (1,NULL),(2,3) on head. The ordinary equality control returned NULLs on all three.

case hour: Hour => clockFallbackReason(hour.child, hour.timeZoneId)
case minute: Minute => clockFallbackReason(minute.child, minute.timeZoneId)
case second: Second => clockFallbackReason(second.child, second.timeZoneId)
case _ if !hasStructSubqueryContext(expr) || !touchesCalendar => None

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve generated Spark concat semantics for newly admitted struct-subquery consumers. Spark’s generated string projection stops evaluating arguments once one is NULL. Native concat evaluates a later failing cast on that row. Account for Spark’s active execution mode when masking arguments, or retain fallback.

Evidence: With ANSI and whole-stage codegen enabled, use items(id,txt)=[(1,'bad'),(2,'2')], nums(k)=[0,1], and z=(SELECT max(named_struct('f',k)) FROM nums WHERE k<0).f. Run SELECT id, concat(IF(id=1,CAST(z AS STRING),CAST(id AS STRING)),CAST(CAST(txt AS INT) AS STRING)) FROM items, substituting z. Spark and base returned (1,NULL),(2,'22'). Head raised CAST_INVALID_INPUT on 'bad'. With whole-stage codegen disabled and spark.sql.codegen.factoryMode=NO_CODEGEN, all three raised the error.

let result = date_part(&array, DatePart::$date_part_variant)?;
Ok(ColumnarValue::Array(result))
let offset = timezone
.offset_from_utc_datetime(&endpoint)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Use the regional offset for the requested instant. Sampling chrono’s endpoint applies standard time to a far-future summer instant in Los Angeles. An ordinary IF can return a runtime scalar, bypassing the Scala guard for subquery consumers. Use Spark-equivalent recurring timezone rules or defer this scalar clock shape to Spark.

Evidence: Set spark.sql.session.timeZone=America/Los_Angeles and put integer IDs 1 and 2 in one Parquet batch. Run SELECT hour(IF(id>0,timestamp_micros(9089364630896000000L),timestamp_micros(CAST(id AS BIGINT)))) FROM items. The instant is July 1 of year 290000 at 12:34:56 UTC. Spark 4.1.3 returned [5,5]. Head returned [4,4]. The base rejected the scalar input, so the new path turns an explicit failure into an incorrect hour.

.downcast_ref::<ListExtract>()
.and_then(ListExtract::default_value);
for child in expr.children().into_iter().chain(default_value) {
let child = consumer_properties(child)?;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep a deterministic condition from disabling width protection for an independent branch. Integral DIV introduces a CheckOverflow node outside the recognized set. The ? propagates that unknown child through the entire IF and its reducing parent, allowing a wide field to expand across the full batch. Handle this row-local node and retain a bounded path or fallback when a child cannot be classified.

Evidence: Use two 262,144-byte strings in Parquet wide(s) and IDs 0–8191 in one Parquet batch o, with ANSI and reuse on and AQE off. Run SELECT id,octet_length(IF(id DIV 2 >= 0,(SELECT max(s) FROM wide),(SELECT min(s) FROM wide))) FROM o. Spark and base returned 8,192 lengths of 262144. Head panicked with offset overflow in GetStructField::broadcast_field. Changing only the condition to id >= 0 made head return all expected lengths.

private def hasVariableWidthStructSubquery(expr: Expression): Boolean =
expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
case subquery: ExecSubqueryExpression =>
subquery.dataType.isInstanceOf[StructType] && containsVariableWidth(subquery.dataType)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Include expanding intermediate values in the fallback decision for stateful or JVM-dispatched consumers. An integer subquery field can control repeat and produce a wide string. Checking only the materialized field type admits the tree, while the opaque or stateful predicate prevents bounded evaluation of its parent. Retain fallback when the intermediate cannot be consumed safely.

Evidence: Use Parquet src(n)=[262144] and IDs 0–8191 in one outer batch, with AQE off and reuse on. Run SELECT sum(y),min(m) FROM (SELECT octet_length(IF(monotonically_increasing_id()>=0,repeat('x',2*(SELECT max(n) FROM src)),'z')) y,(SELECT min(n) FROM src) m FROM outer_rows) q. Spark and base returned (4294967296,262144). Head failed with Scalar subquery consumer result exceeds Arrow's 32-bit offset limit. Replacing the predicate with upper(CAST(id AS STRING)) <> 'missing' reproduced the same failure. Both final outputs fit.

};
return Some(ConsumerProperties {
constant: true,
has_scalar_struct: matches!(literal.value(), ScalarValue::Struct(_)),

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Carry the origin of resolved subquery literals into this check. An ordinary CAST(NULL AS STRUCT<k:BIGINT>) also matches ScalarValue::Struct. An enclosing array constructor consequently switches an existing native projection to one-row execution even when the SQL contains no subquery.

Evidence: Write one Parquet partition from SELECT id%2=0 AS p,named_struct('k',id) AS payload FROM range(262144). Run SELECT array(IF(p,payload,CAST(NULL AS STRUCT<k:BIGINT>))) FROM t with 8,192-row batches. Both revisions used CometProject and produced 262,144 rows. Consuming the projection’s columnar batches on Spark 4.1.3 gave five-run warm medians of 53 ms on base and 1,634 ms on head.

inner,
constant: properties.constant,
max_rows: if properties.unbounded_output {
1

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep integer-to-string casts batched when their output has a known small bound. Formatting a signed BIGINT needs at most 20 bytes. Setting max_rows to one makes the cast and its downstream length kernel run once per row, followed by concatenation. Use the byte bound to choose a safe slice size.

Evidence: Use one Parquet partition with payload.k holding BIGINT IDs 0–262143 and src(n)=[262144]. Run SELECT octet_length(CAST(payload.k+(SELECT max(n) FROM src) AS STRING)),(SELECT min(n) FROM src) FROM t with reuse on, AQE off, and batch size 8,192. Both revisions returned checksum 68721049600 when both columns were consumed. On one Spark 4.1.3 worker, five post-warm-up actions had a median of 114 ms on base and 2,213 ms on head.

| DataType::Binary
| DataType::LargeBinary
| DataType::BinaryView
) || dt.is_nested()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep fixed-cardinality array constructors vectorized when their elements have bounded width. dt.is_nested() marks a three-BIGINT array as unbounded. With a genuine merged subquery, that decision propagates to the enclosing hash and reduces its batch size to one. Derive a bound from the constructor’s arity and element types.

Evidence: Use one Parquet partition t with payload.k holding BIGINT IDs 0–262143 and Parquet v(id)=[0,1]. Run SELECT sum(CAST(hash(array(payload.k,CAST((SELECT max(id) FROM v) AS BIGINT),CAST((SELECT min(id) FROM v) AS BIGINT))) AS BIGINT)) FROM t with AQE off and batch size 8,192. Spark merged the two subqueries. Base and head returned 1095857758320. Five post-warm-up end-to-end actions on one worker had medians of 242 ms on base and 2,837 ms on head.

@github-actions github-actions Bot added the area:shuffle Shuffle (JVM and native) label Oct 26, 2026
@LinSimon-901101
LinSimon-901101 marked this pull request as draft October 26, 2026 08:35
@LinSimon-901101
LinSimon-901101 force-pushed the fix/5834-struct-scalar-subquery branch from dce5b53 to 0209366 Compare October 26, 2026 08:36
LinSimon-901101 added a commit to LinSimon-901101/datafusion-comet that referenced this pull request Oct 26, 2026
Distinguish materialized subqueries from ordinary struct literals and share retained results. Bound wide values across physical operators, aliases and stateful consumers, while keeping bounded numeric casts and fixed-size array/hash consumers batched. Include the related NULL metadata, overlap and regression fixes.

Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review-Items: R10, R08, R07, R09, R11, R12
LinSimon-901101 added a commit to LinSimon-901101/datafusion-comet that referenced this pull request Oct 26, 2026
Preserve Spark evaluation order for dynamic IN, Boolean short-circuiting, generated concat and nullable operands, including Bloom initialization and error ordering. Keep the required ANSI TopK cases native with comparator-driven selection, bounded workspace and owned payloads; include the matching proto, JNI and memory regressions.

Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review-Items: R01, R02, R04, R03
LinSimon-901101 added a commit to LinSimon-901101/datafusion-comet that referenced this pull request Oct 26, 2026
Preserve regional timezone clock extraction and checked scalar Date32-to-epoch conversion through compatible Spark execution. Retain supported native fixed-offset, timestamp-without-timezone and date-column paths, with overflow, regional-midnight and conditional-shape regressions.

Review: apache#5889 (comment)
Review: apache#5889 (comment)
Review-Items: R05, R06
@LinSimon-901101

LinSimon-901101 commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor Author

Thanks @sunchao ,for the detailed review. The follow-up fixes are grouped into three commits:

[Full fork CI] passed on the current PR head, ab9f41b.

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to withdraw my approval. When I approved bb1ce90 on 10-23, the PR was about 5,600 lines, and as I said then, that was already more than it should have been. It is now about 14,400 lines, and the production code alone has more than tripled. Most of the new code hasn't had a human review, and I don't think it can be reviewed well as one change.

Could we cut this back to the struct-typed subquery support and land the rest as separate PRs? The scalar-argument kernel fixes are tracked in #3336 and #7245, and the top-K subquery registration is tracked in #6676. Neither depends on the struct support. Much of the newer code guards against bugs that also hit plain subqueries and columns on main. I've filed those so they don't need to be fixed here: #7246 for operands evaluated on rows Spark skips under ANSI, and #7247 for large scalars that overflow Arrow offsets when expanded to the batch length. Dates outside chrono's range are already tracked in #6843. With those out of scope, I'd expect the bounded consumer evaluation, the provenance tags, the evaluation-order rules, SparkTopKExec and Native.mergeTopK to come out as well.

Several of the asks that pulled this code in were mine, so I'm sorry for the churn.

@andygrove

Copy link
Copy Markdown
Member

@sunchao, to make the scope concrete for the next push, here is where I think each of your open findings belongs.

The IN, AND/OR, nullable operand and concat findings are #7246. The same queries fail on main with a plain subquery or a literal, so they don't need a fix in this PR.

The width findings (field expansion, whole-struct sources, operator boundaries, CheckOverflow and repeat) are #7247. If the field read returns a scalar the way a plain subquery does, instead of expanding to the batch length, the merged form behaves like the plain one and the bounded consumer evaluation isn't needed. The literal origin, integer cast and array constructor slowdowns come from that evaluation, so they go away with it.

The fixed-offset array, padding, regional clock and scalar Date32 findings are about the scalar-argument kernels, so they belong with the fixes for #3336 and #7245. Those should land before the struct support, so that merged subqueries never reach a kernel that rejects scalars.

The TruncDate finding is #6843, which affects plain date columns on main too.

The merged Bloom filter planning failure is specific to struct results, so it still needs to work here. With a scalar field read, BloomFilterMightContain::try_new gets the scalar binary value it expects.

@LinSimon-901101

Copy link
Copy Markdown
Contributor Author

I agree this PR has grown too large. Thanks @andygrove for laying out a clear plan for splitting it up. I'll bring #5889 back to its core scope of supporting struct-typed scalar subqueries and move the independent fixes into separate PRs.

Thanks to both @andygrove and @sunchao for the careful reviews and thoughtful discussions.

@LinSimon-901101
LinSimon-901101 marked this pull request as draft October 28, 2026 08:15
@LinSimon-901101
LinSimon-901101 force-pushed the fix/5834-struct-scalar-subquery branch 2 times, most recently from 2378130 to 37ae130 Compare October 29, 2026 17:17
@LinSimon-901101
LinSimon-901101 force-pushed the fix/5834-struct-scalar-subquery branch from 37ae130 to b97a3be Compare October 29, 2026 17:35

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The scalar field extraction now handles a zero-row planning batch correctly. Reusing the same wide struct still means copying it repeatedly during native planning.

let data_type = to_arrow_datatype(expr.datatype.as_ref().unwrap());
Ok(Arc::new(Subquery::new(self.exec_context_id, id, data_type)))
if matches!(data_type, DataType::Struct(_)) {
let value = Subquery::resolve_struct(self.exec_context_id, id, &data_type)?;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each occurrence of the same struct subquery calls resolve_struct, which serializes the full JVM row and decodes another Arrow struct. Projecting multiple fields from one reused scalar subquery can therefore keep multiple complete copies alive in the same task. Could the planner share the resolved scalar by subquery ID and type within an execution context? Two field projections from one subquery would make a useful serialization-count test.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:expressions Expression evaluation area:ffi Arrow FFI / JNI boundary area:shuffle Shuffle (JVM and native) bug Something isn't working run-spark-4.1-tests Run the Spark 4.1 SQL tests on this pull request instead of waiting for the merge queue run-spark-4.2-tests Run the Spark 4.2 SQL tests on this PR

Projects

None yet

5 participants