fix: preserve projection metadata during optimization

[gene-bordegaray]

• Closes Projection field metadata is lost during physical plan optimization #24721

There were four ways metadata could disappear.

1. Removing a metadata-only identity projection

Consider:

ProjectionExec: i@0 AS i
  output field metadata = {"event_field": "true"}

  DataSourceExec: i
    field metadata = {}

The check to remove the projection asked:

• Is every expression a column?
• Does column 0 remain column 0?
• Does the alias match the column name?
• Does the number of columns match?

All answers yes so optimizer removed projection:

DataSourceExec: i
 metadata = {}

Metadata lost.

2 Collapsing across a metadata boundary

Consider:

ProjectionExec: arrow_metadata(i, 'event_field') AS metadata

  ProjectionExec: i@0 AS i
    output metadata = {"event_field": "true"}

    DataSourceExec: i
      metadata = {}

The correct result is true.

The previous projection colapse logic would substitute the outer expression through the inner projection:

ProjectionExec: arrow_metadata(i, 'event_field') AS metadata

 DataSourceExec: i
   metadata = {}

Now the func sees the scan field instead of the inner projection field giving use result as NULL now.

3 Rebuilding a projection with a new child

Some optimizer paths replace the child of a projection:

Old:
ProjectionExec(metadata={"key": "value"})

 OldChild

---

New:
ProjectionExec(...?)

 NewChild

The previous make_with_child implementation did this:

ProjectionExec::try_new(projection.expr().to_vec(), new_child)

where try_new derives the output schema from the expressions and new child so it woudlnt retain metadata from the original projection.

4 Cast target metadata lost before optimization

This one was a little confusing because main passed the UUID metadata test, but the first version of this PR did not (@gabotechs this is what you called out)

Basically a cast can have an explicit target field with metadata. For example, the UUID type planner produces:

FixedSizeBinary(16)
  metadata = {"ARROW:extension:name": "arrow.uuid"}

But logical cast schema only used the target data type when deriving and kept th source metadata:

Source field:
  raw: FixedSizeBinary(16)
  metadata: {}

Cast target:
  FixedSizeBinary(16)
  metadata: {"ARROW:extension:name": "arrow.uuid"}

Derived cast output:
  FixedSizeBinary(16)
  metadata: {}

This appeared in CI when common sub-expr elimination extracts a repeated cast into
its own projection:

ProjectionExec: arrow_metadata(__common_expr_1, 'ARROW:extension:name')

 ProjectionExec: CAST(raw AS UUID) AS __common_expr_1

The inner projection was initially created with incorrect empty metadata, so the physical optimizer rebuilt that projection and rederived its schema so isthe was accidentally repairing the logical schema bug.

Once this PR started preserving projection metadata correctly had this pop up this other bug.

---

So then I solve the optimizer bugs in this PR

[gabotechs]

Good catch @gene-bordegaray! just to give more context, we were bitten by this in our system while upgrading.

Just left a suggestion for relaxing the requirements, but otherwise LGTM.

[gabotechs]

🤔 there seems to be a CI failure:

1. query result mismatch:
[SQL] SELECT
    CAST(raw AS UUID),
    arrow_metadata(CAST(raw AS UUID), 'ARROW:extension:name')
FROM (
    VALUES (
        arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)')
    )
) AS uuids(raw);
[Diff] (-expected|+actual)
-   00010203040506070809000102030506 arrow.uuid
+   00010203040506070809000102030506 NULL
at /home/runner/work/datafusion/datafusion/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt:36

Do you think it's related to this change?

[gene-bordegaray]

Do you think it's related to this change?

looking into

[gene-bordegaray]

Do you think it's related to this change?

Found issues, this is a bit more involved than I was hoping. Will the variants with the fix

[gene-bordegaray]

@gabotechs ok I figured out what was going on and documented it in the PR description. There is also another bug in the codec / serialization where we need to serialize metadata. I am not solving that in this PR to keep scoped / tracked. I will crete issue for this tmrw or you can if you would like 👍

[gene-bordegaray]

this is also a correctenss issue / regression in 55 so I can note this in the minor version bump

[codecov-commenter]

Codecov Report

❌ Patch coverage is 88.16327% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.60%. Comparing base (124291e) to head (889cdfb).
⚠️ Report is 38 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/projection.rs	88.16%	14 Missing and 15 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24670      +/-   ##
==========================================
- Coverage   81.61%   81.60%   -0.02%     
==========================================
  Files        1123     1123              
  Lines      409392   411751    +2359     
  Branches   409392   411751    +2359     
==========================================
+ Hits       334134   336014    +1880     
- Misses      55637    55945     +308     
- Partials    19621    19792     +171     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[gene-bordegaray]

created codec / serialization follow up here: #24695

Still working on better fix it, so approach might change

[gene-bordegaray]

ok I pushed a change that introduces an enum to differentiation between data type and explicit field casts. It is a larger and public api change but it is what I see as properly tracking this information, not an ad hoc check

[timsaucer]

I've just been looking at this as it pertains to #24462 and it looks like it has a couple of API changes that would make it ineligible for a patch release.

[gene-bordegaray]

hey @timsaucer yes, this PR was originally meant to be stacked on #24725 but because of the breaking changes we are gong to take #23169 approach which avoids this for now. Then I will rebase this on that PR and will not hve these breaking change 👍

[timsaucer]

In addition to the comment here, mostly discovered via agent evaluation, this PR lacks some testing. I believe if you merge in gene-bordegaray#7 it will resolve the testing angle.

[timsaucer]

After some back and forth with an agent, this now seems like a reasonable issue about this check:

This is a precondition, so it exits before try_swapping_with_projection — it blocks relocation as well as removal. The 20 try_swapping_with_projection impls split into two groups:

• 7 relocate the projection via make_with_child — which this PR already made metadata-safe
• 3 absorb it (source.rs:525 covering all DataSource impls, streaming.rs:337) and 1 embeds it (try_embed_projection, covering the FilterExec fallback and all 5 joins) — these re-derive their schema and do drop the metadata

The precondition can't distinguish them, so it pays for the lossy 4 by disabling the safe 7. The catch is who that lands on: overrides_metadata() only fires for projections built through try_new_with_schema_metadata, which after #23169 means embedder-constructed plans, not SQL. So the users who need this fix are exactly the ones who lose the pushdown.

Concretely:

ProjectionExec: a@0 AS a          declared output: a -> {"unit": "ms"}
  FilterExec: a@0 > 5
    DataSourceExec: a, b, c

	metadata	filter reads
pre-PR	lost	1 col
this PR	kept	3 cols
suggestion below	kept	1 col

Would it work to make it a postcondition instead — check the result rather than deciding up front?

projection
    .input()
    .try_swapping_with_projection(projection)?
    .filter(|swapped| swapped.schema() == projection.schema())

Deleting the precondition should be free: the removal path is already covered by the projection.schema() == projection.input().schema() you added to is_projection_removable, so the precondition was only ever guarding the swap. This also means a future try_swapping_with_projection impl is covered without anyone remembering to add a guard — the failure mode becomes a lost optimization rather than lost metadata. And it adds no public API, which I think keeps this patch-eligible.

I tried it against d112760:

Check	Result
Filter relocation (safe path)	swap happens, metadata preserved
try_embed_projection (lossy path)	swap rejected, projection kept, metadata intact
Full sqllogictest suite, 510 files	zero false rejections
slt failures vs. unmodified PR head	identical (same 35 files)
projection::tests	24/24 pass, including all 4 you added

The rejection trace on the lossy path shows the metadata that would have been dropped:

SWAP_REJECTED child=FilterExec
  proj_schema=    [Field { name: "a", metadata: {"unit": "ms"} }]
  swapped_schema= [Field { name: "a" }]

[gene-bordegaray]

The cost concern here is solved with the changes that will be mae to overrides_metadata so i think we are good on that frohnt. It will now read the cached bool

For the condition removal, I don't believe we can do this. There are two checks that I think are being confused as checking the same thing:

1. In is_projection_removable we check if the schemas of the projuection and its input are a match and if they are we can remove it.
2. Then we check if we can swap the projection to be below its child, like in the filter example above. This is a differnt check and the guard is still needed.

Say you had this:

   Projection: a metadata={unit: ms}
     Filter: arrow_metadata(a, "unit") IS NULL
       Input: a,b metadata={}

After the swap:

  Filter: arrow_metadata(a, "unit") IS NULL
    Projection: a metadata={unit: ms}
      Input

The filter outputs the schema it receives so it now receives the schema with the metadata. Meaning that the proposed check will pass, but the filter now also evaluates its predicate using that metadata. Before this it saw rows with no metadata thus returns NULL and the rows pass. Then after it is inspecting it with metadata so the rows dont pass.

I don't know if this is accessible via SQL, but it defnitely can via public APIs. I added a test that checks for this regression.

[gene-bordegaray]

also so I understand, what are the testing concerns that gene-bordegaray#7 addresses? I dont see these tests exposing any of the behavior changes in this PR, more testing #23169 behavior.

[timsaucer]

I think this is very close to ready. You're right to push back on my testing questions, they were directed at the earlier PR.

[timsaucer]

Can we instead just reuse projection.project_schema(input_schema) and compare. the metadata?

[gene-bordegaray]

yes we could. The only thing is project_schema does a ton of extra stuff we don't really need. So with the past wide schema perf concerns I think just deriving exactly what we need is gonna give us better perf here

[gene-bordegaray]

@timsaucer caught a guard on an outer metadata guard that couldn't be reached, merged those in. Thank you 👍

[timsaucer]

Thank you for all the hard work on this!

[timsaucer]

@gabotechs @adriangb Any final comments or thoughts on this PR before we merge?

[adriangb]

lgtm from my end. thanks @timsaucer and @gene-bordegaray for driving this

[gene-bordegaray]

awesome thanks for all the detailed feedback on this 👍

[timsaucer]

@gene-bordegaray Can you create a branch off branch-55 and cherry pick the merge commit, create a PR targeting branch-55?

[gene-bordegaray]

@gene-bordegaray Can you create a branch off branch-55 and cherry pick the merge commit, create a PR targeting branch-55?

yes, done: #24992