-
Notifications
You must be signed in to change notification settings - Fork 1.9k
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
base: main
Are you sure you want to change the base?
Conversation
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
|
hi @gabotechs |
|
Sure! will take a look soon |
|
|
||
| // Convert child plans | ||
| let static_term = Arc::new(consumer.consume_rel(&rel.inputs[0]).await?); | ||
| let recursive_term = Arc::new(consumer.consume_rel(&rel.inputs[1]).await?); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice work 🚀
I know we're only getting to this function with RECURSIVE_QUERY_TYPE_URL, but I think validating that the recursive term references the work table would be good.
Since the Substrait spec doesn't yet have a native recursive relation and ExtensionMultiRel is pretty general purpose for complex operations, a defensive check here would help catch invalid plans.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added the checks
ccb6c72 to
2424884
Compare
Convert sql_interval_to_expr to a thin wrapper that delegates to sql_interval_to_expr_impl. This change encapsulates the original logic in the new function, preserving backward compatibility with the existing public method signature. Existing callers remain unaffected by this refactor.
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.