Refactor: move from_substrait_extended_expr to mod.rs

Author: gabotechs

datafusion/substrait/src/logical_plan/consumer/expr/extended_expression.rs

@@ -17,7 +17,6 @@
 
 mod aggregate_function;
 mod cast;
-mod extended_expression;
 mod field_reference;
 mod function_arguments;
 mod if_then;
@@ -29,7 +28,6 @@ mod window_function;
 
 pub use aggregate_function::*;
 pub use cast::*;
-pub use extended_expression::*;
 pub use field_reference::*;
 pub use function_arguments::*;
 pub use if_then::*;
@@ -39,11 +37,18 @@ pub use singular_or_list::*;
 pub use subquery::*;
 pub use window_function::*;
 
-use crate::logical_plan::consumer::SubstraitConsumer;
-use datafusion::common::{substrait_err, DFSchema};
-use datafusion::logical_expr::Expr;
+use crate::extensions::Extensions;
+use crate::logical_plan::consumer::utils::rename_field;
+use crate::logical_plan::consumer::{
+    from_substrait_named_struct, DefaultSubstraitConsumer, SubstraitConsumer,
+};
+use datafusion::arrow::datatypes::Field;
+use datafusion::common::{not_impl_err, plan_err, substrait_err, DFSchema, DFSchemaRef};
+use datafusion::execution::SessionState;
+use datafusion::logical_expr::{Expr, ExprSchemable};
 use substrait::proto::expression::RexType;
-use substrait::proto::Expression;
+use substrait::proto::expression_reference::ExprType;
+use substrait::proto::{Expression, ExtendedExpression};
 
 /// Convert Substrait Rex to DataFusion Expr
 pub async fn from_substrait_rex(
@@ -92,6 +97,87 @@ pub async fn from_substrait_rex(
     }
 }
 
+/// Convert Substrait ExtendedExpression to ExprContainer
+///
+/// A Substrait ExtendedExpression message contains one or more expressions,
+/// with names for the outputs, and an input schema.  These pieces are all included
+/// in the ExprContainer.
+///
+/// This is a top-level message and can be used to send expressions (not plans)
+/// between systems.  This is often useful for scenarios like pushdown where filter
+/// expressions need to be sent to remote systems.
+pub async fn from_substrait_extended_expr(
+    state: &SessionState,
+    extended_expr: &ExtendedExpression,
+) -> datafusion::common::Result<ExprContainer> {
+    // Register function extension
+    let extensions = Extensions::try_from(&extended_expr.extensions)?;
+    if !extensions.type_variations.is_empty() {
+        return not_impl_err!("Type variation extensions are not supported");
+    }
+
+    let consumer = DefaultSubstraitConsumer {
+        extensions: &extensions,
+        state,
+    };
+
+    let input_schema = DFSchemaRef::new(match &extended_expr.base_schema {
+        Some(base_schema) => from_substrait_named_struct(&consumer, base_schema),
+        None => {
+            plan_err!("required property `base_schema` missing from Substrait ExtendedExpression message")
+        }
+    }?);
+
+    // Parse expressions
+    let mut exprs = Vec::with_capacity(extended_expr.referred_expr.len());
+    for (expr_idx, substrait_expr) in extended_expr.referred_expr.iter().enumerate() {
+        let scalar_expr = match &substrait_expr.expr_type {
+            Some(ExprType::Expression(scalar_expr)) => Ok(scalar_expr),
+            Some(ExprType::Measure(_)) => {
+                not_impl_err!("Measure expressions are not yet supported")
+            }
+            None => {
+                plan_err!("required property `expr_type` missing from Substrait ExpressionReference message")
+            }
+        }?;
+        let expr = consumer
+            .consume_expression(scalar_expr, &input_schema)
+            .await?;
+        let (output_type, expected_nullability) =
+            expr.data_type_and_nullable(&input_schema)?;
+        let output_field = Field::new("", output_type, expected_nullability);
+        let mut names_idx = 0;
+        let output_field = rename_field(
+            &output_field,
+            &substrait_expr.output_names,
+            expr_idx,
+            &mut names_idx,
+            /*rename_self=*/ true,
+        )?;
+        exprs.push((expr, output_field));
+    }
+
+    Ok(ExprContainer {
+        input_schema,
+        exprs,
+    })
+}
+
+/// An ExprContainer is a container for a collection of expressions with a common input schema
+///
+/// In addition, each expression is associated with a field, which defines the
+/// expression's output.  The data type and nullability of the field are calculated from the
+/// expression and the input schema.  However the names of the field (and its nested fields) are
+/// derived from the Substrait message.
+pub struct ExprContainer {
+    /// The input schema for the expressions
+    pub input_schema: DFSchemaRef,
+    /// The expressions
+    ///
+    /// Each item contains an expression and the field that defines the expected nullability and name of the expr's output
+    pub exprs: Vec<(Expr, Field)>,
+}
+
 /// Convert Substrait Expressions to DataFusion Exprs
 pub async fn from_substrait_rex_vec(
     consumer: &impl SubstraitConsumer,