Repository navigation
feat: give PhysicalExtensionCodec the full decode context - #25318
namanjain24-sudo wants to merge 21 commits into
Conversation
PhysicalExtensionCodec::try_decode only receives the TaskContext, so a codec that decodes its own nested expressions has to build a fresh PhysicalPlanDecodeContext. That context has no scalar subquery results, so a ScalarSubqueryExpr inside an extension plan fails to deserialize even under the ScalarSubqueryExec that owns the results. Add try_decode_with_ctx, which receives the full PhysicalPlanDecodeContext and by default calls try_decode. Decode extension plans through it, forward it in ComposedPhysicalExtensionCodec, and call it with a root context on the FFI provider side.
a384c28 to
cf5431e
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25318 +/- ##
========================================
Coverage 82.85% 82.85%
========================================
Files 1151 1152 +1
Lines 456012 456485 +473
Branches 456012 456485 +473
========================================
+ Hits 377824 378222 +398
- Misses 55235 55274 +39
- Partials 22953 22989 +36 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Bumping this — still open for review whenever someone has bandwidth. |
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. The context-aware decode path looks good for local codecs, but the FFI path still loses the active scalar subquery scope. I left one blocking comment and one small test-naming suggestion.
…ndary ForeignPhysicalExtensionCodec now forwards the active ScalarSubqueryResults scope through an ABI-safe FFI_ScalarSubqueryResults handle instead of dropping it, so a ScalarSubqueryExpr decoded by a genuinely foreign codec shares the same results as the surrounding ScalarSubqueryExec. ScalarSubqueryResults gains an alternate backend representation (ScalarSubqueryResultsBackend) to support this proxy, with ScalarValues crossing the boundary as prost-encoded bytes like the rest of the codec's wire format. Adds a forced-foreign end-to-end ScalarSubqueryExec roundtrip test, and clarifies the legacy-decode test's name/intent.
|
@kosiew addressed: |
- Move the new FFI_PhysicalExtensionCodec::try_decode_with_ctx field to the end of the repr(C) struct instead of inserting it in the middle, so every pre-existing field keeps its original offset. - Add RefUnwindSafe as a supertrait bound on ScalarSubqueryResultsBackend so ScalarSubqueryResults keeps implementing UnwindSafe/RefUnwindSafe, with a regression test pinning it.
|
Fixed both: moved |
CI's clippy (ci/scripts/rust_clippy.sh, not --all-features) flagged encode_scalar_value for taking ScalarValue by value while only borrowing it. Takes &ScalarValue instead.
|
CI clippy caught |
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working through the earlier feedback. The decode-context forwarding and related cleanup look improved, but there is still a type-preservation issue in the new FFI scalar-subquery results bridge that needs to be fixed before this is safe to merge.
…uery bridge datafusion_proto::protobuf::ScalarValue has no dedicated variant for every DataType: a non-null Float16 is encoded as Float32Value, so a foreign ScalarSubqueryExpr reading it back got a Float32 instead. Wrap the encoded bytes with the original DataType (as an ArrowType) and cast back to it on decode, so the value crossing FFI always matches the type it was declared with. Adds forced-foreign non-null Float16 get/set regression coverage.
kosiew
left a comment
There was a problem hiding this comment.
Thanks for the updates. The Float16 transport issue is fixed, and the context forwarding is in place. I still have one FFI integration coverage gap that I think should be addressed before merging, plus an ABI metadata/documentation issue.
…bridge The existing unit test only forces the codec foreign through mock_foreign_marker_id, so the results handle still takes the local-bypass fast path and never exercises ForeignScalarSubqueryResultsBackend. Adds an integration-tests-gated test that decodes inside a separately dlopen'd copy of the cdylib, so the handle genuinely crosses the boundary. Shares the scalar-subquery fixtures between the unit test and the new integration test via a fixtures module gated on cfg(any(test, feature = "integration-tests")). Also flags the FFI ABI break from the appended try_decode_with_ctx field for a maintainer to label, since the struct's size changes even though existing try_decode-only implementations keep compiling.
There was a problem hiding this comment.
Thanks for the follow-up changes. The scope forwarding, Float16 type preservation, and ABI documentation have improved.
The new cross-library test confirms that the remote results backend can read host-populated values, but it still bypasses the production context-aware FFI codec callback. Please extend the test to exercise that callback with the active remote results handle, while retaining the host-value assertion.
…ss-library test The cross-library integration test decoded through a bespoke test-only entry point instead of the real FFI_PhysicalExtensionCodec callback a foreign codec actually goes through, so it never exercised try_decode_with_ctx_fn_wrapper or ForeignPhysicalExtensionCodec itself. Now it builds the codec inside the dlopen'd cdylib, converts it to Arc<dyn PhysicalExtensionCodec> (confirming it is really foreign), and calls try_decode_with_ctx on that Arc directly - the same call any caller of a foreign codec makes. The decoded plan comes back as an opaque ForeignExecutionPlan, so the host-value assertion now goes through execute() instead of a downcast: ScalarSubqueryExprExec gets a real execute() that evaluates its expression and emits the result as a RecordBatch, giving the host-populated value a way to cross back out through the FFI_ExecutionPlan/FFI_RecordBatchStream machinery every ExecutionPlan already has, without a generic FFI_PhysicalExpr wrapper. ScalarSubqueryExprExec gains a dedicated single-column PlanProperties (no longer delegating to its child's) and a constructor that computes it, used by the fixture's own decode path and both tests.
kosiew
left a comment
There was a problem hiding this comment.
Thanks for addressing the earlier feedback. The context-aware extension decoding now preserves the active scalar-subquery results scope, including composed and foreign codecs.
The additional Float16 regression tests cover type preservation across the FFI boundary, and the cross-library test verifies that a foreign plan can access host-populated results. The PR also documents the breaking FFI layout change and the need to rebuild both sides.
The changes look good to me. Approved!
…e-context-sync # Conflicts: # datafusion/proto/src/physical_plan/mod.rs
|
@thinkharderdev -- can you please review this PR as you requested / suggested the original feature? |
|
Let's put this as drat until @thinkharderdev has a chance to review and confirm this is what he had in mind. Thank you @kosiew for the review |
# Conflicts: # datafusion/ffi/integration-tests/tests/cases/ffi_physical_extension_codec.rs
…s crate main moved the FFI cross-library tests out of `datafusion/ffi/tests` and into the `datafusion-ffi-integration-tests` crate, where each case is a module under `tests/cases` that `tests/ffi_integration.rs` gates behind the `integration-tests` feature. The merge carried this file to the new directory but left it shaped for the old layout. Drop the per-file `#[cfg(feature = "integration-tests")]`, since `mod cases` is already gated, register the module in `tests/cases/mod.rs`, and name the file after its subject like its siblings. The test needs `ScalarSubqueryExpr`, so add `datafusion-physical-expr` to the crate's dev-dependencies.
|
Synced with main. The FFI cross-library tests moved into their own crate (#26564) while this sat, so the new test went to |
|
Still no word from @thinkharderdev, so this stays in draft as you suggested — no rush from my side. If there's someone else who could confirm the API shape is what was intended, happy to be pointed at them; otherwise I'll leave it parked. |
|
Thanks for the flexibility @namanjain24-sudo This is a case where I want to make sure that if we add an API it is actually usable by users / isn't going to just bitrot / be unused. I appreciate you sticking with it |
|
Added the 56.0.0 upgrade-guide entry for the @thinkharderdev — ping when you have a moment: this is the decode context from #25089, and @kosiew has approved it. @alamb if someone else is a better reviewer, happy to redirect. |
|
Apologies, I somehow missed this but will review today |
Upstream Issue: apache#25089 Decoding `ScalarSubqueryExpr` inside an extension `ExecutionPlan` node fails because we don't have the full context.
thinkharderdev
left a comment
There was a problem hiding this comment.
I'm out of the loop on the FFI stuff but it looks like it has been reviewed already. The other changes lgtm though. Thanks!
|
Synced with main. The license header check had gone red: main switched to the short SPDX header in #27049, so this PR's two new files carry that form now. |
|
Thanks @namanjain24-sudo @thinkharderdev and @kosiew |
alamb
left a comment
There was a problem hiding this comment.
Hi @namanjain24-sudo -- thank you for this PR and thank you for the reviews @thinkharderdev and @kosiew
I went through some of it and I think this PR is much too verbose for what it is doing. There are many seemingly AI generated comments that restate the changes in the PR or re-describe the code for no additional gain. I also have a few questions on the design
Can you please go through (yourself) and clean up the comments?
Also, I left some additional questions -- could you please help me understand this a bit more? There are several design decisions that seem like they are pretty narrow and not well thought through (like ScalarValueEnvelope)
|
|
||
| See [issue #24626](https://github.com/apache/datafusion/issues/24626) for details. | ||
|
|
||
| ### `PhysicalExtensionCodec` can decode with the full decode context |
There was a problem hiding this comment.
actually since we cut the brach, this should now go into the 57 upgrade guide
|
|
||
| ### `PhysicalExtensionCodec` can decode with the full decode context | ||
|
|
||
| `PhysicalExtensionCodec::try_decode` received a `TaskContext`, which carries no |
There was a problem hiding this comment.
I din't think we have to reiterate the problem here -- we need to just say it didn't have enough context and then refer to the issue for more details
| - `ComposedPhysicalExtensionCodec` forwards `try_decode_with_ctx` to the codec | ||
| that wrote the payload, so composing codecs preserves the context. | ||
|
|
||
| **FFI users must rebuild both sides.** `FFI_PhysicalExtensionCodec` gained a |
There was a problem hiding this comment.
is this specific to this change? It seem generic to all FFI
| } | ||
| ``` | ||
|
|
||
| **Who is affected:** |
There was a problem hiding this comment.
Can you please trim this down from AI speak
For example this first sentence doesn't need to restate how the code works, it can just say "no changes are required"
The second can say "if you need the additional context, use the new function"
The detail about ComposedPhysical... is an implementation detail that isn't relevant for implmenters, is it?
| `datafusion_ffi::version()` reports changes with it: a provider and host that | ||
| compare it, which is what that function is for, see the mismatch instead. | ||
|
|
||
| `ScalarSubqueryResults` is no longer backed only by in-process state. It can be |
There was a problem hiding this comment.
is this important to state? This just reiterates how things are implemented. I think we could remove most/all of this last paragraph
| DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, | ||
| }; | ||
|
|
||
| /// The decode context's active scalar-subquery results scope must reach |
There was a problem hiding this comment.
Please cut this comment down to explain the purpose of the test, rahter than restating in english how it is implemented
| /// external integration-test binary, which links this crate as an ordinary | ||
| /// dependency rather than compiling it with `cfg(test)`, can still reach | ||
| /// these types through `datafusion_ffi::proto::physical_extension_codec::fixtures`. | ||
| #[cfg(any(test, feature = "integration-tests"))] |
There was a problem hiding this comment.
Why is this in the main ffi crate (and not in the test binary?)? This doesn't make sense to me
| /// always decodes with a root context (no scope), so the embedded | ||
| /// `ScalarSubqueryExpr` fails to deserialize. | ||
| /// | ||
| /// This only mocks the foreign path in-process (see |
There was a problem hiding this comment.
similar, please trim this comment to describe what this is testing, not restate how it is implemented
Also, why is this test needed if we already have a integration test?
| /// A stable struct for sharing the shared results container of a | ||
| /// `ScalarSubqueryExec` across FFI boundaries. | ||
| /// | ||
| /// Unlike most FFI wrappers in this crate, which stand in for a whole |
There was a problem hiding this comment.
Why is it important that this is different? Please clarify this comment
There was a problem hiding this comment.
And remove details about its implementation -- instead focus on its relevance
| /// The wire form a [`ScalarValue`] crosses the FFI boundary in. | ||
| /// | ||
| /// `datafusion_proto::protobuf::ScalarValue` alone is not a type-preserving | ||
| /// transport: it has no dedicated variant for every `DataType` (`Float16`, |
There was a problem hiding this comment.
THis seems like a rationale for expanding the coverage of datafusion_proto::protobuf::ScalarValue, ot adding a whole new wrapper
|
Also the comments say "give the full decode context" but this seems mostly to be about adding an API to pass scalar subquery results around 🤔 |
|
Thanks for the detailed review. The comments were helpful. I've gone through each one and the PR is now 210 lines smaller. The enum: Removed. The in-crate Fixtures in the FFI crate: These need to stay in the crate because the test loads this crate's cdylib and asks it for the codec. The cdylib needs to be able to build the fixtures itself; types defined only in the test binary aren't available inside it. I've narrowed the gate from The upgrade note: Moved to a new The remaining comments: Trimmed them down to explain the purpose of each thing: the FFI field, the cross-library test, and There are two things I'd still like your input on:
The framing: I agree that the PR's framing could be clearer. The trait change itself is 42 changed lines in |
alamb
left a comment
There was a problem hiding this comment.
thanks -- this is getting closer but still has a lot of unecessary complexity -- please re-read the whole thing before asking for a re-review and either call out or clean up anything you don't understand
|
|
||
| /// An extension plan and codec for the cross-library integration test. | ||
| /// | ||
| /// These live in this crate because the test `dlopen`s this crate's cdylib and |
There was a problem hiding this comment.
I don't understand this
Why can't it be in datafusion/ffi/integration-tests/tests/cases/physical_extension_codec.rs ?
This seems a lot of code to be testing the availablity of subquery contexts over FFI
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { |
There was a problem hiding this comment.
what do these tests cover that the integration tests do not?
| FFI_ExecutionPlan::new(plan, None) | ||
| } | ||
|
|
||
| /// Builds a [`FFI_PhysicalExtensionCodec`] wrapping |
There was a problem hiding this comment.
Have you read these comments carefully? It seems like they are restating the details of how the function works -- before asking for my review again, please carefully read all the code and comments your LLM has written and determine if it is needed or not
|
Thanks - that was a fair call. I went back through the whole diff rather than just the three spots, and removed 55 lines of comments plus the code below. Net 175 lines out, 20 in. The comments. You were right that they were restating the code. The block you pointed at on I also stopped exposing the fixtures' internals - The two tests in One honest gap: the other test also exercised The fixtures' location. They cannot live in The cost is that On the volume. The fixtures module is 195 lines, and about 100 of those are the plan itself. Its And the structural version of the same question, which I asked last time and is still open: the trait change itself is 42 lines in |
Which issue does this PR close?
PhysicalExtensionCodecshould take aPhysicalPlanDecodeContextinstead of aTaskContext#25089.Rationale for this change
A custom
ExecutionPlanwhose codec decodes its own expressions fails to deserialize when one of those expressions is aScalarSubqueryExpr. It fails even when the plan sits inside theScalarSubqueryExecthat owns the results:The decoder does track the results of the enclosing
ScalarSubqueryExecinPhysicalPlanDecodeContext. ButPhysicalExtensionCodec::try_decodeonly receivesctx.task_ctx(), so a codec has to build a freshPhysicalPlanDecodeContext::new(..)to decode its expressions, and that fresh context has no results.What changes are included in this PR?
This takes the non-breaking option from the issue:
PhysicalExtensionCodec::try_decode_with_ctxreceives the fullPhysicalPlanDecodeContext. Its default implementation callstry_decode(ctx.task_ctx()), so existing codecs compile and behave as before.try_decode_with_ctx.ComposedPhysicalExtensionCodecforwardstry_decode_with_ctxto the codec that encoded the node. Without this, composing codecs would drop the context again.FFI_PhysicalExtensionCodecgains atry_decode_with_ctxfunction pointer that also forwards the caller's active scalar-subquery results scope (if any) across the boundary through a newFFI_ScalarSubqueryResultshandle, so aScalarSubqueryExprdecoded by a foreign codec shares the same populated results as the host'sScalarSubqueryExec. A codec that only implements the new method therefore also works through FFI. For codecs that only implementtry_decode, nothing changes.I went with the new method rather than changing
try_decode's signature because of the API health policy. The breaking change is simpler if reviewers prefer it, and I'm happy to switch.This touches the same function as the draft #24631, which adds a per-type registry in front of the codec. That PR's registry path already passes decoders the full context (through
ConverterPlanDecoder), but its codec fallback still callstry_decode(ctx.task_ctx()), which is the line this PR changes. Whichever lands second needs a one-line rebase.What is the testing strategy for this PR?
New tests in
datafusion/proto/tests/cases/plans/scalar_subquery.rsuse an extension plan whose codec decodes its own expressions throughtry_decode_with_ctx. Itstry_decodereturns an error, so a caller that drops the context fails the test:ScalarSubqueryExprinside the extension plan shares the enclosingScalarSubqueryExec's results, with bothDefaultPhysicalProtoConverterandDeduplicatingProtoConverterComposedPhysicalExtensionCodecScalarSubqueryExecs: each expression binds to its own scope and not to the otherScalarSubqueryExec, decoding still fails with the existing errorA new FFI test checks that a codec that only decodes through
try_decode_with_ctxworks behindFFI_PhysicalExtensionCodec. A second, cross-library integration test (datafusion/ffi/tests/ffi_physical_extension_codec.rs, gated by theintegration-testsfeature) builds that codec inside a separatelydlopen'd copy of thedatafusion-fficdylib, so the scalar-subquery results handle crosses a genuine FFI boundary and exercisesForeignScalarSubqueryResultsBackendfor real, rather than only the in-processmock_foreign_marker_idpath the unit test uses.I checked that each part of the change is needed:
mainwith the error from the issue.try_decode, all four new proto tests fail.ComposedPhysicalExtensionCodec::try_decode_with_ctxremoved, only the composed-codec test fails.Existing suites pass:
datafusion-proto(17 lib, 265 integration, 4 doc tests) anddatafusion-ffi(123 lib tests, plus the new cross-library test under--features integration-tests)../dev/rust_lint.shis clean.Are there any user-facing changes?
This has two separate kinds of API impact:
PhysicalExtensionCodec::try_decode_with_ctxis a new provided method. Existing implementations keep compiling and behave the same.FFI_PhysicalExtensionCodecis#[repr(C)]and gains atry_decode_with_ctxfunction-pointer field, appended after the existing fields so a producer compiled against the old layout still sees every field it knows about at its original offset. The struct's size still changes, so this is an ABI break: a producer and a consumer built against different layouts of this struct must not be mixed — both sides need to be rebuilt against the same DataFusion version. This should carry theapi changelabel; I don't have permission to add labels on this repo, so flagging it here for a maintainer to apply.The same two-sided split applies to the
FFI_ScalarSubqueryResultshandle added for the scalar-subquery results forwarding described above.