feat(ffi): Rust bindings for sirius::ffi::Fragment and Context - #1702
Draft
aocsa wants to merge 1 commit into
Draft
feat(ffi): Rust bindings for sirius::ffi::Fragment and Context#1702aocsa wants to merge 1 commit into
aocsa wants to merge 1 commit into
Conversation
dev's sirius-sys binds only Context + execute_substrait; every StarRocks-CN layer imports Fragment. Bind the C++ Fragment surface already on dev (#1481) in the cxx bridge and wrap it as Fragment<'ctx>: SiriusContext::fragment(&self) hands out several fragments of one query by moving the UniquePtr behind a RefCell, execute_substrait relaxes to &self, and stream_view_name is shared so a front end and the engine agree. The result drain is result_to_arrow (matches the C++; it borrows and does not consume). Tests: broadcast/hash-key (INT64 + DECIMAL) fan-out, relay schema guard, and the three #1598 API tests. GPU-only; CI compiles them with cargo test --no-run. Supersedes #1598. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Layer 2 of the ffi stack, on top of #1697 (
stacked/ffi-transaction-scope). The C++sirius::ffi::Fragmentis already in the base branch. This PR binds it from Rust so a caller can drive a distributed plan one fragment at a time, moving native GPU batches between fragments with no Arrow or file in between. Three files change. Two are the Rust crates. The third is a three-line comment insrc/include/sirius_ffi.hppsaying whystream_view_namereturns aunique_ptr<std::string>. cxx can bind that directly, and the view-name convention then has exactly one definition that both languages read.sirius-sys. The cxx bridge gainstype Fragment,make_fragment,stream_view_name, and the fragment methodsdeclare_input_column,declare_input_sender,declare_output,declare_output_broadcast,declare_output_hash_key,build,relay_from,close_input,run,result_to_arrow,output_batch_countandoutput_types.result_to_arrowisunsafeat this level because it takes the address of a caller-ownedArrowArrayStream.Fragment,make_fragmentandstream_view_nameare re-exported from the crate root.sirius.SiriusContext::fragment(&self)returns aFragment<'ctx>tied to the context throughPhantomData<&'ctx SiriusContext>, so the compiler, rather than the C++ side, enforces that a fragment cannot outlive its engine. The safe methods take&strfor column names and DuckDB type names and&[u8]for the Substrait plan.relay_fromreturns the number of batches moved.result_to_arrowreturns aSubstraitResultand owns theFFI_ArrowArrayStreamon its own stack frame for the duration of the call, the same patternexecute_substrait_resultalready used. I pulled the drain into onecollect_arrow_streamhelper that both call.output_typesreturnsVec<String>. The crate also exportsstream_view_name(u64) -> Stringso a front end can emit a read of the view the engine will create atbuild.Why
fragment()takes&self. A distributed plan keeps several fragments alive at once. Senders sit parked until their receiver relays from them. With a&mut selffactory the secondctx.fragment()would not borrow-check. So the context now holds itsUniquePtr<Context>in aRefCelland takes the mutable borrow inside each call that needsPin<&mut Context>, releasing it before the call returns. The context is neitherSendnorSync, so there is no cross-thread aliasing to think about. That costs existing callers one signature change.execute_substraitandexecute_substrait_resultnow take&selfinstead of&mut self, and the existing test dropped itslet mut ctxto match.Naming.
result_to_arrowkeeps the C++ verb rather thaninto_*. Rust reservesinto_*for by-value conversions, and this one borrows and could in principle be called more than once.Tests. Seven new tests in
rust/crates/sirius/src/lib.rs. All but the first take the existingGPU_CONTEXT_LOCKand need a GPU.stream_view_name_matches_the_engine_convention. No GPU, no context. Assertssirius_stream_0,sirius_stream_42and theu64::MAXform.routing_modes_are_mutually_exclusive. GPU only for context bring-up. Broadcast then hash key errors, and hash key then broadcast errors. The "needs at least two destinations" rule is not checked here because it cannot be known untilbuild.context_makes_several_fragments_at_once. Two live fragments from one context, thenrelay_frombeforebuildreturnsErr. This only trips the build-ordering guard. The "source must have run" guard needs two built fragments and a real plan, so the fan-out tests cover it instead.broadcast_fragment_feeds_every_destination. A sender with outputs 0 and 1 under broadcast. Each receiver gets all three rows of the users fixture.hash_partitioned_fragment_routes_keys_deterministically. Two parquet files with the same 20000BIGINTkeys but row groups of 5000 vs 7000. Each sender's two partitions are disjoint and union to 20000, and the two independently built senders assign every key to the same stream.hash_partitioned_fragment_routes_decimal_keys_deterministically. Same contract for aDECIMAL(15,2)key, the shape of TPC-H q10's shuffle key. Precision 15 keeps the parquet physical typeINT64. The engine hashes decimals through aFLOAT64cast. What a shuffle needs from that cast is determinism, equal values landing in equal buckets on every sender, not injectivity, so the assertions match theBIGINTtest.relay_from_rejects_a_mismatched_schema.output_types()errs beforebuildand returns["BIGINT", "VARCHAR"]after. A receiver that declaresidasDOUBLEfails atrelay_from, before any batch moves, and the error namescolumn 0,DOUBLEandBIGINT.Test-only helpers:
write_users_parquet,write_multi_row_group_parquet,write_multi_row_group_decimal_parquet,stream_read_plan,decimal_stream_read_plan,stream_read_plan_f64,rows,decimal_rows.Not covered here.
declare_input_sender,close_inputandoutput_batch_countare bound but no test in this PR calls them. Bothfragment()andexecute_substraittake&self, so the Rust types allow a whole-plan execute while a fragment sits betweenbuildandrun. There is no test for that combination either.How I tested it. On a GB200 box (aarch64), a full C++ build, then cargo fmt, clippy with warnings as errors, the
siriuscrate's test suite on one GPU (10 tests, 7 of them new), and the Catch2 tag[sirius_ffi]. The seven new tests need a GPU, so CI does not run them. The barecargo test --no-runthat CI runs fails to link on this aarch64 box for a pre-existing toolchain reason (it needs an rpath-link flag for the conda libs); x64 CI is the authority for that command.Intentionally not handled. Each of these is its own layer or PR on the same two files: the exchange staging arena (
staging_*,StagingArena,PackedBatch,export_packedandpush_packed);declare_input_cardinalityandoutput_row_count(the Rust half of #1694);pin_table,unpin_tableandPinTableSpec; the byte-rangelocal_files_plan_rangedhelper and its test, which need the scan-side PRs first;fragment_over_an_empty_input_stream_terminateswithunder_watchdog, a hang-class test that lands above the pipeline-completion work; and thederive_key_cast_typeFLOAT/DOUBLE widening insrc/exec/streaming_fragment.cpp, which nothing here needs.rust/Cargo.lockis untouched and byte-identical todev; I did not take #1598's uuid downgrade. No docs change:docs/super-sirius/streaming-fragments.mdandstreaming-sessions.mdalready describe the cross-language lifecycle. One runtime caveat: with theUniquePtrbehind aRefCell, a re-entrant call into the context from inside another call panics withBorrowMutErrorinstead of failing to compile. No such caller exists in-tree or inexperimental/starrocks.What I want eyes on. The
RefCell. Every borrow is taken and released inside one call, andFragmentholds no reference to the context beyondPhantomData, so I do not see a re-entrant borrow. If you can find one, that is the review comment I most want.Checklist
References
aocsa/stream/16-rust-fragmentat 516e033, 47 commits stale). Its doc comments and three tests are carried here; its lockfile hunk is not.sirius::ffi::FragmentandContext. Every symbol this bridge binds already exists ondev.stacked/ffi-transaction-scope).