-
Notifications
You must be signed in to change notification settings - Fork 1.8k
Add Substrait roundtrip support for RecursiveQuery and recursive CTE scans
#19179
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
kosiew
wants to merge
18
commits into
apache:main
Choose a base branch
from
kosiew:substrait-recursive-16274
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Conversation
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
Implement recursive_query_rel module for both producer and consumer side. Serialize RecursiveQuery into Substrait's ExtensionMultiRel and deserialize it back. Update handling and routing in respective mod files. Add tests to verify metadata and child plan integrity during round-trip conversion. (cherry picked from commit 5b31577)
(cherry picked from commit 9969102)
(cherry picked from commit 101a68a)
Implement recursive scan extension type with encoder/decoder to distinguish recursive work-table scans. Update Substrait ReadRel production to detect CTE work tables and emit an advanced extension marker. Enhance Substrait consumer to rebuild work-table scans without catalog resolution and include an integration test for round-tripping the balances recursive CTE. (cherry picked from commit b80d55f)
(cherry picked from commit 55283a5)
(cherry picked from commit 7eed1da)
Implement validation in recursive query serialization to return a Substrait error when the query name is empty. This change prevents the creation of invalid ExtensionMultiRel payloads. Introduce a regression test to ensure that serializing a RecursiveQuery without a name fails as expected. (cherry picked from commit a11e12c)
(cherry picked from commit 0a73223)
(cherry picked from commit 23a7450)
(cherry picked from commit bf0eeb2)
Implement a unit test in the read_rel.rs file that constructs a TableScan backed by a CteWorkTable wrapped in a DefaultTableSource. The test calls from_table_scan(...) to ensure ReadRel.advanced_extension is present, validates the type_url for RECURSIVE_SCAN_TYPE_URL, and confirms the encoded payload decodes correctly to the work table name.
RecursiveQuery and recursive CTE scans
22bcc44 to
6b2b9da
Compare
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.
Which issue does this PR close?
Rationale for this change
Substrait roundtrip mode currently fails for plans that include
RecursiveQuery, resulting innot_impl_err!("Unsupported plan type: RecursiveQuery")during SQL logic tests. This prevents recursive CTE queries from being roundtripped through Substrait, causing multiplecte.sltcases to fail and reducing confidence in Substrait interoperability.This PR adds end-to-end support for serializing and deserializing
RecursiveQuerylogical plans and their associated recursive work-table scans. This unblocks the failing SQL logic tests and improves parity between the DataFusion logical plan representation and its Substrait encoding.What changes are included in this PR?
New
logical_plan::recursivehelper moduleIntroduces
RECURSIVE_QUERY_TYPE_URLandRECURSIVE_SCAN_TYPE_URLto tag recursive structures in Substrait extensions.Defines small prost messages for:
RecursiveQueryDetail { name, is_distinct }used to carryRecursiveQuerymetadata.RecursiveScanDetail { name }used to identify recursive work-table scans.Provides helpers to encode/decode these messages:
encode_recursive_query_detail/decode_recursive_query_detail.encode_recursive_scan_detail/decode_recursive_scan_detail.Adds validation so that empty names and malformed payloads are reported as Substrait errors.
Producer-side support for
RecursiveQueryAdds
logical_plan/producer/rel/recursive_query_rel.rsimplementing:from_recursive_query, which serializesLogicalPlan::RecursiveQueryintoExtensionMultiRelwith two inputs:inputs[0]:static_term.inputs[1]:recursive_term.Encodes
{ name, is_distinct }into thedetailfield usingRECURSIVE_QUERY_TYPE_URL.Wires this into the generic Substrait producer:
SubstraitProducerwithhandle_recursive_query.to_substrait_relto delegateLogicalPlan::RecursiveQuerytohandle_recursive_queryinstead of returningnot_impl_err!.Consumer-side support for
RecursiveQueryAdds
logical_plan/consumer/rel/recursive_query_rel.rsimplementing:from_recursive_query_rel, which reconstructsLogicalPlan::RecursiveQueryfromExtensionMultiRel.Validates that:
detailfield is present and decodes successfully.Rebuilds
RecursiveQuery { name, is_distinct, static_term, recursive_term }.Integrates with
DefaultSubstraitConsumer:RECURSIVE_QUERY_TYPE_URLinExtensionMultiReland routes tofrom_recursive_query_rel.Support for recursive CTE work-table scans
Producer (TableScan → ReadRel)
Detects
CteWorkTable-backed scans viaDefaultTableSource.Adds a helper
recursive_scan_name(&TableScan) -> Option<String>that:DefaultTableSourceand then toCteWorkTable.When a
CteWorkTableis detected, setsReadRel.advanced_extensionwith:type_url = RECURSIVE_SCAN_TYPE_URL.value = encode_recursive_scan_detail(name).Consumer (ReadRel → LogicalPlan::TableScan)
ReadRel.advanced_extensionand, whentype_url == RECURSIVE_SCAN_TYPE_URL, decodes the recursive scan detail.TableReferencetable name and returns a Substrait error if they differ.CteWorkTablewrapped in aTableSourceinstead of resolving a regular catalog table for recursive scans.resolve_table_reflogic for non-recursive scans.Refactoring in
ReadRelconsumerread_with_schemahelper with a more flexibleread_with_sourcehelper that accepts aTableSource(either a catalog table or aCteWorkTable).Tests
Adds unit tests for producer-side recursive scan encoding:
from_table_scan_sets_advanced_extension_for_cte_work_tableensures that:TableScanover aCteWorkTableusesRECURSIVE_SCAN_TYPE_URLinadvanced_extension.Adds roundtrip and error-coverage tests for recursive queries and scan details in
roundtrip_logical_plan.rs:roundtrip_recursive_queryverifies that aRecursiveQueryroundtrips through Substrait and back, preserving the name andis_distinctflag.serialize_recursive_query_with_empty_name_errorschecks that encoding fails with a clear error when theRecursiveQueryname is empty.decode_recursive_query_detail_malformed_bytes_errorsensures malformed bytes produce a descriptive Substrait error.decode_recursive_scan_detail_malformed_bytes_errorssimilarly validates error handling for malformed recursive scan detail.roundtrip_recursive_query_distinctconfirms thatis_distinct = trueis preserved across the roundtrip.roundtrip_recursive_query_preserves_child_plansdoes a structural sanity check that the main characteristics of the child plans (projections, filters, table scans) survive the roundtrip.roundtrip_recursive_query_with_work_table_scan_executesruns a real recursive CTE (balances example) through Substrait roundtrip and executes it, asserting non-empty results.Are these changes tested?
Before
After
Yes.
New and updated tests have been added to cover:
RecursiveQuerySubstrait serialization and deserialization.RecursiveQueryDetailandRecursiveScanDetail, including malformed-byte and empty-name error paths.CteWorkTablescans viaadvanced_extension.Existing Substrait roundtrip tests continue to run and provide regression coverage for non-recursive plans.
Are there any user-facing changes?
Behavioral improvement (no breaking API changes):
RecursiveQueryand recursive CTE work-table scans. Previously, such plans failed withnot_impl_err!("Unsupported plan type: RecursiveQuery").No public Rust API signatures are changed, and no configuration flags or SQL syntax are modified.
LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated content has been manually reviewed and tested.