From b82c123f699c74f702feb32722dc1b06f346d898 Mon Sep 17 00:00:00 2001 From: Brian Myers Date: Wed, 11 Mar 2026 18:11:11 -0500 Subject: [PATCH 01/15] turn off submodule updating --- .github/actions/setup-submodules/action.yaml | 15 ++++++++ .github/workflows/dependencies.yml | 3 +- .github/workflows/docs_pr.yaml | 3 +- .github/workflows/extended.yml | 37 ++++++++++++++++++-- .github/workflows/rust.yml | 36 ++++++++++++------- .gitmodules | 3 ++ 6 files changed, 80 insertions(+), 17 deletions(-) create mode 100644 .github/actions/setup-submodules/action.yaml diff --git a/.github/actions/setup-submodules/action.yaml b/.github/actions/setup-submodules/action.yaml new file mode 100644 index 0000000000000..28f3c08052b88 --- /dev/null +++ b/.github/actions/setup-submodules/action.yaml @@ -0,0 +1,15 @@ +name: "Setup Submodules" +description: "Initialize and update git submodules for testing" +runs: + using: "composite" + steps: + - name: Initialize and update submodules + shell: bash + run: | + git config --global --add safe.directory $GITHUB_WORKSPACE + git submodule init + # Override update=none setting for CI + git config submodule.testing.update checkout + git config submodule.parquet-testing.update checkout + git config submodule.datafusion-testing.update checkout + git submodule update --depth 1 diff --git a/.github/workflows/dependencies.yml b/.github/workflows/dependencies.yml index 2f3a127ef98c4..ab8803ede053d 100644 --- a/.github/workflows/dependencies.yml +++ b/.github/workflows/dependencies.yml @@ -43,8 +43,9 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: diff --git a/.github/workflows/docs_pr.yaml b/.github/workflows/docs_pr.yaml index 4b8d25b0611eb..7726ccd0ea142 100644 --- a/.github/workflows/docs_pr.yaml +++ b/.github/workflows/docs_pr.yaml @@ -44,8 +44,9 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup uv uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 - name: Install doc dependencies diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index a143cb49fd35b..eaed0b215cb2d 100644 --- a/.github/workflows/extended.yml +++ b/.github/workflows/extended.yml @@ -57,6 +57,34 @@ permissions: checks: write jobs: + + # Check crate compiles and base cargo check passes + linux-build-lib: + name: linux build test + runs-on: ${{ github.repository_owner == 'apache' && format('runs-on={0},family=m8a,cpu=8,image=ubuntu24-full-x64,extras=s3-cache,disk=large,tag=datafusion', github.run_id) || 'ubuntu-latest' }} + # note: do not use amd/rust container to preserve disk space + steps: + - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # v2.0.3 + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push + fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules + - name: Install Rust + run: | + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + source $HOME/.cargo/env + rustup toolchain install + - name: Install Protobuf Compiler + run: | + sudo apt-get update + sudo apt-get install -y protobuf-compiler + - name: Prepare cargo build + run: | + cargo check --profile ci --all-targets + cargo clean + # Run extended tests (with feature 'extended_tests') linux-test-extended: name: cargo test 'extended_tests' (amd64) @@ -67,8 +95,9 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Free Disk Space (Ubuntu) uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1 - name: Install Rust @@ -114,8 +143,9 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -136,8 +166,9 @@ jobs: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules # Don't use setup-builder to avoid configuring RUST_BACKTRACE which is expensive - name: Install protobuf compiler run: | diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 5ff1f6467bbf1..569f83c519d46 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -279,8 +279,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -325,8 +326,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain run: rustup toolchain install stable - name: Rust Dependency Cache @@ -357,8 +359,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -388,8 +391,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -451,8 +455,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -499,8 +504,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -524,8 +530,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -567,8 +574,9 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-macos-aarch64-builder - name: Run datafusion-ffi tests @@ -660,8 +668,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -685,8 +694,9 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -707,8 +717,9 @@ jobs: - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -742,8 +753,9 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - submodules: true fetch-depth: 1 + - name: Setup Submodules + uses: ./.github/actions/setup-submodules - name: Mark repository as safe for git # Required for git commands inside container (avoids "dubious ownership" error) diff --git a/.gitmodules b/.gitmodules index 037accdbe4241..14f7ce2ec08d4 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,10 +1,13 @@ [submodule "parquet-testing"] path = parquet-testing url = https://github.com/apache/parquet-testing.git + update = none [submodule "testing"] path = testing url = https://github.com/apache/arrow-testing + update = none [submodule "datafusion-testing"] path = datafusion-testing url = https://github.com/apache/datafusion-testing.git branch = main + update = none From 854e21899f5a403c11669497155bbf9ea5e39c82 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Tue, 21 Apr 2026 21:36:50 +0200 Subject: [PATCH 02/15] Fix ArrayCompact incompatibility with branch-53 --- datafusion/functions-nested/src/array_compact.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/datafusion/functions-nested/src/array_compact.rs b/datafusion/functions-nested/src/array_compact.rs index 11be494b5b20f..d4d0f12830d28 100644 --- a/datafusion/functions-nested/src/array_compact.rs +++ b/datafusion/functions-nested/src/array_compact.rs @@ -32,6 +32,7 @@ use datafusion_expr::{ Volatility, }; use datafusion_macros::user_doc; +use std::any::Any; use std::sync::Arc; make_udf_expr_and_func!( @@ -85,6 +86,10 @@ impl ScalarUDFImpl for ArrayCompact { "array_compact" } + fn as_any(&self) -> &dyn Any { + self + } + fn signature(&self) -> &Signature { &self.signature } From 38c2b397d09b837e0cf8742917e937f6270b345a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=ADa=20Adriana?= Date: Fri, 22 May 2026 14:03:29 +0200 Subject: [PATCH 03/15] fix(substrait): dedupe names of aggregate measures, not just groupings (#22453) (#126) ## Which issue does this PR close? - Closes #. ## Rationale for this change When the substrait consumer hits an `Aggregate` with two identical measures (e.g. `sum(a)` present twice), planning fails with `Schema contains duplicate unqualified field name`. Substrait carries column names at the plan root rather than on the measures themselves, so the measures arrive at `Aggregate` schema construction without aliases -- and two identical exprs produce two identical field names. PR #20539 fixed the `NameTracker` to dedupe duplicate names in the consumer, but it was only applied to grouping expressions, not to the measures. The planner sees: ``` field 1: (qualifier: None, name: "sum(data.a)") field 2: (qualifier: None, name: "sum(data.a)") ``` which is rejected when constructing the Aggregate's output schema. ## What changes are included in this PR? Run aggregate measures through the same `NameTracker` like the grouping expressions in `from_aggregate_rel` ## Are these changes tested? Yes -- added a roundtrip test `aggregate_identical_measures`. Without the fix it produces `Error: SchemaError(DuplicateUnqualifiedField { name: "sum(data.a)" }, Some(""))` ## Are there any user-facing changes? No. (cherry picked from commit 097efae26c7d5ca8e6f124f27545a79eb227636f) --- .../consumer/rel/aggregate_rel.rs | 12 +- .../tests/cases/roundtrip_logical_plan.rs | 21 ++++ ...ggregate_identical_measures.substrait.json | 103 ++++++++++++++++++ 3 files changed, 133 insertions(+), 3 deletions(-) create mode 100644 datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index ac7d2479c397a..413ee4b537c29 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -109,11 +109,17 @@ pub async fn from_aggregate_rel( aggr_exprs.push(std::sync::Arc::unwrap_or_clone(agg_func?)); } - // Ensure that all expressions have a unique name + // Ensure that all expressions have a unique name. Both grouping and + // aggregate expressions become fields in the aggregate's output schema, + // so they share a single namespace. let mut name_tracker = NameTracker::new(); let group_exprs = group_exprs - .iter() - .map(|e| name_tracker.get_uniquely_named_expr(e.clone())) + .into_iter() + .map(|e| name_tracker.get_uniquely_named_expr(e)) + .collect::, _>>()?; + let aggr_exprs = aggr_exprs + .into_iter() + .map(|e| name_tracker.get_uniquely_named_expr(e)) .collect::, _>>()?; input.aggregate(group_exprs, aggr_exprs)?.build() diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 1d65256d76420..872c2d0cd2a81 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -1116,6 +1116,27 @@ async fn aggregate_identical_grouping_expressions() -> Result<()> { Ok(()) } +#[tokio::test] +async fn aggregate_identical_measures() -> Result<()> { + // Two identical aggregate measures share the same schema_name; without + // NameTracker dedup over measures, building the Aggregate's output + // DFSchema fails with "Schema contains duplicate unqualified field name". + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_identical_measures.substrait.json", + ); + + let plan = generate_plan_from_substrait(proto_plan).await?; + assert_snapshot!( + plan, + @r" + Projection: __common_expr_1 AS sum_a_1, __common_expr_1 AS sum(data.a)__temp__0 AS sum_a_2 + Aggregate: groupBy=[[]], aggr=[[sum(data.a) AS __common_expr_1]] + TableScan: data projection=[a] + " + ); + Ok(()) +} + #[tokio::test] async fn simple_intersect_consume() -> Result<()> { let proto_plan = read_json("tests/testdata/test_plans/intersect.substrait.json"); diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json new file mode 100644 index 0000000000000..620d55e93ee1e --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json @@ -0,0 +1,103 @@ +{ + "extensionUris": [{ + "extensionUriAnchor": 1, + "uri": "/functions_arithmetic.yaml" + }], + "extensions": [{ + "extensionFunction": { + "extensionUriReference": 1, + "functionAnchor": 0, + "name": "sum:i64" + } + }], + "relations": [{ + "root": { + "input": { + "aggregate": { + "common": { + "direct": {} + }, + "input": { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": ["a"], + "struct": { + "types": [{ + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": ["data"] + } + } + }, + "groupings": [{ + "groupingExpressions": [] + }], + "measures": [ + { + "measure": { + "functionReference": 0, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "invocation": "AGGREGATION_INVOCATION_ALL", + "arguments": [{ + "value": { + "selection": { + "directReference": { + "structField": { + "field": 0 + } + }, + "rootReference": {} + } + } + }] + } + }, + { + "measure": { + "functionReference": 0, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "invocation": "AGGREGATION_INVOCATION_ALL", + "arguments": [{ + "value": { + "selection": { + "directReference": { + "structField": { + "field": 0 + } + }, + "rootReference": {} + } + } + }] + } + } + ] + } + }, + "names": ["sum_a_1", "sum_a_2"] + } + }], + "version": { + "minorNumber": 54, + "producer": "manual" + } +} From 826f5c6c5a0b6de78b20b154704ac69fb618083e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=ADa=20Adriana?= Date: Wed, 27 May 2026 11:09:30 +0200 Subject: [PATCH 04/15] Add lambda substrait support (#21193) (#134) Part of #21172 Substrait support wasn't implemented in the core lambda support to reduce PR size Substrait consuming and producing of higher-order functions, lambdas and lambda variables Unit tests added to `datafusion/substrait/tests/cases/roundtrip_logical_plan.rs` None --------- (cherry picked from commit 9a6f67e202423ec1feee256045f6aa256a04daae) (cherry picked from commit 1ac2df106955382ee843f2adc030174d2f6b9492) Co-authored-by: gstvg <28798827+gstvg@users.noreply.github.com> Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Co-authored-by: Ben Bellick <36523439+benbellick@users.noreply.github.com> --- datafusion-testing | 2 +- .../consumer/expr/field_reference.rs | 88 +++- .../src/logical_plan/consumer/expr/lambda.rs | 102 ++++ .../src/logical_plan/consumer/expr/mod.rs | 9 +- .../consumer/expr/scalar_function.rs | 24 +- .../consumer/substrait_consumer.rs | 210 ++++++++- .../src/logical_plan/producer/expr/lambda.rs | 48 ++ .../producer/expr/lambda_variable.rs | 49 ++ .../src/logical_plan/producer/expr/mod.rs | 10 +- .../producer/expr/scalar_function.rs | 78 +++- .../producer/substrait_producer.rs | 214 ++++++++- .../substrait/tests/cases/logical_plans.rs | 22 + .../tests/cases/roundtrip_logical_plan.rs | 242 +++++++++- datafusion/substrait/tests/cases/serialize.rs | 154 +++++- .../test_plans/higher_order_function.json | 438 ++++++++++++++++++ 15 files changed, 1658 insertions(+), 32 deletions(-) create mode 100644 datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs create mode 100644 datafusion/substrait/src/logical_plan/producer/expr/lambda.rs create mode 100644 datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs create mode 100644 datafusion/substrait/tests/testdata/test_plans/higher_order_function.json diff --git a/datafusion-testing b/datafusion-testing index 7833a65d5b08b..13bbae38776c2 160000 --- a/datafusion-testing +++ b/datafusion-testing @@ -1 +1 @@ -Subproject commit 7833a65d5b08be2ca484ea938f471cf01df54e18 +Subproject commit 13bbae38776c2bfbc1fab1be7e7220222d4284bf diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs index dae6c625ef55b..aab9bde949203 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs @@ -21,7 +21,7 @@ use datafusion::logical_expr::Expr; use std::sync::Arc; use substrait::proto::expression::FieldReference; use substrait::proto::expression::field_reference::ReferenceType::DirectReference; -use substrait::proto::expression::field_reference::RootType; +use substrait::proto::expression::field_reference::{LambdaParameterReference, RootType}; use substrait::proto::expression::reference_segment::ReferenceType::StructField; pub async fn from_field_reference( @@ -56,9 +56,9 @@ pub(crate) fn from_substrait_field_reference( Some(RootType::Expression(_)) => not_impl_err!( "Expression root type in field reference is not supported" ), - Some(RootType::LambdaParameterReference(_)) => not_impl_err!( - "Lambda parameter reference in field reference is not yet supported" - ), + Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )) => consumer.lambda_variable(*steps_out as usize, field_idx), } } _ => not_impl_err!( @@ -85,3 +85,83 @@ fn resolve_outer_reference( let col = Column::from((qualifier, field)); Ok(Expr::OuterReferenceColumn(Arc::clone(field), col)) } + +#[cfg(test)] +mod tests { + use datafusion::{ + common::{DFSchema, assert_contains}, + prelude::SessionContext, + }; + use substrait::proto::{ + Type, + expression::{ + FieldReference, ReferenceSegment, + field_reference::{self, LambdaParameterReference, RootType}, + reference_segment::{ReferenceType, StructField}, + }, + r#type::{I64, Kind}, + }; + + use crate::{ + extensions::Extensions, + logical_plan::consumer::{ + DefaultSubstraitConsumer, SubstraitConsumer, from_field_reference, + }, + }; + + #[tokio::test] + async fn test_lambda_variable_invalid_steps_out() { + let lambda_field_ref = lambda_field_ref(0, 99); + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_field_reference(&consumer, &lambda_field_ref, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!(err.to_string(), "No lambda at 99 steps out, got only 0"); + } + + #[tokio::test] + async fn test_lambda_variable_invalid_field_idx() { + let lambda_field_ref = lambda_field_ref(1, 0); + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + let _names = consumer + .push_lambda_parameters( + &[Type { + kind: Some(Kind::I64(I64::default())), + }], + &DFSchema::empty(), + ) + .unwrap(); + + let err = from_field_reference(&consumer, &lambda_field_ref, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "At lambda 0 steps out, no field at index 1, got only 1" + ); + } + + fn lambda_field_ref(field: i32, steps_out: u32) -> FieldReference { + FieldReference { + reference_type: Some(field_reference::ReferenceType::DirectReference( + ReferenceSegment { + reference_type: Some(ReferenceType::StructField(Box::new( + StructField { field, child: None }, + ))), + }, + )), + root_type: Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )), + } + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs new file mode 100644 index 0000000000000..2fb71024b054f --- /dev/null +++ b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs @@ -0,0 +1,102 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::{ + common::{DFSchema, substrait_err}, + prelude::{Expr, lambda}, +}; +use substrait::proto; + +use crate::logical_plan::consumer::SubstraitConsumer; + +pub async fn from_lambda( + consumer: &impl SubstraitConsumer, + expr: &proto::expression::Lambda, + input_schema: &DFSchema, +) -> datafusion::common::Result { + let Some(parameters) = expr.parameters.as_ref() else { + return substrait_err!("Lambda expression without parameters is not allowed"); + }; + + let names = consumer.push_lambda_parameters(¶meters.types, input_schema)?; + + let Some(body) = expr.body.as_ref() else { + return substrait_err!("Lambda expression without body is not allowed"); + }; + + let body = consumer.consume_expression(body, input_schema).await?; + + consumer.pop_lambda_parameters(); + + Ok(lambda(names, body)) +} + +#[cfg(test)] +mod tests { + use datafusion::{ + common::{DFSchema, assert_contains}, + prelude::SessionContext, + }; + use substrait::proto::{self, Expression, r#type::Struct}; + + use crate::{ + extensions::Extensions, + logical_plan::consumer::{DefaultSubstraitConsumer, from_lambda}, + }; + + #[tokio::test] + async fn test_lambda_without_body() { + let lambda = proto::expression::Lambda { + parameters: Some(Struct::default()), + body: None, + }; + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_lambda(&consumer, &lambda, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Lambda expression without body is not allowed" + ); + } + + #[tokio::test] + async fn test_lambda_without_parameters() { + let lambda = proto::expression::Lambda { + parameters: None, + body: Some(Box::new(Expression::default())), + }; + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_lambda(&consumer, &lambda, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Lambda expression without parameters is not allowed" + ); + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs index 295456e95f9f3..2fcc11f4e417d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs @@ -20,6 +20,7 @@ mod cast; mod field_reference; mod function_arguments; mod if_then; +mod lambda; mod literal; mod nested; mod scalar_function; @@ -32,6 +33,7 @@ pub use cast::*; pub use field_reference::*; pub use function_arguments::*; pub use if_then::*; +pub use lambda::*; pub use literal::*; pub use nested::*; pub use scalar_function::*; @@ -95,8 +97,11 @@ pub async fn from_substrait_rex( RexType::DynamicParameter(expr) => { consumer.consume_dynamic_parameter(expr, input_schema).await } - RexType::Lambda(_) | RexType::LambdaInvocation(_) => { - not_impl_err!("Lambda expressions are not yet supported") + RexType::Lambda(lambda) => { + consumer.consume_lambda(lambda.as_ref(), input_schema).await + } + RexType::LambdaInvocation(_) => { + not_impl_err!("Lambda invocations are not supported") } }, None => substrait_err!("Expression must set rex_type: {expression:?}"), diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs index 1a0fb3f55f609..4cd856fc562e8 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs @@ -30,7 +30,6 @@ pub async fn from_scalar_function( f: &ScalarFunction, input_schema: &DFSchema, ) -> Result { - //TODO: handle higher order functions, as they are also encoded as scalar functions let Some(fn_signature) = consumer .get_extensions() .functions @@ -45,6 +44,20 @@ pub async fn from_scalar_function( let fn_name = substrait_fun_name(fn_signature); let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?; + let higher_order_func = consumer + .get_function_registry() + .higher_order_function(fn_name) + .or_else(|e| { + if let Some(alt_name) = substrait_to_df_name(fn_name) { + consumer + .get_function_registry() + .higher_order_function(alt_name) + .or(Err(e)) + } else { + Err(e) + } + }); + let udf_func = consumer.get_function_registry().udf(fn_name).or_else(|e| { if let Some(alt_name) = substrait_to_df_name(fn_name) { consumer.get_function_registry().udf(alt_name).or(Err(e)) @@ -53,9 +66,14 @@ pub async fn from_scalar_function( } }); - // try to first match the requested function into registered udfs, then built-in ops + // try to first match the requested function into registered higher-order functions, then udfs, built-in ops // and finally built-in expressions - if let Ok(func) = udf_func { + if let Ok(func) = higher_order_func { + Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + func.to_owned(), + args, + ))) + } else if let Ok(func) = udf_func { Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf( func.to_owned(), args, diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 65bc53ce0834e..266916756e017 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -23,21 +23,24 @@ use super::{ from_substrait_rex, from_window_function, }; use crate::extensions::Extensions; +use crate::logical_plan::consumer::{from_lambda, from_substrait_type_without_names}; use async_trait::async_trait; -use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; use datafusion::catalog::TableProvider; use datafusion::common::{ DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err, }; use datafusion::execution::{FunctionRegistry, SessionState}; +use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use std::collections::VecDeque; use std::sync::{Arc, RwLock}; -use substrait::proto; use substrait::proto::expression as substrait_expression; use substrait::proto::expression::{ Enum, FieldReference, IfThen, Literal, MultiOrList, Nested, ScalarFunction, SingularOrList, SwitchExpression, WindowFunction, }; +use substrait::proto::{self, Type}; use substrait::proto::{ AggregateRel, ConsistentPartitionWindowRel, CrossRel, DynamicParameter, ExchangeRel, Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel, @@ -62,17 +65,19 @@ use substrait::proto::{ /// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; /// # use std::sync::Arc; /// # use substrait::proto; -/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel}; +/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel, Type}; /// # use datafusion::arrow::datatypes::DataType; /// # use datafusion::logical_expr::expr::ScalarFunction; /// # use datafusion_substrait::extensions::Extensions; /// # use datafusion_substrait::logical_plan::consumer::{ -/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer +/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer, DefaultSubstraitLambdaConsumer /// # }; /// /// struct CustomSubstraitConsumer { /// extensions: Arc, /// state: Arc, +/// // You can reuse existing consumer code related to lambdas +/// lambda_consumer: DefaultSubstraitLambdaConsumer, /// } /// /// #[async_trait] @@ -95,6 +100,30 @@ use substrait::proto::{ /// self.state.as_ref() /// } /// +/// fn push_lambda_parameters( +/// &self, +/// lambda_parameters: &[Type], +/// input_schema: &DFSchema, +/// ) -> datafusion::common::Result> { +/// self.lambda_consumer.push_lambda_parameters( +/// self, +/// lambda_parameters, +/// input_schema, +/// ) +/// } +/// +/// fn pop_lambda_parameters(&self) { +/// self.lambda_consumer.pop_lambda_parameters(); +/// } +/// +/// fn lambda_variable( +/// &self, +/// steps_out: usize, +/// field_idx: usize, +/// ) -> datafusion::common::Result { +/// self.lambda_consumer.lambda_variable(steps_out, field_idx) +/// } +/// /// // You can reuse existing consumer code to assist in handling advanced extensions /// async fn consume_project(&self, rel: &ProjectRel) -> Result { /// let df_plan = from_project_rel(self, rel).await?; @@ -384,6 +413,14 @@ pub trait SubstraitConsumer: Send + Sync + Sized { )) } + async fn consume_lambda( + &self, + expr: &proto::expression::Lambda, + input_schema: &DFSchema, + ) -> datafusion::common::Result { + from_lambda(self, expr, input_schema).await + } + // Outer Schema Stack // These methods manage a stack of outer schemas for correlated subquery support. // When entering a subquery, the enclosing query's schema is pushed onto the stack. @@ -481,6 +518,35 @@ pub trait SubstraitConsumer: Send + Sync + Sized { }; substrait_err!("Missing handler for user-defined literals {}", type_ref) } + + // Lambda related methods + + /// Push the given lambda parameters onto the stack when entering a lambda and + /// returns the names they got assigned + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it + fn push_lambda_parameters( + &self, + _lambda_parameters: &[Type], + _input_schema: &DFSchema, + ) -> datafusion::common::Result> { + not_impl_err!("SubstraitConsumer::push_lambda_parameters") + } + + /// Pop lambda parameters from the stack when leaving a lambda. + fn pop_lambda_parameters(&self) {} + + /// Returns an expression corresponding to the lambda variable with the given field_idx within the lambda it originates from, + /// at the lambda `step_outs` of the current scope + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it + fn lambda_variable( + &self, + _steps_out: usize, + _field_idx: usize, + ) -> datafusion::common::Result { + not_impl_err!("SubstraitConsumer::lambda_variable") + } } /// Default SubstraitConsumer for converting standard Substrait without user-defined extensions. @@ -490,6 +556,7 @@ pub struct DefaultSubstraitConsumer<'a> { pub(super) extensions: &'a Extensions, pub(super) state: &'a SessionState, outer_schemas: RwLock>>, + lambda_consumer: DefaultSubstraitLambdaConsumer, } impl<'a> DefaultSubstraitConsumer<'a> { @@ -498,6 +565,7 @@ impl<'a> DefaultSubstraitConsumer<'a> { extensions, state, outer_schemas: RwLock::new(Vec::new()), + lambda_consumer: DefaultSubstraitLambdaConsumer::new(), } } } @@ -594,6 +662,140 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; Ok(LogicalPlan::Extension(Extension { node: plan })) } + + fn push_lambda_parameters( + &self, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> datafusion::common::Result> { + self.lambda_consumer + .push_lambda_parameters(self, lambda_parameters, input_schema) + } + + fn pop_lambda_parameters(&self) { + self.lambda_consumer.pop_lambda_parameters() + } + + fn lambda_variable( + &self, + steps_out: usize, + field_idx: usize, + ) -> datafusion::common::Result { + self.lambda_consumer.lambda_variable(steps_out, field_idx) + } +} + +/// Default implementation of lambda related methods of the [SubstraitConsumer] trait +/// +/// Can be embedded into a custom [SubstraitConsumer] to implement them +pub struct DefaultSubstraitLambdaConsumer { + inner: RwLock, +} + +struct DefaultSubstraitLambdaConsumerInner { + /// Parameters of the lambdas currently in scope, ordered from innermost + /// to outermost. Index 0 is the lambda being consumed; higher indices + /// are enclosing lambdas, matching the `steps_out` value used by + /// [`DefaultSubstraitLambdaConsumer::lambda_variable`] and `LambdaParameterReference`. + lambda_parameters: VecDeque>, + next_lambda_parameter: usize, +} + +impl Default for DefaultSubstraitLambdaConsumer { + fn default() -> Self { + Self::new() + } +} + +impl DefaultSubstraitLambdaConsumer { + pub fn new() -> Self { + Self { + inner: RwLock::new(DefaultSubstraitLambdaConsumerInner { + lambda_parameters: VecDeque::new(), + next_lambda_parameter: 0, + }), + } + } + + pub fn push_lambda_parameters( + &self, + consumer: &impl SubstraitConsumer, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> datafusion::common::Result> { + let mut inner = self.inner.write().unwrap(); + + let lambda_parameters: Vec = lambda_parameters + .iter() + .map(|ty| { + let (assigned_number, default_name) = + next_lambda_parameter_name(inner.next_lambda_parameter, input_schema); + + inner.next_lambda_parameter = assigned_number + 1; + + let data_type = from_substrait_type_without_names(consumer, ty)?; + Ok(Arc::new(Field::new(default_name, data_type, true))) + }) + .collect::>>()?; + + let names = lambda_parameters.iter().map(|f| f.name().clone()).collect(); + + inner.lambda_parameters.push_front(lambda_parameters); + + Ok(names) + } + + pub fn pop_lambda_parameters(&self) { + self.inner.write().unwrap().lambda_parameters.pop_front(); + } + + pub fn lambda_variable( + &self, + steps_out: usize, + field_idx: usize, + ) -> datafusion::common::Result { + let lambda_parameters = &self.inner.read().unwrap().lambda_parameters; + + let Some(lambda_parameters) = lambda_parameters.get(steps_out) else { + return substrait_err!( + "No lambda at {steps_out} steps out, got only {}", + lambda_parameters.len() + ); + }; + + let Some(var) = lambda_parameters.get(field_idx) else { + return substrait_err!( + "At lambda {steps_out} steps out, no field at index {field_idx}, got only {}", + lambda_parameters.len() + ); + }; + + Ok(Expr::LambdaVariable(LambdaVariable::new( + var.name().clone(), + Some(Arc::clone(var)), + ))) + } +} + +/// Returns the next available lambda parameter name and the index it was assigned. +/// +/// Names follow the pattern `pN` where `N` starts at `next_lambda_parameter`. If `pN` +/// conflicts with an existing column name in `input_schema`, `N` is incremented until +/// a free name is found. +fn next_lambda_parameter_name( + mut next_lambda_parameter: usize, + input_schema: &DFSchema, +) -> (usize, String) { + loop { + let name = format!("p{next_lambda_parameter}"); + + // avoid conflicts with column names + if !input_schema.has_column_with_unqualified_name(&name) { + return (next_lambda_parameter, name); + } + + next_lambda_parameter += 1; + } } #[cfg(test)] diff --git a/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs new file mode 100644 index 0000000000000..0d32dab2ccadd --- /dev/null +++ b/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs @@ -0,0 +1,48 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::{common::DFSchemaRef, logical_expr::expr::Lambda}; +use substrait::proto::{ + Expression, + expression::RexType, + r#type::{Nullability, Struct}, +}; + +use crate::logical_plan::producer::SubstraitProducer; + +pub fn from_lambda( + producer: &mut impl SubstraitProducer, + lambda: &Lambda, + schema: &DFSchemaRef, +) -> Result { + Ok(Expression { + rex_type: Some(RexType::Lambda(Box::new( + substrait::proto::expression::Lambda { + parameters: Some(Struct { + nullability: Nullability::Required as i32, + type_variation_reference: 0, + types: lambda + .params + .iter() + .map(|p| producer.lambda_parameter_type(p)) + .collect::>()?, + }), + body: Some(Box::new(producer.handle_expr(&lambda.body, schema)?)), + }, + ))), + }) +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs b/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs new file mode 100644 index 0000000000000..3d7f06e2332a1 --- /dev/null +++ b/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs @@ -0,0 +1,49 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::logical_expr::expr::LambdaVariable; +use substrait::proto::{ + Expression, + expression::{ + FieldReference, ReferenceSegment, RexType, + field_reference::{LambdaParameterReference, ReferenceType, RootType}, + reference_segment::{self, StructField}, + }, +}; + +use crate::logical_plan::producer::SubstraitProducer; + +pub fn from_lambda_variable( + producer: &mut impl SubstraitProducer, + lambda_variable: &LambdaVariable, + _schema: &datafusion::common::DFSchema, +) -> Result { + let (steps_out, field) = producer.lambda_variable(&lambda_variable.name)?; + + Ok(Expression { + rex_type: Some(RexType::Selection(Box::new(FieldReference { + reference_type: Some(ReferenceType::DirectReference(ReferenceSegment { + reference_type: Some(reference_segment::ReferenceType::StructField( + Box::new(StructField { field, child: None }), + )), + })), + root_type: Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )), + }))), + }) +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs index 6e053f0d90a96..c728af2f1458d 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs @@ -19,6 +19,8 @@ mod aggregate_function; mod cast; mod field_reference; mod if_then; +mod lambda; +mod lambda_variable; mod literal; mod placeholder; mod scalar_function; @@ -30,6 +32,8 @@ pub use aggregate_function::*; pub use cast::*; pub use field_reference::*; pub use if_then::*; +pub use lambda::*; +pub use lambda_variable::*; pub use literal::*; pub use placeholder::*; pub use scalar_function::*; @@ -154,10 +158,8 @@ pub fn to_substrait_rex( Expr::HigherOrderFunction(expr) => { producer.handle_higher_order_function(expr, schema) } - Expr::Lambda(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"), - Expr::LambdaVariable(expr) => { - not_impl_err!("Cannot convert {expr:?} to Substrait") - } + Expr::Lambda(expr) => producer.handle_lambda(expr, schema), + Expr::LambdaVariable(expr) => producer.handle_lambda_variable(expr, schema), } } diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs index e36d5128cd293..e7dd2af13f9ca 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs @@ -16,7 +16,10 @@ // under the License. use crate::logical_plan::producer::{SubstraitProducer, to_substrait_literal_expr}; -use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err}; +use datafusion::common::datatype::FieldExt; +use datafusion::common::{ + DFSchemaRef, ScalarValue, internal_datafusion_err, not_impl_err, substrait_err, +}; use datafusion::logical_expr::{Between, BinaryExpr, Expr, Like, Operator, expr}; use substrait::proto::expression::{RexType, ScalarFunction}; use substrait::proto::function_argument::ArgType; @@ -35,7 +38,78 @@ pub fn from_higher_order_function( fun: &expr::HigherOrderFunction, schema: &DFSchemaRef, ) -> datafusion::common::Result { - from_function(producer, fun.name(), &fun.args, schema) + let mut lambda_parameters = fun.lambda_parameters(schema)?.into_iter(); + + let num_lambdas = fun + .args + .iter() + .filter(|arg| matches!(arg, Expr::Lambda(_))) + .count(); + + if lambda_parameters.len() != num_lambdas { + return substrait_err!( + "{} returned {} lambdas but {num_lambdas} expected", + fun.name(), + lambda_parameters.len() + ); + } + + let arguments = fun + .args + .iter() + .map(|arg| { + let arg = match arg { + Expr::Lambda(l) => { + let lambda_parameters = + lambda_parameters.next().ok_or_else(|| { + internal_datafusion_err!( + "lambda_parameters len should have been checked above" + ) + })?; + + if l.params.len() > lambda_parameters.len() { + return substrait_err!( + "Lambda defined {} parameters ({}) but function {} supports only {}", + l.params.len(), + l.params.join(","), + fun.name(), + lambda_parameters.len() + ) + } + + let named_lambda_parameters = + std::iter::zip(&l.params, lambda_parameters) + .map(|(name, parameter)| parameter.renamed(name)) + .collect(); + + producer.push_lambda_parameters(named_lambda_parameters)?; + + let arg = producer.handle_lambda(l, schema); + + producer.pop_lambda_parameters()?; + + arg + } + _ => producer.handle_expr(arg, schema), + }?; + + Ok(FunctionArgument { + arg_type: Some(ArgType::Value(arg)), + }) + }) + .collect::>()?; + + let function_anchor = producer.register_function(fun.name().to_string()); + #[expect(deprecated)] + Ok(Expression { + rex_type: Some(RexType::ScalarFunction(ScalarFunction { + function_reference: function_anchor, + arguments, + output_type: None, + options: vec![], + args: vec![], + })), + }) } fn from_function( diff --git a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs index 4228c32435897..8f1d5bc6ac475 100644 --- a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs +++ b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs @@ -20,18 +20,23 @@ use crate::logical_plan::producer::{ from_aggregate, from_aggregate_function, from_alias, from_between, from_binary_expr, from_case, from_cast, from_column, from_distinct, from_empty_relation, from_exists, from_filter, from_higher_order_function, from_in_list, from_in_subquery, from_join, - from_like, from_limit, from_literal, from_placeholder, from_projection, - from_repartition, from_scalar_function, from_scalar_subquery, from_set_comparison, - from_sort, from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr, - from_union, from_values, from_window, from_window_function, to_substrait_rel, - to_substrait_rex, + from_lambda, from_lambda_variable, from_like, from_limit, from_literal, + from_placeholder, from_projection, from_repartition, from_scalar_function, + from_scalar_subquery, from_set_comparison, from_sort, from_subquery_alias, + from_table_scan, from_try_cast, from_unary_expr, from_union, from_values, + from_window, from_window_function, to_substrait_rel, to_substrait_rex, + to_substrait_type, +}; +use datafusion::arrow::datatypes::FieldRef; +use datafusion::common::{ + Column, DFSchemaRef, HashMap, ScalarValue, not_impl_err, substrait_err, }; -use datafusion::common::{Column, DFSchemaRef, ScalarValue, substrait_err}; use datafusion::execution::SessionState; use datafusion::execution::registry::SerializerRegistry; use datafusion::logical_expr::Subquery; use datafusion::logical_expr::expr::{ - Alias, Exists, InList, InSubquery, Placeholder, SetComparison, WindowFunction, + Alias, Exists, InList, InSubquery, Lambda, LambdaVariable, Placeholder, + SetComparison, WindowFunction, }; use datafusion::logical_expr::{ Aggregate, Between, BinaryExpr, Case, Cast, Distinct, EmptyRelation, Expr, Extension, @@ -57,16 +62,19 @@ use substrait::proto::{ /// # use std::sync::Arc; /// # use substrait::proto::{Expression, Rel}; /// # use substrait::proto::rel::RelType; +/// # use datafusion::arrow::datatypes::FieldRef; /// # use datafusion::common::DFSchemaRef; /// # use datafusion::error::Result; /// # use datafusion::execution::SessionState; /// # use datafusion::logical_expr::{Between, Extension, Projection}; /// # use datafusion_substrait::extensions::Extensions; -/// # use datafusion_substrait::logical_plan::producer::{from_projection, SubstraitProducer}; +/// # use datafusion_substrait::logical_plan::producer::{from_projection, SubstraitProducer, DefaultSubstraitLambdaProducer, lambda_parameters_map}; /// /// struct CustomSubstraitProducer { /// extensions: Extensions, /// state: Arc, +/// // You can reuse existing producer code related to lambdas +/// lambda_producer: DefaultSubstraitLambdaProducer, /// } /// /// impl SubstraitProducer for CustomSubstraitProducer { @@ -83,6 +91,33 @@ use substrait::proto::{ /// self.extensions /// } /// +/// fn push_lambda_parameters( +/// &mut self, +/// lambda_parameters: Vec, +/// ) -> datafusion::common::Result<()> { +/// let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?; +/// +/// self.lambda_producer +/// .push_lambda_parameters(lambda_parameters_map); +/// +/// Ok(()) +/// } +/// +/// fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { +/// self.lambda_producer.pop_lambda_parameters() +/// } +/// +/// fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { +/// self.lambda_producer.lambda_variable(name) +/// } +/// +/// fn lambda_parameter_type( +/// &self, +/// name: &str, +/// ) -> datafusion::common::Result { +/// self.lambda_producer.lambda_parameter_type(name) +/// } +/// /// // You can set additional metadata on the Rels you produce /// fn handle_projection(&mut self, plan: &Projection) -> Result> { /// let mut rel = from_projection(self, plan)?; @@ -405,11 +440,65 @@ pub trait SubstraitProducer: Send + Sync + Sized { ) -> datafusion::common::Result { from_placeholder(self, placeholder) } + + fn handle_lambda( + &mut self, + lambda: &Lambda, + schema: &DFSchemaRef, + ) -> datafusion::common::Result { + from_lambda(self, lambda, schema) + } + + fn handle_lambda_variable( + &mut self, + lambda_variable: &LambdaVariable, + schema: &DFSchemaRef, + ) -> datafusion::common::Result { + from_lambda_variable(self, lambda_variable, schema) + } + + // Lambda related methods + + /// Push the given `lambda_parameters` into this producer so they can be referenced by lambda variables + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn push_lambda_parameters( + &mut self, + _lambda_parameters: Vec, + ) -> datafusion::common::Result<()> { + not_impl_err!("SubstraitProducer::push_lambda_parameters") + } + + /// Pop the last pushed `lambda_parameters` so that it unshadow any previously shadowed lambda parameter + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + not_impl_err!("SubstraitProducer::pop_lambda_parameters") + } + + /// Get the (`steps_out`, `field_idx`) of the lambda variable with the given `name`. `steps_out` refers to the number + /// of lambda boundaries to traverse (0 = current lambda), and `field_idx` refers to the index within the lambda parameters + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn lambda_variable(&self, _name: &str) -> datafusion::common::Result<(u32, i32)> { + not_impl_err!("SubstraitProducer::lambda_variable") + } + + /// Get the type of the lambda parameter with the given `name` + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn lambda_parameter_type( + &self, + _name: &str, + ) -> datafusion::common::Result { + not_impl_err!("SubstraitProducer::lambda_parameter_type") + } } pub struct DefaultSubstraitProducer<'a> { extensions: Extensions, serializer_registry: &'a dyn SerializerRegistry, + lambda_producer: DefaultSubstraitLambdaProducer, } impl<'a> DefaultSubstraitProducer<'a> { @@ -417,6 +506,7 @@ impl<'a> DefaultSubstraitProducer<'a> { DefaultSubstraitProducer { extensions: Extensions::default(), serializer_registry: state.serializer_registry().as_ref(), + lambda_producer: DefaultSubstraitLambdaProducer::new(), } } } @@ -471,4 +561,112 @@ impl SubstraitProducer for DefaultSubstraitProducer<'_> { rel_type: Some(rel_type), })) } + + fn push_lambda_parameters( + &mut self, + lambda_parameters: Vec, + ) -> datafusion::common::Result<()> { + let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?; + + self.lambda_producer + .push_lambda_parameters(lambda_parameters_map); + + Ok(()) + } + + fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + self.lambda_producer.pop_lambda_parameters() + } + + fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { + self.lambda_producer.lambda_variable(name) + } + + fn lambda_parameter_type( + &self, + name: &str, + ) -> datafusion::common::Result { + self.lambda_producer.lambda_parameter_type(name) + } +} + +/// Default implementation of lambda related methods of the [SubstraitProducer] trait +/// +/// Can be embedded into a custom [SubstraitProducer] to implement them +pub struct DefaultSubstraitLambdaProducer { + lambdas_variables: Vec>, +} + +impl Default for DefaultSubstraitLambdaProducer { + fn default() -> Self { + Self::new() + } +} + +impl DefaultSubstraitLambdaProducer { + pub fn new() -> Self { + Self { + lambdas_variables: Vec::new(), + } + } + + /// Note you can construct the `lambda_parameters` argument using [lambda_parameters_map] + pub fn push_lambda_parameters( + &mut self, + lambda_parameters: HashMap, + ) { + self.lambdas_variables.push(lambda_parameters); + } + + pub fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + match self.lambdas_variables.pop() { + Some(_) => Ok(()), + None => substrait_err!("no lambda_parameters to pop"), + } + } + + pub fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { + for (steps_out, lambda_parameters) in + self.lambdas_variables.iter().rev().enumerate() + { + if let Some((field_idx, _type)) = lambda_parameters.get(name) { + return Ok((steps_out as u32, *field_idx as i32)); + } + } + + substrait_err!("unknown lambda variable {name}") + } + + pub fn lambda_parameter_type( + &self, + name: &str, + ) -> datafusion::common::Result { + for lambda_parameters in self.lambdas_variables.iter().rev() { + if let Some((_field_idx, type_)) = lambda_parameters.get(name) { + return Ok(type_.clone()); + } + } + + substrait_err!("unknown lambda variable {name}") + } +} + +/// Produces a map of lambda parameters as expected by [DefaultSubstraitLambdaProducer::push_lambda_parameters] +pub fn lambda_parameters_map( + producer: &mut impl SubstraitProducer, + lambda_parameters: Vec, +) -> datafusion::common::Result> { + lambda_parameters + .into_iter() + .enumerate() + .map(|(field_idx, field)| { + Ok(( + field.name().clone(), + ( + field_idx, + to_substrait_type(producer, field.data_type(), field.is_nullable())?, + ), + )) + }) + .collect::>() } diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 663a372fe2e4f..d4ac01462c879 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -19,6 +19,7 @@ #[cfg(test)] mod tests { + use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; use datafusion::common::test_util::format_batches; use std::collections::HashSet; @@ -293,4 +294,25 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn higher_order_function() -> Result<()> { + let proto_plan = + read_json("tests/testdata/test_plans/higher_order_function.json"); + // ctx already contains the queried table + let ctx = higher_order_function_ctx().await?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_concat(array_transform2(p0, (p3, p4) -> p3 * p2 * p4), array_transform2(p0, (p5, p6) -> p5 * p2 * p6))) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j))) + TableScan: data3 + " + ); + + // Trigger execution to ensure plan validity + DataFrame::new(ctx.state(), plan).show().await?; + Ok(()) + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 872c2d0cd2a81..ebe4eb5ee7a28 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -17,8 +17,14 @@ use crate::utils::test::read_json; use datafusion::arrow::array::ArrayRef; +use datafusion::config::Dialect; +use datafusion::execution::FunctionRegistry; use datafusion::functions_nested::map::map; -use datafusion::logical_expr::LogicalPlanBuilder; +use datafusion::logical_expr::{ + ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, LogicalPlanBuilder, + ValueOrLambda, +}; use datafusion::physical_plan::Accumulator; use datafusion::scalar::ScalarValue; use datafusion_substrait::logical_plan::{ @@ -27,7 +33,9 @@ use datafusion_substrait::logical_plan::{ use std::cmp::Ordering; use std::mem::size_of_val; -use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit}; +use datafusion::arrow::datatypes::{ + DataType, Field, FieldRef, IntervalUnit, Schema, TimeUnit, +}; use datafusion::common::tree_node::Transformed; use datafusion::common::{DFSchema, DFSchemaRef, Spans, not_impl_err, plan_err}; use datafusion::error::Result; @@ -1941,6 +1949,215 @@ async fn roundtrip_placeholder_typed_utf8() -> Result<()> { assert_eq!(plan.schema(), plan2.schema()); Ok(()) + +#[tokio::test] +async fn roundtrip_array_transform_higher_order_function() -> Result<()> { + let ctx = higher_order_function_ctx().await?; + + // simple + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], p0 -> p0 * 2) from data3", + ctx.clone(), + ) + .await?; + + // dont use the parameter + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], p0 -> 3) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters using both + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> p0 * p2) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters only last + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> 2 * p2) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters use none + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> 3) from data3", + ctx.clone(), + ) + .await?; + + // nested without variable shadowing + roundtrip_with_ctx("SELECT array_transform2([[data3.p1]], p0 -> array_transform2(p0, p2 -> p2 * 2)) from data3", ctx.clone()) + .await?; + + // nested with multiple parameters without variable shadowing + roundtrip_with_ctx("SELECT array_transform2([[data3.p1]], (p0, p2) -> array_transform2(p0, (p3, p4) -> p2 * p3 * p4)) from data3", ctx.clone()) + .await?; + + // since substrait doesn't encode lambda parameters names, they got generated, non-conflicting names during consumption + // testing name shadowing requires to assert against the generated plan and check the correct parameter usage instead of round tripping + + // nested with variable shadowing. + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2([[data3.p1]], v -> array_transform2(v, v -> v * 2)) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_transform2(p0, (p2) -> p2 * Int64(2))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_transform2(v,(v) -> v * Int64(2))) + TableScan: data3 projection=[p1] + " + ); + + // nested with variable shadowing with multiple parameters + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2([[data3.p1]], (v, i) -> array_transform2(v, (v, i) -> v * i)) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_transform2(p0, (p3, p4) -> p3 * p4)) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_transform2(v,(v, i) -> v * i)) + TableScan: data3 projection=[p1] + " + ); + + // nested with variable shadowing and later reuse of the shadowed var after exiting the shadowing expression + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2( + [[data3.p1]], + v -> array_concat( + -- when entering this expression, inner v is pushed into the producer and shadows outer v, but after exiting this, + -- it should be removed and unshadow the outer v, so that it can be used in the next expression + array_transform2(v, v -> v * 2), + array_transform2(v, v -> v * 2) + ) + ) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_concat(array_transform2(p0, (p2) -> p2 * Int64(2)), array_transform2(p0, (p3) -> p3 * Int64(2)))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_concat(array_transform2(v,(v) -> v * Int64(2)),array_transform2(v,(v) -> v * Int64(2)))) + TableScan: data3 projection=[p1] + " + ); + + Ok(()) +} + +pub(crate) async fn higher_order_function_ctx() -> Result { + let mut ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; + + ctx.register_higher_order_function(Arc::new(ArrayTransform::new()))?; + + let data3_fields = vec![ + Field::new("p1", DataType::Int64, true), // lambda parameters should not conflict with this column + ]; + let data3 = Schema::new(data3_fields); + let mut data3_options = CsvReadOptions::new(); + data3_options.schema = Some(&data3); + data3_options.has_header = false; + ctx.register_csv("data3", "tests/testdata/empty.csv", data3_options) + .await?; + + Ok(ctx) +} + +// todo use core array_transform when it supports multiple lambda parameters +#[derive(Debug, PartialEq, Eq, Hash)] +struct ArrayTransform { + signature: HigherOrderSignature, +} + +impl ArrayTransform { + fn new() -> Self { + Self { + signature: HigherOrderSignature::variadic_any(Volatility::Immutable), + } + } +} + +impl HigherOrderUDF for ArrayTransform { + fn name(&self) -> &str { + "array_transform2" + } + + fn aliases(&self) -> &[String] { + &[] + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(_)] = fields else { + unreachable!() + }; + + let field = match list.data_type() { + DataType::List(field) => field, + _ => unreachable!(), + }; + + Ok(LambdaParametersProgress::Complete(vec![vec![ + Arc::clone(field), + Arc::new(Field::new("", DataType::Int64, true)), + ]])) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)] = args.arg_fields + else { + unreachable!() + }; + + let field = Arc::new(Field::new( + Field::LIST_FIELD_DEFAULT_NAME, + lambda.data_type().clone(), + lambda.is_nullable(), + )); + + let return_type = match list.data_type() { + DataType::List(_) => DataType::List(field), + _ => unreachable!(), + }; + + Ok(Arc::new(Field::new("", return_type, list.is_nullable()))) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + // this function is only tested with roundtrip_with_ctx, which only prints the output + // and generate_plan_from_sql_with_ctx which doesn't execute nothing, so the output doesn't matter + Ok(ColumnarValue::Scalar(ScalarValue::new_default( + args.return_type(), + )?)) + } + } fn check_post_join_filters(rel: &Rel) -> Result<()> { @@ -2084,6 +2301,15 @@ async fn generate_plan_from_sql( optimized: bool, ) -> Result { let ctx = create_context().await?; + generate_plan_from_sql_with_ctx(sql, assert_schema, optimized, &ctx).await +} + +async fn generate_plan_from_sql_with_ctx( + sql: &str, + assert_schema: bool, + optimized: bool, + ctx: &SessionContext, +) -> Result { let df: DataFrame = ctx.sql(sql).await?; let plan = if optimized { @@ -2437,8 +2663,18 @@ async fn roundtrip_all_types(sql: &str) -> Result<()> { } async fn create_context() -> Result { + create_context_with_dialect(None).await +} + +async fn create_context_with_dialect(dialect: Option) -> Result { + let mut session_config = SessionConfig::default(); + + if let Some(dialect) = dialect { + session_config.options_mut().sql_parser.dialect = dialect; + } + let mut state = SessionStateBuilder::new() - .with_config(SessionConfig::default()) + .with_config(session_config) .with_runtime_env(Arc::new(RuntimeEnv::default())) .with_default_features() .with_serializer_registry(Arc::new(MockSerializerRegistry)) diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 2d7257fad3394..1981ef66db377 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -28,9 +28,16 @@ mod tests { use insta::assert_snapshot; use std::fs; + use substrait::proto::expression::field_reference::{ReferenceType, RootType}; + use substrait::proto::expression::reference_segment; + use substrait::proto::expression::{ReferenceSegment, RexType}; + use substrait::proto::function_argument::ArgType; use substrait::proto::plan_rel::RelType; use substrait::proto::rel_common::{Emit, EmitKind}; - use substrait::proto::{RelCommon, rel}; + use substrait::proto::r#type::{I64, Kind as TypeKind, List, Nullability, Struct}; + use substrait::proto::{Expression, RelCommon, Type, rel}; + + use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; #[tokio::test] async fn serialize_to_file() -> Result<()> { @@ -196,6 +203,101 @@ mod tests { panic!("plan did not match expected structure") } + #[tokio::test] + async fn higher_order_function() -> Result<()> { + let ctx = higher_order_function_ctx().await?; + let df = ctx + .sql( + "SELECT array_transform2( + [[data3.p1]], + (v, i) -> array_concat( + -- when entering this expression, inner v is pushed into the producer and shadows outer v, but after exiting this, + -- it should be removed and unshadow the outer v, so that it can be used in the next expression + array_transform2(v, (v, j) -> v * i * j), + array_transform2(v, (v, j) -> v * i * j) + ) + ) from data3" + ) + .await?; + let datafusion_plan = df.into_optimized_plan()?; + let plan = to_substrait_plan(&datafusion_plan, &ctx.state())? + .as_ref() + .clone(); + + let relation = plan.relations.first().unwrap().rel_type.as_ref(); + let root_rel = match relation { + Some(RelType::Root(root)) => root.input.as_ref().unwrap(), + _ => panic!("expected Root"), + }; + + let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() else { + panic!("expected Project at top of plan") + }; + + let mut params = vec![]; + let mut lambda_refs = vec![]; + + collect_lambda_ref(&p.expressions[0], &mut params, &mut lambda_refs); + + let nullable_i64 = Type { + kind: Some(TypeKind::I64(I64 { + type_variation_reference: 0, + nullability: Nullability::Nullable as i32, + })), + }; + + let inner_lambda_struct = Struct { + // v, j + types: vec![nullable_i64.clone(); 2], + type_variation_reference: 0, + nullability: Nullability::Required as i32, + }; + + assert_eq!( + params, + vec![ + Struct { + types: vec![ + // v + Type { + kind: Some(TypeKind::List(Box::new(List { + r#type: Some(Box::new(nullable_i64.clone())), + type_variation_reference: 0, + nullability: Nullability::Nullable as i32 + }))) + }, + // i + nullable_i64, + ], + type_variation_reference: 0, + nullability: Nullability::Required as i32, + }, + inner_lambda_struct.clone(), + inner_lambda_struct, + ] + ); + + assert_eq!( + lambda_refs, + vec![ + // first inner array_transform2 argument: outer v + (0, 0), + // first inner lambda body: v * i * j + (0, 0), + (1, 1), + (0, 1), + // second inner array_transform2 argument: outer v + (0, 0), + // second inner lambda body: v * i * j + (0, 0), + (1, 1), + (0, 1), + ] + ); + + Ok(()) + } + fn assert_emit(rel_common: Option<&RelCommon>, output_mapping: Vec) { assert_eq!( rel_common.unwrap().emit_kind.clone(), @@ -211,4 +313,54 @@ mod tests { .await?; Ok(ctx) } + + // Recursively walks a expression tree depth-first, collecting in visit order: + // - `params`: the parameter struct of each Lambda encountered + // - `lambda_refs`: every field reference whose root is a LambdaParameterReference, + // recorded as (steps_out, field_index) so tests can assert which enclosing + // lambda each reference resolves to and which parameter within it. + fn collect_lambda_ref( + expr: &Expression, + params: &mut Vec, + lambda_refs: &mut Vec<(u32, i32)>, + ) { + if let Some(rex_type) = &expr.rex_type { + match rex_type { + RexType::Selection(field_reference) => { + if let ( + Some(ReferenceType::DirectReference(ReferenceSegment { + reference_type: + Some(reference_segment::ReferenceType::StructField( + struct_field, + )), + })), + Some(RootType::LambdaParameterReference(lambda_param_ref)), + ) = (&field_reference.reference_type, &field_reference.root_type) + { + lambda_refs.push((lambda_param_ref.steps_out, struct_field.field)) + } + } + RexType::ScalarFunction(scalar_function) => { + for arg in &scalar_function.arguments { + match &arg.arg_type { + Some(ArgType::Value(value)) => { + collect_lambda_ref(value, params, lambda_refs) + } + _ => unreachable!(), + } + } + } + RexType::Lambda(lambda) => { + if let Some(parameters) = &lambda.parameters { + params.push(parameters.clone()); + } + if let Some(body) = &lambda.body { + collect_lambda_ref(body, params, lambda_refs); + } + } + RexType::Literal(_literal) => {} + _ => unreachable!(), + } + } + } } diff --git a/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json b/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json new file mode 100644 index 0000000000000..da613b2573447 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json @@ -0,0 +1,438 @@ +{ + "version": { + "minorNumber": 85, + "producer": "datafusion" + }, + "extensions": [ + { + "extensionFunction": { + "extensionUrnReference": 2, + "functionAnchor": 2, + "name": "array_transform2" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 2, + "name": "make_array" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 2, + "functionAnchor": 3, + "name": "array_concat" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 1, + "functionAnchor": 1, + "name": "multiply" + } + } + ], + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 1 + ] + } + }, + "input": { + "read": { + "baseSchema": { + "names": [ + "p1" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "projection": { + "select": { + "structItems": [ + {} + ] + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + }, + "expressions": [ + { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "rootReference": {} + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 3, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": { + "stepsOut": 1 + } + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": {} + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": { + "stepsOut": 1 + } + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": {} + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + ] + } + }, + "names": [ + "array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j)))" + ] + } + } + ], + "extensionUrns": [ + { + "extensionUrnAnchor": 1, + "urn": "extension:io.substrait:functions_arithmetic" + }, + { + "extensionUrnAnchor": 2, + "urn": "extension:io.substrait:functions_list" + } + ] +} From 76d436155bc01bf202805a5e6f2bf38c0c33c37c Mon Sep 17 00:00:00 2001 From: Pablo Abad Rubio Date: Tue, 2 Jun 2026 15:17:19 +0200 Subject: [PATCH 05/15] fix: make PushDownLeafProjections work with unnest (#22620) --- .../optimizer/src/extract_leaf_expressions.rs | 51 +++++++++++++++++++ datafusion/sqllogictest/test_files/unnest.slt | 44 ++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index c5c5610aeaed9..185f9d045f10f 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -1143,6 +1143,13 @@ fn try_push_into_inputs( return Ok(None); } + // Unnest may output a column with the same name but different value/type + // than its input column. Name-based routing cannot distinguish those. + // On top of that Unnest can't go through the `node.with_new_exprs(node.expressions(), new_inputs)` rebuild + if matches!(node, LogicalPlan::Unnest(_)) { + return Ok(None); + } + // SubqueryAlias remaps qualifiers between input and output. // Rewrite pairs/columns from alias-space to input-space before routing. let remapped = if let LogicalPlan::SubqueryAlias(sa) = node { @@ -3035,4 +3042,48 @@ mod tests { Ok(()) } + + /// Regression test for the `Assertion failed: expr.is_empty(): Unnest` + /// internal error. + /// + /// `try_push_into_inputs` rebuilds the parent node via + /// `node.with_new_exprs(node.expressions(), new_inputs)`. For `Unnest`, + /// `apply_expressions` exposes the `exec_columns` as `Expr::Column`s + /// (so `expressions()` is **non-empty**), but `with_new_exprs` for + /// `Unnest` immediately calls `assert_no_expressions(expr)?` and errors + /// out. The optimizer should treat `Unnest` as a barrier and bail + /// instead of attempting to push through it. + #[test] + fn test_no_push_through_unnest() -> Result<()> { + use arrow::datatypes::{DataType, Field, Schema}; + + let schema = Schema::new(vec![ + Field::new("list_col", DataType::new_list(DataType::Int32, true), true), + Field::new("other_col", DataType::Int32, true), + ]); + let table_scan = + datafusion_expr::logical_plan::table_scan(Some("t"), &schema, None)? + .build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .unnest_column("list_col")? + .filter(leaf_udf(col("list_col"), "x").eq(lit(1i32)))? + .build()?; + + let ctx = OptimizerContext::new().with_max_passes(1); + let optimizer = Optimizer::with_rules(vec![ + Arc::new(ExtractLeafExpressions::new()), + Arc::new(PushDownLeafProjections::new()), + ]); + let optimized = optimizer.optimize(plan, &ctx, |_, _| {})?; + + insta::assert_snapshot!(format!("{optimized}"), @r#" + Projection: list_col, t.other_col + Filter: __datafusion_extracted_1 = Int32(1) + Projection: leaf_udf(list_col, Utf8("x")) AS __datafusion_extracted_1, list_col, t.other_col + Unnest: lists[t.list_col|depth=1] structs[] + TableScan: t + "#); + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index faeb5d59578e5..46d7b5b58b903 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -1419,3 +1419,47 @@ FROM ( statement ok DROP TABLE unused_unnest_pruning; +## Regression: pushing a leaf-extracted projection (containing get_field, +## which has MoveTowardsLeafNodes placement) through an `Unnest` used to +## trip `Assertion failed: expr.is_empty(): Unnest` inside +## `PushDownLeafProjections`. The optimizer must not try to pushdown these +## projections through an `Unnest` and should produce a valid plan. + +statement ok +CREATE TABLE struct_and_list_table +AS VALUES + (struct(1, 2), [10, 20, 30]), + (struct(3, 4), [40, 50]); + +query I +SELECT sum(get_field(s, 'c0')) +FROM (SELECT s, unnest(arr) + FROM (SELECT column1 AS s, column2 AS arr + FROM struct_and_list_table)); +---- +9 + +statement ok +DROP TABLE struct_and_list_table; + +## Regression: get_field directly references the struct produced by unnest. +## This covers the case where the leaf-extracted expression depends on the +## unnested column itself rather than a sibling input column below the Unnest. + +statement ok +CREATE TABLE list_struct_table +AS VALUES + ([struct(1, 'a'), struct(2, 'b')]), + ([struct(3, 'c')]); + +query IT +SELECT get_field(unnest(column1), 'c0'), get_field(unnest(column1), 'c1') +FROM list_struct_table; +---- +1 a +2 b +3 c + +statement ok +DROP TABLE list_struct_table; + From 532011f19a787c014a8ea9515afc8c0c16d76f81 Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Wed, 10 Jun 2026 17:12:23 +0200 Subject: [PATCH 06/15] fix a version 54 migration issue --- datafusion/functions-nested/src/array_compact.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/datafusion/functions-nested/src/array_compact.rs b/datafusion/functions-nested/src/array_compact.rs index d4d0f12830d28..11be494b5b20f 100644 --- a/datafusion/functions-nested/src/array_compact.rs +++ b/datafusion/functions-nested/src/array_compact.rs @@ -32,7 +32,6 @@ use datafusion_expr::{ Volatility, }; use datafusion_macros::user_doc; -use std::any::Any; use std::sync::Arc; make_udf_expr_and_func!( @@ -86,10 +85,6 @@ impl ScalarUDFImpl for ArrayCompact { "array_compact" } - fn as_any(&self) -> &dyn Any { - self - } - fn signature(&self) -> &Signature { &self.signature } From 9d170913c078abdb7bcda7bc9f47f665f8cf31ba Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Wed, 10 Jun 2026 17:32:52 +0200 Subject: [PATCH 07/15] fix a version 54 migration issue --- .../substrait/src/logical_plan/consumer/substrait_consumer.rs | 4 ++-- datafusion/substrait/tests/cases/roundtrip_logical_plan.rs | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 266916756e017..e3101ccc21dcc 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -403,8 +403,8 @@ pub trait SubstraitConsumer: Send + Sync + Sized { .r#type .as_ref() .map(|t| { - super::from_substrait_type_without_names(self, t).map(|dt| { - Arc::new(datafusion::arrow::datatypes::Field::new(&id, dt, true)) + from_substrait_type_without_names(self, t).map(|dt| { + Arc::new(Field::new(&id, dt, true)) }) }) .transpose()?; diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index ebe4eb5ee7a28..5351d4472c701 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -1949,6 +1949,7 @@ async fn roundtrip_placeholder_typed_utf8() -> Result<()> { assert_eq!(plan.schema(), plan2.schema()); Ok(()) +} #[tokio::test] async fn roundtrip_array_transform_higher_order_function() -> Result<()> { From 3b0030ba752ae21dd9fcef8c4aeee446c569cdf8 Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Thu, 11 Jun 2026 09:10:38 +0200 Subject: [PATCH 08/15] fix a version 54 migration issue --- .../src/logical_plan/consumer/substrait_consumer.rs | 5 ++--- .../substrait/tests/cases/roundtrip_logical_plan.rs | 11 ++++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index e3101ccc21dcc..8ff25561dbaee 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -403,9 +403,8 @@ pub trait SubstraitConsumer: Send + Sync + Sized { .r#type .as_ref() .map(|t| { - from_substrait_type_without_names(self, t).map(|dt| { - Arc::new(Field::new(&id, dt, true)) - }) + from_substrait_type_without_names(self, t) + .map(|dt| Arc::new(Field::new(&id, dt, true))) }) .transpose()?; Ok(Expr::Placeholder( diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 5351d4472c701..52da1237bac8a 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -22,8 +22,8 @@ use datafusion::execution::FunctionRegistry; use datafusion::functions_nested::map::map; use datafusion::logical_expr::{ ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, - HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, LogicalPlanBuilder, - ValueOrLambda, + HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, + LogicalPlanBuilder, ValueOrLambda, }; use datafusion::physical_plan::Accumulator; use datafusion::scalar::ScalarValue; @@ -2066,7 +2066,9 @@ async fn roundtrip_array_transform_higher_order_function() -> Result<()> { pub(crate) async fn higher_order_function_ctx() -> Result { let mut ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; - ctx.register_higher_order_function(Arc::new(ArrayTransform::new()))?; + ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + ArrayTransform::new(), + )))? let data3_fields = vec![ Field::new("p1", DataType::Int64, true), // lambda parameters should not conflict with this column @@ -2095,7 +2097,7 @@ impl ArrayTransform { } } -impl HigherOrderUDF for ArrayTransform { +impl HigherOrderUDFImpl for ArrayTransform { fn name(&self) -> &str { "array_transform2" } @@ -2158,7 +2160,6 @@ impl HigherOrderUDF for ArrayTransform { args.return_type(), )?)) } - } fn check_post_join_filters(rel: &Rel) -> Result<()> { From 2d0ec3e387d2e9dd66a42972d4bc29b2a92c9eff Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Thu, 11 Jun 2026 09:15:21 +0200 Subject: [PATCH 09/15] fix a version 54 migration issue --- datafusion/substrait/tests/cases/roundtrip_logical_plan.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 52da1237bac8a..3b466a293297a 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -2068,7 +2068,7 @@ pub(crate) async fn higher_order_function_ctx() -> Result { ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( ArrayTransform::new(), - )))? + )))?; let data3_fields = vec![ Field::new("p1", DataType::Int64, true), // lambda parameters should not conflict with this column From 4f02d66c66d20e2910bc92c8f4fe1552dc9ea1e9 Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Thu, 11 Jun 2026 09:24:59 +0200 Subject: [PATCH 10/15] fix a version 54 migration issue --- datafusion/substrait/tests/cases/roundtrip_logical_plan.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 3b466a293297a..e48bd2e10f8ed 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -2068,7 +2068,7 @@ pub(crate) async fn higher_order_function_ctx() -> Result { ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( ArrayTransform::new(), - )))?; + ))); let data3_fields = vec![ Field::new("p1", DataType::Int64, true), // lambda parameters should not conflict with this column From 2aa22745216cdad9a4bc69e24c90878f52fe9545 Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Thu, 11 Jun 2026 10:10:35 +0200 Subject: [PATCH 11/15] fix a version 54 migration issue --- datafusion/substrait/tests/cases/roundtrip_logical_plan.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index e48bd2e10f8ed..018e1aef80ea1 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -18,7 +18,6 @@ use crate::utils::test::read_json; use datafusion::arrow::array::ArrayRef; use datafusion::config::Dialect; -use datafusion::execution::FunctionRegistry; use datafusion::functions_nested::map::map; use datafusion::logical_expr::{ ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, @@ -2064,7 +2063,7 @@ async fn roundtrip_array_transform_higher_order_function() -> Result<()> { } pub(crate) async fn higher_order_function_ctx() -> Result { - let mut ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; + let ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( ArrayTransform::new(), From 61357ab37581d2459c4fb8a70845022cbb94a301 Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Thu, 11 Jun 2026 11:02:03 +0200 Subject: [PATCH 12/15] fix a version 54 migration issue --- datafusion/sqllogictest/test_files/unnest.slt | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 46d7b5b58b903..9c7e78cbe008b 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -1419,6 +1419,7 @@ FROM ( statement ok DROP TABLE unused_unnest_pruning; + ## Regression: pushing a leaf-extracted projection (containing get_field, ## which has MoveTowardsLeafNodes placement) through an `Unnest` used to ## trip `Assertion failed: expr.is_empty(): Unnest` inside From 93ba65a41644443c83b416c7cf0d412a821e6c2a Mon Sep 17 00:00:00 2001 From: Fred Thomas Date: Tue, 16 Jun 2026 15:25:34 +0200 Subject: [PATCH 13/15] add downcast_delegate --- datafusion/datasource/src/source.rs | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index af4bc09504937..766aea0e5159f 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -123,6 +123,18 @@ use datafusion_physical_plan::filter_pushdown::{ /// └─────────────────────┘ /// ``` pub trait DataSource: Any + Send + Sync + Debug { + /// Optionally delegates downcast identity to an inner [`DataSource`]. + /// + /// Wrapper types (e.g. a partitioning adapter that wraps a [`FileScanConfig`]) + /// can implement this to make `is` and `downcast_ref` transparently see + /// through to the inner source, exactly as [`ExecutionPlan::downcast_delegate`] + /// does for execution plans. + /// + /// Implementations should return the inner source, not `self`. + fn downcast_delegate(&self) -> Option<&dyn DataSource> { + None + } + /// Open the specified output partition and return its stream of /// [`RecordBatch`]es. /// @@ -292,11 +304,17 @@ impl OpenArgs { impl dyn DataSource { pub fn is(&self) -> bool { - (self as &dyn Any).is::() + match self.downcast_delegate() { + Some(delegate) => delegate.is::(), + None => (self as &dyn Any).is::(), + } } pub fn downcast_ref(&self) -> Option<&T> { - (self as &dyn Any).downcast_ref() + match self.downcast_delegate() { + Some(delegate) => delegate.downcast_ref::(), + None => (self as &dyn Any).downcast_ref(), + } } } From e146411f2b44933f67e2237c066489938325ecf6 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 24 Jun 2026 13:14:24 +0200 Subject: [PATCH 14/15] Merge pull request #149 from DataDog/codex/branch-54-proto-downcast-delegate fix(proto): honor ExecutionPlan downcast_delegate during serialization --- datafusion/proto/src/physical_plan/mod.rs | 2 +- .../tests/cases/roundtrip_physical_plan.rs | 71 ++++++++++++++++++- 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 9a5489177319d..159e485fe7b87 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -424,7 +424,7 @@ impl protobuf::PhysicalPlanNode { Self: Sized, { let plan_clone = Arc::clone(&plan); - let plan = plan.as_ref() as &dyn Any; + let plan = plan.as_ref(); if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec); diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 0cb6068af3b29..1810c69da36ff 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -88,7 +88,8 @@ use datafusion::physical_plan::windows::{ create_udwf_window_expr, }; use datafusion::physical_plan::{ - ExecutionPlan, InputOrderMode, Partitioning, PhysicalExpr, Statistics, displayable, + DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, + PhysicalExpr, PlanProperties, SendableRecordBatchStream, Statistics, displayable, }; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion::scalar::ScalarValue; @@ -228,6 +229,74 @@ fn roundtrip_empty() -> Result<()> { roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty())))) } +#[derive(Debug)] +struct DowncastDelegatingExec { + inner: Arc, +} + +impl DowncastDelegatingExec { + fn new(inner: Arc) -> Self { + Self { inner } + } +} + +impl DisplayAs for DowncastDelegatingExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.inner.fmt_as(t, f) + } +} + +impl ExecutionPlan for DowncastDelegatingExec { + fn name(&self) -> &str { + self.inner.name() + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + self.inner.children() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let inner = Arc::clone(&self.inner).with_new_children(children)?; + Ok(Arc::new(Self::new(inner))) + } + + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + Some(self.inner.as_ref()) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } +} + +#[test] +fn serialize_uses_downcast_delegate() -> Result<()> { + let inner: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + let plan: Arc = Arc::new(DowncastDelegatingExec::new(inner)); + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; + + assert!(matches!( + proto.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(_)) + )); + + Ok(()) +} + #[test] fn roundtrip_date_time_interval() -> Result<()> { let schema = Schema::new(vec![ From d33d7f1ac3fd740f7b5484b90012f054acf037fd Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Sun, 5 Jul 2026 07:15:27 +0200 Subject: [PATCH 15/15] fix: Avoid panicing when stats are not available for a file group split (#23277) ## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/23219. ## Rationale for this change The query from the issue: ```sql SELECT (((Cast(id AS BIGINT) % 1024) + 1024) % 1024) AS computed_bucket FROM profile ORDER BY computed_bucket, Cast(id AS BIGINT) limit 10; ``` panics: ``` thread 'main' panicked at .../datafusion-datasource-54.0.0/src/statistics.rs:100:48: index out of bounds: the len is 0 but the index is 0 ``` The underlying issue is that the current code panics when files are split by statistics and there are no statistics available for the column where the sort order is defined in this case `computed_bucket`. ## What changes are included in this PR? - Fix in `MinMaxStatistics` to check if there are stats available for a given column - Test ## Are these changes tested? Yes ## Are there any user-facing changes? No --- datafusion/datasource/src/statistics.rs | 39 +++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/datafusion/datasource/src/statistics.rs b/datafusion/datasource/src/statistics.rs index 6abfafe9d39d4..5063bf7785816 100644 --- a/datafusion/datasource/src/statistics.rs +++ b/datafusion/datasource/src/statistics.rs @@ -97,8 +97,16 @@ impl MinMaxStatistics { .zip(s.column_statistics[i].max_value.get_value().cloned()) .ok_or_else(|| plan_datafusion_err!("statistics not found")) } else { - let partition_value = &pv[i - s.column_statistics.len()]; - Ok((partition_value.clone(), partition_value.clone())) + if let Some(partition_value) = + pv.get(i - s.column_statistics.len()) + { + Ok((partition_value.clone(), partition_value.clone())) + } else { + Err(plan_datafusion_err!( + "statistics not found for partition, expected at most {}", + s.column_statistics.len() + )) + } } }) .collect::>>()? @@ -890,4 +898,31 @@ mod tests { Ok(()) } + + #[test] + fn min_max_statistics_missing_column_stats_returns_error() { + let schema = test_schema(); + let sort_order = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + let files = [ + file_with_stats("f1.parquet", Statistics::default()), + file_with_stats("f2.parquet", Statistics::default()), + ]; + + let err = match MinMaxStatistics::new_from_files( + &sort_order, + &schema, + None, + files.iter(), + ) { + Ok(_) => panic!("expected missing statistics error"), + Err(err) => err, + }; + + assert!( + err.to_string() + .contains("statistics not found for partition"), + "unexpected error: {err:?}" + ); + } }