datafusion_substrait/logical_plan/consumer/expr/
mod.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
18mod aggregate_function;
19mod cast;
20mod field_reference;
21mod function_arguments;
22mod if_then;
23mod lambda;
24mod literal;
25mod nested;
26mod scalar_function;
27mod singular_or_list;
28mod subquery;
29mod window_function;
30
31pub use aggregate_function::*;
32pub use cast::*;
33pub use field_reference::*;
34pub use function_arguments::*;
35pub use if_then::*;
36pub use lambda::*;
37pub use literal::*;
38pub use nested::*;
39pub use scalar_function::*;
40pub use singular_or_list::*;
41pub use subquery::*;
42pub use window_function::*;
43
44use crate::extensions::Extensions;
45use crate::logical_plan::consumer::{
46    DefaultSubstraitConsumer, SubstraitConsumer, from_substrait_named_struct,
47    rename_field,
48};
49use datafusion::arrow::datatypes::Field;
50use datafusion::common::{DFSchema, DFSchemaRef, not_impl_err, plan_err, substrait_err};
51use datafusion::execution::SessionState;
52use datafusion::logical_expr::{Expr, ExprSchemable};
53use substrait::proto::expression::RexType;
54use substrait::proto::expression_reference::ExprType;
55use substrait::proto::{Expression, ExtendedExpression};
56
57/// Convert Substrait Rex to DataFusion Expr
58pub async fn from_substrait_rex(
59    consumer: &impl SubstraitConsumer,
60    expression: &Expression,
61    input_schema: &DFSchema,
62) -> datafusion::common::Result<Expr> {
63    match &expression.rex_type {
64        Some(t) => match t {
65            RexType::Literal(expr) => consumer.consume_literal(expr).await,
66            RexType::Selection(expr) => {
67                consumer.consume_field_reference(expr, input_schema).await
68            }
69            RexType::ScalarFunction(expr) => {
70                consumer.consume_scalar_function(expr, input_schema).await
71            }
72            RexType::WindowFunction(expr) => {
73                consumer.consume_window_function(expr, input_schema).await
74            }
75            RexType::IfThen(expr) => consumer.consume_if_then(expr, input_schema).await,
76            RexType::SwitchExpression(expr) => {
77                consumer.consume_switch(expr, input_schema).await
78            }
79            RexType::SingularOrList(expr) => {
80                consumer.consume_singular_or_list(expr, input_schema).await
81            }
82
83            RexType::MultiOrList(expr) => {
84                consumer.consume_multi_or_list(expr, input_schema).await
85            }
86
87            RexType::Cast(expr) => {
88                consumer.consume_cast(expr.as_ref(), input_schema).await
89            }
90
91            RexType::Subquery(expr) => {
92                consumer.consume_subquery(expr.as_ref(), input_schema).await
93            }
94            RexType::Nested(expr) => consumer.consume_nested(expr, input_schema).await,
95            #[expect(deprecated)]
96            RexType::Enum(expr) => consumer.consume_enum(expr, input_schema).await,
97            RexType::DynamicParameter(expr) => {
98                consumer.consume_dynamic_parameter(expr, input_schema).await
99            }
100            RexType::Lambda(lambda) => {
101                consumer.consume_lambda(lambda.as_ref(), input_schema).await
102            }
103            RexType::LambdaInvocation(_) => {
104                not_impl_err!("Lambda invocations are not supported")
105            }
106        },
107        None => substrait_err!("Expression must set rex_type: {expression:?}"),
108    }
109}
110
111/// Convert Substrait ExtendedExpression to ExprContainer
112///
113/// A Substrait ExtendedExpression message contains one or more expressions,
114/// with names for the outputs, and an input schema.  These pieces are all included
115/// in the ExprContainer.
116///
117/// This is a top-level message and can be used to send expressions (not plans)
118/// between systems.  This is often useful for scenarios like pushdown where filter
119/// expressions need to be sent to remote systems.
120pub async fn from_substrait_extended_expr(
121    state: &SessionState,
122    extended_expr: &ExtendedExpression,
123) -> datafusion::common::Result<ExprContainer> {
124    // Register function extension
125    let extensions = Extensions::try_from(&extended_expr.extensions)?;
126    if !extensions.type_variations.is_empty() {
127        return not_impl_err!("Type variation extensions are not supported");
128    }
129
130    let consumer = DefaultSubstraitConsumer::new(&extensions, state);
131
132    let input_schema = DFSchemaRef::new(match &extended_expr.base_schema {
133        Some(base_schema) => from_substrait_named_struct(&consumer, base_schema),
134        None => {
135            plan_err!(
136                "required property `base_schema` missing from Substrait ExtendedExpression message"
137            )
138        }
139    }?);
140
141    // Parse expressions
142    let mut exprs = Vec::with_capacity(extended_expr.referred_expr.len());
143    for (expr_idx, substrait_expr) in extended_expr.referred_expr.iter().enumerate() {
144        let scalar_expr = match &substrait_expr.expr_type {
145            Some(ExprType::Expression(scalar_expr)) => Ok(scalar_expr),
146            Some(ExprType::Measure(_)) => {
147                not_impl_err!("Measure expressions are not yet supported")
148            }
149            None => {
150                plan_err!(
151                    "required property `expr_type` missing from Substrait ExpressionReference message"
152                )
153            }
154        }?;
155        let expr = consumer
156            .consume_expression(scalar_expr, &input_schema)
157            .await?;
158        let output_field = expr.to_field(&input_schema)?.1;
159        let mut names_idx = 0;
160        let output_field = rename_field(
161            &output_field,
162            &substrait_expr.output_names,
163            expr_idx,
164            &mut names_idx,
165        )?;
166        exprs.push((expr, output_field));
167    }
168
169    Ok(ExprContainer {
170        input_schema,
171        exprs,
172    })
173}
174
175/// An ExprContainer is a container for a collection of expressions with a common input schema
176///
177/// In addition, each expression is associated with a field, which defines the
178/// expression's output.  The data type and nullability of the field are calculated from the
179/// expression and the input schema.  However the names of the field (and its nested fields) are
180/// derived from the Substrait message.
181pub struct ExprContainer {
182    /// The input schema for the expressions
183    pub input_schema: DFSchemaRef,
184    /// The expressions
185    ///
186    /// Each item contains an expression and the field that defines the expected nullability and name of the expr's output
187    pub exprs: Vec<(Expr, Field)>,
188}
189
190/// Convert Substrait Expressions to DataFusion Exprs
191pub async fn from_substrait_rex_vec(
192    consumer: &impl SubstraitConsumer,
193    exprs: &Vec<Expression>,
194    input_schema: &DFSchema,
195) -> datafusion::common::Result<Vec<Expr>> {
196    let mut expressions: Vec<Expr> = vec![];
197    for expr in exprs {
198        let expression = consumer.consume_expression(expr, input_schema).await?;
199        expressions.push(expression);
200    }
201    Ok(expressions)
202}
203
204#[cfg(test)]
205mod tests {
206    use crate::extensions::Extensions;
207    use crate::logical_plan::consumer::utils::tests::test_consumer;
208    use crate::logical_plan::consumer::*;
209    use datafusion::common::DFSchema;
210    use datafusion::logical_expr::Expr;
211    use substrait::proto::Expression;
212    use substrait::proto::expression::RexType;
213    use substrait::proto::expression::window_function::BoundsType;
214
215    #[tokio::test]
216    async fn window_function_with_range_unit_and_no_order_by()
217    -> datafusion::common::Result<()> {
218        let substrait = Expression {
219            rex_type: Some(RexType::WindowFunction(
220                substrait::proto::expression::WindowFunction {
221                    function_reference: 0,
222                    bounds_type: BoundsType::Range as i32,
223                    sorts: vec![],
224                    ..Default::default()
225                },
226            )),
227        };
228
229        let mut consumer = test_consumer();
230
231        // Just registering a single function (index 0) so that the plan
232        // does not throw a "function not found" error.
233        let mut extensions = Extensions::default();
234        extensions.register_function("count");
235        consumer.extensions = &extensions;
236
237        match from_substrait_rex(&consumer, &substrait, &DFSchema::empty()).await? {
238            Expr::WindowFunction(window_function) => {
239                assert_eq!(window_function.params.order_by.len(), 1)
240            }
241            _ => panic!("expr was not a WindowFunction"),
242        };
243
244        Ok(())
245    }
246
247    #[tokio::test]
248    async fn window_function_with_count() -> datafusion::common::Result<()> {
249        let substrait = Expression {
250            rex_type: Some(RexType::WindowFunction(
251                substrait::proto::expression::WindowFunction {
252                    function_reference: 0,
253                    ..Default::default()
254                },
255            )),
256        };
257
258        let mut consumer = test_consumer();
259
260        let mut extensions = Extensions::default();
261        extensions.register_function("count");
262        consumer.extensions = &extensions;
263
264        match from_substrait_rex(&consumer, &substrait, &DFSchema::empty()).await? {
265            Expr::WindowFunction(window_function) => {
266                assert_eq!(window_function.params.args.len(), 1)
267            }
268            _ => panic!("expr was not a WindowFunction"),
269        };
270
271        Ok(())
272    }
273}
datafusion_substrait/logical_plan/consumer/expr/mod.rs

datafusion_substrait/logical_plan/consumer/expr/
mod.rs
