Repository navigation
fix: preserve CASE result field metadata - #26145
osipovartem wants to merge 7 commits into
Conversation
github-actions
Bot
added
logical-expr
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #26145 +/- ##
==========================================
- Coverage 82.77% 82.76% -0.02%
==========================================
Files 1148 1148
Lines 451075 451632 +557
Branches 451075 451632 +557
==========================================
+ Hits 373397 373810 +413
- Misses 54930 54976 +46
- Partials 22748 22846 +98 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
osipovartem
changed the title
|
PTAL at #24865 I originally had this strict metadata policy like you do here but changed it to be much more permissive on the basis of @paleolimbot 's comments. |
|
Thanks for pointing me to #24865. I agree CASE should use the same policy: merge unrelated metadata keys permissively, keep extension name/parameters paired, and reject conflicting extension names. The current strict full-map equality in this PR is inconsistent with that direction. I will align it with the shared helper and add CASE-specific tests for missing/unrelated metadata and conflicting extension names before this is merged. |
|
While I don't necessarily mind the approach here, I think it needs some justification for why it's needed over the simpler fix. This is a good workaround for when you don't control the ScalarUDF source code, but the actual bug here is that I put a note on the other PR as well, but I recently discovered |
|
Thanks. CASE is a separate logical drop site: direct |
ee2aa10 to
d7aebcc
Compare
| /// Return the field and column nullability together. | ||
| /// | ||
| /// The default preserves custom implementations of [`Self::nullable`]. | ||
| fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> { |
There was a problem hiding this comment.
the field has a nullability column, right? Why can't this just use the existing nullability from Field?
There was a problem hiding this comment.
DFSchema does use Field::is_nullable(). The default preserves custom ExprSchema::nullable() overrides, which may disagree with the field or return errors. The paired method avoids a second DFSchema lookup without changing that contract. Added a regression for both disagreement directions in e0753c3.
…#24865) ## Which issue does this PR close? - Part of apache#24860. ## Rationale for this change `COALESCE`, `NVL`/`IFNULL`, and `NVL2` drop the field metadata of their arguments when computing their return field: `return_field_from_args` builds the output `Field` from the argument *data types* only. For columns carrying Arrow extension types (`ARROW:extension:name` / `ARROW:extension:metadata` — e.g. `arrow.uuid`, `geoarrow.wkb`), the planned schema of any expression containing these functions silently loses the extension-type identity, and metadata-aware UDFs that dispatch on their argument `Field`s fail to plan over them (see apache#24860 for a full reproducer). This is one of several independent metadata drop sites listed in apache#24860. It is the plan-time schema fix for the conditional functions. These functions simplify into `CASE`, whose logical and physical drop sites are addressed separately in apache#26145 and apache#26158. ## What changes are included in this PR? - `datafusion_common::metadata::coerced_fields_metadata`: computes the metadata for an expression whose result may come from any one of several inputs coerced to a common type. It lives next to the related metadata utilities so that `CASE` (in `datafusion-expr`) can share it: - Metadata is merged across the inputs whose data type is not `Null` (an untyped NULL literal carries no metadata and cannot contribute a typed value), with earlier inputs taking precedence for repeated keys. Unrelated keys are not required to agree, so metadata that arrived with the data (e.g. from an Arrow file) cannot cause an extension type to be dropped. - The extension type keys are handled as a unit, as in `Expr::Cast` output field derivation: differing values of `ARROW:extension:name` are a planning error, and `ARROW:extension:metadata` always comes from the input that supplied the extension name (byte-wise differences in it are tolerated, since extension parameters are often JSON). - `coalesce`: `return_field_from_args` attaches the merged metadata of its arguments. Return type and nullability are unchanged. `NVL`/`IFNULL` delegate to `coalesce` and are fixed by the same change. - `nvl2`: same, over the second and third arguments only — the first argument is only tested for NULL and never contributes a value, so its metadata is excluded. ## What is the testing strategy for this PR? - Unit tests for `coerced_fields_metadata` in `metadata.rs`: merge precedence, untyped NULL inputs skipped, empty result, error on differing extension names, and extension name/metadata kept paired. - Unit tests in `coalesce.rs` and `nvl2.rs` exercise `return_field_from_args`: metadata propagated, untyped NULL arguments not blocking propagation, conflicting extension types rejected, NVL2's test argument excluded, and nullability unchanged. There is deliberately no sqllogictest yet: `arrow_metadata()` reads its argument field at execution time, and by then these functions have been simplified into `CASE`, so end-to-end coverage depends on the `CASE` fixes (apache#26145, apache#26158). This ordering is safe: the optimizer's schema invariant ignores metadata (`assert_expected_schema` / `logically_equivalent_names_and_types`), so a plan whose `coalesce` carries metadata that the simplified `CASE` does not still plans. ## Are there any user-facing changes? - The planned schema (e.g. `DataFrame::schema()`) of expressions containing `COALESCE`/`NVL`/`IFNULL`/`NVL2` over metadata-bearing columns now preserves that metadata. - `COALESCE`/`NVL`/`IFNULL`/`NVL2` over arguments with *differing* `ARROW:extension:name` values (e.g. an `arrow.uuid` column and an `arrow.json` column) is now a planning error. Previously these planned and silently dropped the extension type. - New public function `datafusion_common::metadata::coerced_fields_metadata`. No other API changes.
d7aebcc to
e0753c3
Compare
| } | ||
| } | ||
|
|
||
| // Resolve metadata bottom-up for branches whose metadata can be preserved. |
There was a problem hiding this comment.
this seems quite complicated for resolving metadata -- can you explain why it needs a custom worklist?
There was a problem hiding this comment.
The bottom-up path avoids re-running type/nullability inference for every suffix of a nested CASE chain. I kept it for CASE, forwarding wrappers and binary expressions; other branches use native to_field without the depth cutoff. Added counted regressions and nested benchmarks; CASE/UDF re-inference remains a separate limitation.
| assert_eq!(meta, outer_ref.metadata(&schema).unwrap()); | ||
| } | ||
|
|
||
| #[test] |
There was a problem hiding this comment.
what cases does this cover that the .slt tests do not? Can we conslidate to slt?
osipovartem
changed the title
|
I tried this stack (#26145 + #26158) with a Parquet Variant column, since the slt tests here and in #26895 only cover extension metadata on On
Would it be worth adding one struct-backed extension case to |
|
Yes, please. A struct-backed extension case in metadata.slt would be valuable. The execution-level coverage is in #26158, which uses this logical fix as well. Thanks for validating the stack. |
@osipovartem I've opened the struct-backed cases against your branch: Embucket#186. |
Which issue does this PR address?
Part of #24860, complementary to merged #24865 and physical follow-up #26158.
Rationale
Logical CASE field inference loses result metadata needed by later planning and metadata-aware function signatures. The shared
coerced_fields_metadatapolicy should apply to THEN/ELSE fields, not the CASE operand or WHEN conditions.What changes are included?
to_fieldfor ordinary branches. Remove the depth-8 cutoff and whole-CASE metadata erasure for unhandled expression variants. Preserve the documented non-extension cast metadata, including alias-wrapped NULL CAST/TRY_CAST expressions.ExprSchema::field_and_nullablemethod preserves customnullable()overrides; DFSchema obtains both results from one field lookup.What is the testing strategy?
Ordinary SQL/extension-policy assertions were removed from the large Rust unit test. Execution-level acceptance cases are consolidated in
metadata.sltin #26158, including the cases proposed in #26895. They exercise physical metadata and are not presented as end-to-end fixes supplied by this logical-only PR.Focused Rust tests retain API-only field attachments, pre-coercion type guards, custom schema nullability/errors, and inference-count checks. The aliased NULL cast regression was reproduced failing before its fix. The separate binary/CASE negative control required 97 schema lookups at depth 8 against a bound of 56 before bottom-up binary inference; its final coverage is consolidated into the existing counted test rather than duplicating fixtures. A further negative control verifies that
NOT(CAST(NULL AS BOOLEAN))remains a typed Boolean contributor rather than bypassing the pre-coercion storage-type guard.Local logical-only validation at
66121cee1: 305 expr unit tests; the documented extended workspace run passed 13,351 Rust tests (8 existing ignored) and 545/545 SLT files. Workspace all-target/all-feature Clippy with-D warnings, rustfmt and diff checks passed. Main was subsequently merged into this branch asbd6ab71dd; the two changed expr source/benchmark files are byte-identical. The main-synced combined stack with physical #26158 at3018e535aadditionally passed 13,383 Rust tests (8 ignored), 545/545 SLT files and strict workspace Clippy. That combined run is not a logical-only run. The public SQL integration checks original/optimized logical metadata, returned values and extension-conflict errors; physical CASE metadata remains #26158's scope.Performance
The Criterion benchmark now includes nested columns, aliases/casts, binary/CASE chains and alternating CASE/UDFs, in addition to the existing flat literal/column cases. Fixtures are constructed outside timing. Measurements use the same host and
ci-benchprofile (optimization level 1, Rust 1.99.1); these are field-inference microbenchmarks, not query-runtime or production-release performance claims.The binary-path screen at
66121cee1, before the main merge, used 20 samples and one-second warmup/measurement. To exclude stale artifacts when different source workspaces share a target directory, the entire benchmark profile was cleaned and rebuilt between revisions. Initial runs without that isolation were discarded. Medians in us:49c67c2961f70d3d9066121cee1Bottom-up binary inference removes the previous PR's quadratic schema-inference work, but this is not parity with main: marked-last planning costs about 2.9x main, and the nested-binary fixtures remain about 2.1-2.2x main. Main does not provide the added CASE metadata semantics, including the marked-last output; the table makes that cost explicit rather than presenting correctness fixes as free. The three flat controls differ about -1.5% to +1.7% from the previous PR in this single sequential screen. These are not whole-query or zero-regression claims, and API/performance acceptance remains a maintainer merge decision.
A straightforward recursive prototype that obtains each result branch through
to_fieldrepeatedly revisits nested CASE inference. In the earlier screen against logical1f70d3d90, the narrowed traversal measured about 40 us versus 365 us for aliases and 59 us versus 2,740 us for casts at depth 128. The lookup-count regression rejects that recursive prototype. This historical comparison justifies retaining the traversal; it is not a speedup over the prior worklist or a fresh timing of the final head.There is a correctness/performance tradeoff in deeply alternating CASE/UDFs: the previous cutoff silently discarded metadata. An earlier short exploratory run compared about 418/1,686 us at depths 6/8 against 910/7,917 us for logical
1f70d3d90preserving metadata. This is historical evidence, not a fresh timing of the final head. Repeated type/nullability/field inference remains broader follow-up work; the worklist does not solve it. No zero-regression claim is made for that path.Two earlier sequential flat-path A/B repetitions used 20 samples with one-second warmup and measurement, comparing original PR head
adb0656against logical1f70d3d90, before the binary follow-up, not Apache main. Representative historical medians in us:1f70d3d90A / BAcross those six 100/1000-branch flat cases, paired medians increased about 0.9-8.0%. Unchanged fast-path source does not justify dismissing measurements or claiming zero regression. These costs accompany correct metadata preservation and removal of unsupported/depth-based erasure; their acceptance remains a maintainer decision. Query-level latency was not measured, and the final single-pair screen does not supersede these repetitions.
Are there user-facing or API changes?
Logical CASE metadata is retained rather than silently discarded.
ExprSchemagains a defaulted method without changing the meaning of existing custom nullability implementations. This does not change execution kernels or claim to fixarrow_metadata(CASE ...)without the physical follow-up.