Repository navigation
fix: propagate field metadata through coalesce, nvl, and nvl2 - #24865
Conversation
|
In this PR, if the branches of the coalesce, etc don't agree on the metadata, we drop extension metadata defensively. My opinion is that this should actually be a planning time error unless all rows will hit only one branch but wanted to take the more conservative approach for now. I want to solicit opinions from the reviewers on what the desried behavior is here. |
These functions built their return Field from argument data types only, dropping field metadata (e.g. Arrow extension types) from the planned schema. Attach the metadata shared by every value-contributing argument instead. Part of apache#24860.
82761e8 to
1ea88c0
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24865 +/- ##
==========================================
+ Coverage 82.49% 82.58% +0.09%
==========================================
Files 1140 1142 +2
Lines 439034 440994 +1960
Branches 439034 440994 +1960
==========================================
+ Hits 362165 364182 +2017
+ Misses 54962 54813 -149
- Partials 21907 21999 +92 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
Thank you for this fix! I'm not a committer so I can't help with the CI but I'll review from a metadata perspective.
I think the family of functions that are affected here are ones that have >1 "compatible" (castable to a common type) inputs. While there are some extension aware type systems that can handle this (notably: vctrs via vec_ptype2() and vec_cast()), datafusion has no mechanism for that and so probably the best we can do for now is to propagate agressively and rely on engines to write analyzer or optimizer rules to enforce other behaviour.
One other place I've seen this pattern is combining values in a values literal clause. I would personally love if the logic were the same there but that might insert too much scope here.
datafusion/datafusion/expr/src/logical_plan/builder.rs
Lines 309 to 320 in d02bd78
I am curious what @timsaucer would expect to happen to metadata in CASE WHEN since that is a metadata usage I'm less familiar with!
…on-common Requiring the arguments to agree on metadata was too strict: extension types with JSON-encoded parameters can differ byte-wise, and unrelated metadata that arrived with the data should not drop an extension type from the result. Propagate the first non-empty metadata among the value-contributing arguments instead, and move the helper next to the related utilities in datafusion_common::metadata so the upcoming CASE fix in datafusion-expr can reuse it.
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
I know you're still working here, just making sure I didn't forget to review!
Adopt the suggested compromise: merge metadata across the value-contributing arguments (earlier arguments win for repeated keys) instead of taking only the first non-empty metadata, but refuse differing values of 'ARROW:extension:name' with a planning error, keeping the extension name and its parameters paired. Also accept impl IntoIterator and inline the field construction in the datafusion-common tests.
|
@james-willis Can you address the merge conflicts? I can bring up this PR in the weekly sync tomorrow to get a second review if it's ready. |
# Conflicts: # datafusion/functions/src/core/nvl2.rs
Drop the field/ext_field/return_field test helpers in favor of Field::new().with_metadata() and direct return_field_from_args calls, matching the datafusion-common tests in this PR and the existing tests in with_metadata.rs.
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
Thank you!
This would benefit from a second review to double check my take on metadata merging here, but in general I think this a standalone improvement and fixes a hard to work around bug.
| /// [`Expr::Cast`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html#variant.Cast | ||
| pub fn coerced_fields_metadata<'a>( |
There was a problem hiding this comment.
I recently discovered intersect_metadata_for_union(), which silently discards non-matching metadata keys. I like this version better (I want a louder error when the types don't match), but either is OK with me if we want to keep the existing logic for consistency.
There was a problem hiding this comment.
I don't totally understand the UNION plumbing but intersect_metadata_for_union is only used for the first stage and is not what UNION will ultimately give for the query result's type. The ultimate behavior of UNION is also not great here.
I'm inclined to keep this PR as-is because I think probably the different branches need to agree at least on the extension type and it is important to enforce that. Even better we do it at plan time.
I would like to be more strict with the other metadata because often that is part of the type (e.g., the CRS is part of the type of a geometry, randomly picking one of the CRSes will give incorrect query results) but I'm not confident in that being universally the right call to enforce it here (e.g., in the CRS case we would want to convert to item-level CRS when branches diverge).
| // Propagate field metadata (e.g. Arrow extension types) from the | ||
| // arguments that can supply a value; see `coerced_fields_metadata` | ||
| // for the merge semantics. | ||
| let metadata = | ||
| coerced_fields_metadata(args.arg_fields.iter().map(|f| f.as_ref()))?; | ||
| Ok(metadata | ||
| .add_to_field(Field::new(self.name(), return_type, nullable)) | ||
| .into()) |
There was a problem hiding this comment.
Happy to hear another opinion, but I do think this is a bug worth fixing regardless of whether #26145 is a good idea.
mprammer
left a comment
mprammer
left a comment
There was a problem hiding this comment.
Thanks for the PR! Chiming in to say that I ran into the same problem. @paleolimbot's suggested way of merging metadata seems right to me.
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you for this PR @james-willis and for the reviews @paleolimbot and @mprammer
I think we need "end to end" tests of this -- namely to verify that in an actual query the metadata is preserved.
I am pretty sure this PR won't actually fix the problem as it nvl etc are rewritten macro-style to CASE which loses the metadata. In fact I wrote up some tests in a separate PR showing this (this PR doesn't fix the problem):
Here is the kind of tests I think are needed
With those tests I don't think there is any need for the rust unit tests in this PR
coalesce, nvl, ifnull and nvl2 are rewritten to CASE before execution, so arrow_metadata() cannot observe the field they compute during planning. - metadata.slt: conflicting extension types are a planning error, and nested cases observe the planning-time field through that error (an outer coalesce only reports the inner extension type if it was propagated). - core/tests/sql/select.rs: the planned schema (DataFrame::schema()) carries merged non-extension metadata for each function. - Remove the coalesce.rs and nvl2.rs unit tests these replace. The coerced_fields_metadata unit tests remain.
github-actions
Bot
added
core
|
Thanks for taking a look @alamb. I agree that this PR is not sufficient to fix the metadata dropping issues in these functions, but it is necessary: COALESCE’s (et al's) output field is computed at initial planning, before that rewrite, and a metadata-aware UDF over COALESCE checks its argument there. Without this change the query fails at that point even once CASE is fixed. Since this PR fixes something only observable at plan time, I added a test in select.rs to capture that planning went as expected. The style in your linked PR wouldn't let me capture that. |
Thanks -- can you be clear what your usecase is? In other words, how are you hitting the code that this PR solves? Is it like a third-party frontend ? Why isn't the simplifier being run in your usecase? |
It's being run, but the return type is also being computed before it runs (e.g., asking for the schema of the unoptimized logical plan). For something like Agreed that we need to fix more to get this to work end to end, but we do also need to fix the |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Got it -- makes sense to me. Thank you
|
My intent is to work toward end to end metadata conservation in these branching functions. I just started small with this PR as a new contributor. |
|
Thank you @james-willis You have probably seen #26158 from @osipovartem -- perhaps you can help me review / take a look at that and figure out if we can find a simpler way to accomplish the same aims |
Which issue does this PR close?
CASEexpressions drop field metadata, silently stripping Arrow extension types and breaking metadata-aware UDFs #24860.Rationale for this change
COALESCE,NVL/IFNULL, andNVL2drop the field metadata of their arguments when computing their return field:return_field_from_argsbuilds the outputFieldfrom 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 argumentFields fail to plan over them (see #24860 for a full reproducer).This is one of several independent metadata drop sites listed in #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 #26145 and #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 thatCASE(indatafusion-expr) can share it: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.Expr::Castoutput field derivation: differing values ofARROW:extension:nameare a planning error, andARROW:extension:metadataalways 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_argsattaches the merged metadata of its arguments. Return type and nullability are unchanged.NVL/IFNULLdelegate tocoalesceand 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?
coerced_fields_metadatainmetadata.rs: merge precedence, untyped NULL inputs skipped, empty result, error on differing extension names, and extension name/metadata kept paired.coalesce.rsandnvl2.rsexercisereturn_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 intoCASE, so end-to-end coverage depends on theCASEfixes (#26145, #26158). This ordering is safe: the optimizer's schema invariant ignores metadata (assert_expected_schema/logically_equivalent_names_and_types), so a plan whosecoalescecarries metadata that the simplifiedCASEdoes not still plans.Are there any user-facing changes?
DataFrame::schema()) of expressions containingCOALESCE/NVL/IFNULL/NVL2over metadata-bearing columns now preserves that metadata.COALESCE/NVL/IFNULL/NVL2over arguments with differingARROW:extension:namevalues (e.g. anarrow.uuidcolumn and anarrow.jsoncolumn) is now a planning error. Previously these planned and silently dropped the extension type.datafusion_common::metadata::coerced_fields_metadata. No other API changes.