[branch-53] Add lambda substrait support (#21193)

datafusion/expr/src/logical_plan/plan.rs

@@ -4532,49 +4532,49 @@ mod tests {
         [
           {
             "Plan": {
+              "Node Type": "Projection",
               "Expressions": [
                 "employee_csv.id"
               ],
-              "Node Type": "Projection",
-              "Output": [
-                "id"
-              ],
               "Plans": [
                 {
-                  "Condition": "employee_csv.state IN (<subquery>)",
                   "Node Type": "Filter",
-                  "Output": [
-                    "id",
-                    "state"
-                  ],
+                  "Condition": "employee_csv.state IN (<subquery>)",
                   "Plans": [
                     {
                       "Node Type": "Subquery",
-                      "Output": [
-                        "state"
-                      ],
                       "Plans": [
                         {
                           "Node Type": "TableScan",
+                          "Relation Name": "employee_csv",
+                          "Plans": [],
                           "Output": [
                             "state"
-                          ],
-                          "Plans": [],
-                          "Relation Name": "employee_csv"
+                          ]
                         }
+                      ],
+                      "Output": [
+                        "state"
                       ]
                     },
                     {
                       "Node Type": "TableScan",
+                      "Relation Name": "employee_csv",
+                      "Plans": [],
                       "Output": [
                         "id",
                         "state"
-                      ],
-                      "Plans": [],
-                      "Relation Name": "employee_csv"
+                      ]
                     }
+                  ],
+                  "Output": [
+                    "id",
+                    "state"
                   ]
                 }
+              ],
+              "Output": [
+                "id"
               ]
             }
           }

datafusion/sqllogictest/test_files/explain.slt

@@ -642,11 +642,11 @@ logical_plan
 02)--{
 03)----"Plan": {
 04)------"Node Type": "Values",
-05)------"Output": [
-06)--------"column1"
-07)------],
-08)------"Plans": [],
-09)------"Values": "(Int64(1))"
+05)------"Values": "(Int64(1))",
+06)------"Plans": [],
+07)------"Output": [
+08)--------"column1"
+09)------]
 10)----}
 11)--}
 12)]

datafusion/substrait/Cargo.toml

@@ -44,7 +44,7 @@ object_store = { workspace = true }
 # We need to match the version in substrait, so we don't use the workspace version here
 pbjson-types = { version = "0.8.0" }
 prost = { workspace = true }
-substrait = { version = "=0.62.2", features = ["serde"] }
+substrait = { version = "0.63.0", features = ["serde"] }
 url = { workspace = true }
 tokio = { workspace = true, features = ["fs"] }
 

datafusion/substrait/src/extensions.rs

@@ -113,14 +113,10 @@ impl TryFrom<&Vec<SimpleExtensionDeclaration>> for Extensions {
 }
 
 impl From<Extensions> for Vec<SimpleExtensionDeclaration> {
-    // Silence deprecation warnings for `extension_uri_reference` during the uri -> urn migration
-    // See: https://github.com/substrait-io/substrait/issues/856
-    #[expect(deprecated)]
     fn from(val: Extensions) -> Vec<SimpleExtensionDeclaration> {
         let mut extensions = vec![];
         for (f_anchor, f_name) in val.functions {
             let function_extension = ExtensionFunction {
-                extension_uri_reference: u32::MAX,
                 extension_urn_reference: u32::MAX,
                 function_anchor: f_anchor,
                 name: f_name,
@@ -133,7 +129,6 @@ impl From<Extensions> for Vec<SimpleExtensionDeclaration> {
 
         for (t_anchor, t_name) in val.types {
             let type_extension = ExtensionType {
-                extension_uri_reference: u32::MAX, // https://github.com/apache/datafusion/issues/11545
                 extension_urn_reference: u32::MAX, // https://github.com/apache/datafusion/issues/11545
                 type_anchor: t_anchor,
                 name: t_name,
@@ -146,8 +141,7 @@ impl From<Extensions> for Vec<SimpleExtensionDeclaration> {
 
         for (tv_anchor, tv_name) in val.type_variations {
             let type_variation_extension = ExtensionTypeVariation {
-                extension_uri_reference: u32::MAX, // We don't register proper extension URIs yet
-                extension_urn_reference: u32::MAX, // We don't register proper extension URIs yet
+                extension_urn_reference: u32::MAX, // We don't register proper extension URNs yet
                 type_variation_anchor: tv_anchor,
                 name: tv_name,
             };

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,6 +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(
+                        LambdaParameterReference { steps_out },
+                    )) => consumer.lambda_variable(*steps_out as usize, field_idx),
                 }
             }
             _ => not_impl_err!(
@@ -82,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 },
+            )),
+        }
+    }
+}