test: fail Comet tests when a native plan mislabels a timestamp's timezone

[andygrove]

Which issue does this PR close?

There is no issue for this. It follows up the timezone audit (#6335), whose label bugs, such as #6328 and #6330, shipped with passing tests.

Rationale for this change

Native plans label every TimestampType value Timestamp(µs, "UTC") and every TimestampNTZType value with no timezone. The label never changes the values, so checkSparkAnswer and friends cannot see a wrong one: they compare collected rows, export_batch hands each array to the JVM with whatever type the plan produced, and CometPlainVector.getLong reads the raw microseconds. The casts at ScanExec and in the shuffle writer's SchemaAlignExec repair a wrong label as well.

Inside a native plan the label does matter. Another zone breaks comparisons and CASE/coalesce coercion against "UTC" arrays, and a missing one makes kernels such as hour and casts to string read the value as wall-clock time. So a test of an expression's direct output passes even when the expression breaks a query that composes it, which is how #6328 and #6330 shipped with passing tests.

What changes are included in this PR?

An internal, test-only config, spark.comet.testing.strictTimestampLabels.enabled, which CometTestBase turns on for every Comet suite. It defaults to false, so nothing changes outside the test suites. With it on:

• CometExecIterator checks each batch a native plan returns against the plan's Spark output types, including nested and dictionary-encoded columns (NativeUtil.checkTimestampLabels). The callers that know their output types pass them in a new outputTypes parameter: CometNativeExec, CometNativeScanExec, CometIcebergNativeScanExec, the native limit, and CometTakeOrderedAndProjectExec.
• ScanExec fails, instead of casting, on an input timestamp whose label differs from the declared type.
• The native shuffle writer fails planning, instead of wrapping its child in SchemaAlignExec, when the child labels a timestamp differently from the declared schema. The error includes the native plan, so it points at the expression.

The native checks share describe_timestamp_label_mismatch in datafusion-comet-common.

These are boundary checks. A wrong label that is consumed entirely inside one native plan, such as hour(f(ts)) where f drops the label, never reaches a boundary, and still shows up only as wrong values in a non-UTC session.

How are these changes tested?

• Unit tests: Rust tests for describe_timestamp_label_mismatch, the strict ScanExec, and the strict shuffle-writer alignment. NativeUtilSuite tests for the JVM check, covering nested types, maps, dictionary values imported over FFI, and closing a batch it rejects. CometExecSuite's config serde test covers the new natively read key.
• Planted bugs, to show the check fires: with timestamp_trunc changed back to labelling its result with the session timezone (its behavior before fix: keep the input's timezone label on native date_trunc results #6345), SELECT c0, date_trunc('HOUR', c0) FROM tbl and a GROUP BY over the truncated value both passed checkSparkAnswer with the check off. With it on, they failed at the plan output and at the shuffle writer. With CometSparkToColumnarExec labelling its timestamps Etc/UTC, the ScanExec check failed the same way. On unmodified code the same probes raised nothing.
• With the check on, 2,446 tests in 26 timestamp-heavy suites pass on Spark 4.1, so main has no label drift on those paths. They cover the SQL file suite, casts, Parquet reads, expressions, joins, aggregates, windows, shuffle, the in-memory cache, fuzz, the Python runner, and the Iceberg read and write suites.

[sunchao]

Reviewed the full 19-file diff from 0691bd39f75da61e02aa23ce188de141bd623252 to 5e758b23c1932d623e224d191fa9f51c37fb5800. The PR is not a draft. The snapshot and live discussion contain no existing reviews or comments. No introduced P1/P2 issues found within this review.

Summary

• Prior state and problem: Row comparisons could miss incorrect timestamp timezone labels, while scan casts and shuffle schema alignment could erase the evidence. Such labels can affect subsequent native expressions even when collected values look correct.
• Design approach: Add an internal spark.comet.testing.strictTimestampLabels.enabled switch, enabled by CometTestBase, to reject label mismatches at native output, scan-input, and shuffle-writer boundaries.
• Correctness: Checked recursive timestamp matching, dictionary-value lookup, output-type propagation, config serialization, and rejected-batch cleanup. The checks preserve timestamp values and leave ordinary casts unchanged when disabled. No reproducible P1/P2 correctness issue was identified.
• Compatibility analysis: Compared Spark’s ArrowUtils sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. All distinguish microsecond timestamps with a timezone from TimestampNTZType without one. Requiring exactly UTC enforces Comet’s existing internal convention. No version-specific shim change is needed for this check.
• Key design decisions: Keep the switch false by default, serialize its resolved boolean value, inspect dictionary values rather than indices, and report nested field paths. The JVM rejection path closes the batch before throwing.
• Implementation sketch: NativeUtil compares imported Arrow fields with Spark output types supplied through CometExecRDD and iterator callers. Rust shares describe_timestamp_label_mismatch between ScanExec and shuffle-input alignment, with the flag propagated through execution settings and the planner.
• Performance: Validation traverses type metadata without scanning rows or copying payload buffers. Recursive checks are disabled by default. Output-type metadata and flag plumbing still add some per-plan work, but there is no evidence of a P1/P2 performance regression. No throughput benchmark was run.
• Design: Checking boundaries before label-repairing casts directly addresses the testing blind spot and keeps production behavior stable. The boundary checks are straightforward to reason about.
• Abstraction & complexity: The shared Rust helper avoids duplicating recursive matching. A separate JVM helper is appropriate for Arrow Java fields, Spark types, and dictionary providers. No unnecessary framework or actionable complexity problem was identified.
• Behavioral changes worth calling out: Strict Comet tests now reject labels such as Etc/UTC or +00:00 where UTC is expected, despite equivalent offsets. Comparison with branch-1.1 at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1 confirms this PR adds intended test-only enforcement while retaining default boundary casts. Incorrect labels consumed entirely inside one native plan remain outside these checks’ coverage.
• Suggested improvements: No additional change meeting the P1/P2 reporting bar was identified. The existing nested-type, dictionary, cleanup, and strict-versus-default tests cover the behavior introduced here.

Routed skills: review-comet-pr, review-comet-ffi-pr, review-comet-shuffle-pr, and review-comet-memory-pr, with the timezone and relevant contributor guides.

Exact-head CI: All executed check runs passed, including required-check aggregates, Linux Comet suites across Spark 3.4–4.2, Rust tests, and TPC-H/TPC-DS verification. The inspected Rust and JVM logs identify the merge of this exact head into the requested base and show the new regression tests passing. See the primary run and additional Spark-profile run. Spark’s own SQL suites, external Iceberg suites, and macOS tests were skipped.

Local validation: Three Rust tests using the exact timestamp-label helper and two Scala probe groups passed, covering nested fields and dictionary lookup. These were isolated helper checks, not full integration runs. The normal Cargo test command stopped because the offline dependency cache lacked DataFusion 55.2.0. Full native/JVM integration validation therefore relies on CI. The project checkout remains unchanged.