datafusion_substrait/logical_plan/consumer/expr/
lambda.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use datafusion::{
19    common::{DFSchema, substrait_err},
20    prelude::{Expr, lambda},
21};
22use substrait::proto;
23
24use crate::logical_plan::consumer::SubstraitConsumer;
25
26pub async fn from_lambda(
27    consumer: &impl SubstraitConsumer,
28    expr: &proto::expression::Lambda,
29    input_schema: &DFSchema,
30) -> datafusion::common::Result<Expr> {
31    let Some(parameters) = expr.parameters.as_ref() else {
32        return substrait_err!("Lambda expression without parameters is not allowed");
33    };
34
35    let names = consumer.push_lambda_parameters(&parameters.types, input_schema)?;
36
37    let Some(body) = expr.body.as_ref() else {
38        return substrait_err!("Lambda expression without body is not allowed");
39    };
40
41    let body = consumer.consume_expression(body, input_schema).await?;
42
43    consumer.pop_lambda_parameters();
44
45    Ok(lambda(names, body))
46}
47
48#[cfg(test)]
49mod tests {
50    use datafusion::{
51        common::{DFSchema, assert_contains},
52        prelude::SessionContext,
53    };
54    use substrait::proto::{self, Expression, r#type::Struct};
55
56    use crate::{
57        extensions::Extensions,
58        logical_plan::consumer::{DefaultSubstraitConsumer, from_lambda},
59    };
60
61    #[tokio::test]
62    async fn test_lambda_without_body() {
63        let lambda = proto::expression::Lambda {
64            parameters: Some(Struct::default()),
65            body: None,
66        };
67
68        let extensions = Extensions::default();
69        let session_state = SessionContext::new().state();
70        let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state);
71
72        let err = from_lambda(&consumer, &lambda, DFSchema::empty_ref())
73            .await
74            .unwrap_err();
75
76        assert_contains!(
77            err.to_string(),
78            "Lambda expression without body is not allowed"
79        );
80    }
81
82    #[tokio::test]
83    async fn test_lambda_without_parameters() {
84        let lambda = proto::expression::Lambda {
85            parameters: None,
86            body: Some(Box::new(Expression::default())),
87        };
88
89        let extensions = Extensions::default();
90        let session_state = SessionContext::new().state();
91        let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state);
92
93        let err = from_lambda(&consumer, &lambda, DFSchema::empty_ref())
94            .await
95            .unwrap_err();
96
97        assert_contains!(
98            err.to_string(),
99            "Lambda expression without parameters is not allowed"
100        );
101    }
102}
datafusion_substrait/logical_plan/consumer/expr/lambda.rs

datafusion_substrait/logical_plan/consumer/expr/
lambda.rs
