Factor out Substrait consumers into separate files (#15794)

* Factor out Substrait consumers into separate files

* Move relations and expressions into their own modules

* Refactor: rename rex to expr

* Refactor: move from_substrait_extended_expr to mod.rs

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: gabotechs

datafusion/substrait/src/logical_plan/consumer.rs

@@ -0,0 +1,71 @@
+// 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 crate::logical_plan::consumer::{
+    from_substrait_func_args, substrait_fun_name, SubstraitConsumer,
+};
+use datafusion::common::{not_impl_datafusion_err, plan_err, DFSchema, ScalarValue};
+use datafusion::execution::FunctionRegistry;
+use datafusion::logical_expr::{expr, Expr, SortExpr};
+use std::sync::Arc;
+use substrait::proto::AggregateFunction;
+
+/// Convert Substrait AggregateFunction to DataFusion Expr
+pub async fn from_substrait_agg_func(
+    consumer: &impl SubstraitConsumer,
+    f: &AggregateFunction,
+    input_schema: &DFSchema,
+    filter: Option<Box<Expr>>,
+    order_by: Option<Vec<SortExpr>>,
+    distinct: bool,
+) -> datafusion::common::Result<Arc<Expr>> {
+    let Some(fn_signature) = consumer
+        .get_extensions()
+        .functions
+        .get(&f.function_reference)
+    else {
+        return plan_err!(
+            "Aggregate function not registered: function anchor = {:?}",
+            f.function_reference
+        );
+    };
+
+    let fn_name = substrait_fun_name(fn_signature);
+    let udaf = consumer.get_function_registry().udaf(fn_name);
+    let udaf = udaf.map_err(|_| {
+        not_impl_datafusion_err!(
+            "Aggregate function {} is not supported: function anchor = {:?}",
+            fn_signature,
+            f.function_reference
+        )
+    })?;
+
+    let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?;
+
+    // Datafusion does not support aggregate functions with no arguments, so
+    // we inject a dummy argument that does not affect the query, but allows
+    // us to bypass this limitation.
+    let args = if udaf.name() == "count" && args.is_empty() {
+        vec![Expr::Literal(ScalarValue::Int64(Some(1)))]
+    } else {
+        args
+    };
+
+    Ok(Arc::new(Expr::AggregateFunction(
+        expr::AggregateFunction::new_udf(udaf, args, distinct, filter, order_by, None),
+    )))
+}

datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.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 crate::logical_plan::consumer::types::from_substrait_type_without_names;
+use crate::logical_plan::consumer::SubstraitConsumer;
+use datafusion::common::{substrait_err, DFSchema};
+use datafusion::logical_expr::{Cast, Expr, TryCast};
+use substrait::proto::expression as substrait_expression;
+use substrait::proto::expression::cast::FailureBehavior::ReturnNull;
+
+pub async fn from_cast(
+    consumer: &impl SubstraitConsumer,
+    cast: &substrait_expression::Cast,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Expr> {
+    match cast.r#type.as_ref() {
+        Some(output_type) => {
+            let input_expr = Box::new(
+                consumer
+                    .consume_expression(
+                        cast.input.as_ref().unwrap().as_ref(),
+                        input_schema,
+                    )
+                    .await?,
+            );
+            let data_type = from_substrait_type_without_names(consumer, output_type)?;
+            if cast.failure_behavior() == ReturnNull {
+                Ok(Expr::TryCast(TryCast::new(input_expr, data_type)))
+            } else {
+                Ok(Expr::Cast(Cast::new(input_expr, data_type)))
+            }
+        }
+        None => substrait_err!("Cast expression without output type is not allowed"),
+    }
+}

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

@@ -0,0 +1,53 @@
+// 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 crate::logical_plan::consumer::SubstraitConsumer;
+use datafusion::common::{not_impl_err, Column, DFSchema};
+use datafusion::logical_expr::Expr;
+use substrait::proto::expression::field_reference::ReferenceType::DirectReference;
+use substrait::proto::expression::reference_segment::ReferenceType::StructField;
+use substrait::proto::expression::FieldReference;
+
+pub async fn from_field_reference(
+    _consumer: &impl SubstraitConsumer,
+    field_ref: &FieldReference,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Expr> {
+    from_substrait_field_reference(field_ref, input_schema)
+}
+
+pub(crate) fn from_substrait_field_reference(
+    field_ref: &FieldReference,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Expr> {
+    match &field_ref.reference_type {
+        Some(DirectReference(direct)) => match &direct.reference_type.as_ref() {
+            Some(StructField(x)) => match &x.child.as_ref() {
+                Some(_) => not_impl_err!(
+                    "Direct reference StructField with child is not supported"
+                ),
+                None => Ok(Expr::Column(Column::from(
+                    input_schema.qualified_field(x.field as usize),
+                ))),
+            },
+            _ => not_impl_err!(
+                "Direct reference with types other than StructField is not supported"
+            ),
+        },
+        _ => not_impl_err!("unsupported field ref type"),
+    }
+}

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

@@ -0,0 +1,39 @@
+// 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 crate::logical_plan::consumer::SubstraitConsumer;
+use datafusion::common::{not_impl_err, DFSchema};
+use datafusion::logical_expr::Expr;
+use substrait::proto::function_argument::ArgType;
+use substrait::proto::FunctionArgument;
+
+/// Convert Substrait FunctionArguments to DataFusion Exprs
+pub async fn from_substrait_func_args(
+    consumer: &impl SubstraitConsumer,
+    arguments: &Vec<FunctionArgument>,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Vec<Expr>> {
+    let mut args: Vec<Expr> = vec![];
+    for arg in arguments {
+        let arg_expr = match &arg.arg_type {
+            Some(ArgType::Value(e)) => consumer.consume_expression(e, input_schema).await,
+            _ => not_impl_err!("Function argument non-Value type not supported"),
+        };
+        args.push(arg_expr?);
+    }
+    Ok(args)
+}

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

@@ -0,0 +1,69 @@
+// 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 crate::logical_plan::consumer::SubstraitConsumer;
+use datafusion::common::DFSchema;
+use datafusion::logical_expr::{Case, Expr};
+use substrait::proto::expression::IfThen;
+
+pub async fn from_if_then(
+    consumer: &impl SubstraitConsumer,
+    if_then: &IfThen,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Expr> {
+    // Parse `ifs`
+    // If the first element does not have a `then` part, then we can assume it's a base expression
+    let mut when_then_expr: Vec<(Box<Expr>, Box<Expr>)> = vec![];
+    let mut expr = None;
+    for (i, if_expr) in if_then.ifs.iter().enumerate() {
+        if i == 0 {
+            // Check if the first element is type base expression
+            if if_expr.then.is_none() {
+                expr = Some(Box::new(
+                    consumer
+                        .consume_expression(if_expr.r#if.as_ref().unwrap(), input_schema)
+                        .await?,
+                ));
+                continue;
+            }
+        }
+        when_then_expr.push((
+            Box::new(
+                consumer
+                    .consume_expression(if_expr.r#if.as_ref().unwrap(), input_schema)
+                    .await?,
+            ),
+            Box::new(
+                consumer
+                    .consume_expression(if_expr.then.as_ref().unwrap(), input_schema)
+                    .await?,
+            ),
+        ));
+    }
+    // Parse `else`
+    let else_expr = match &if_then.r#else {
+        Some(e) => Some(Box::new(
+            consumer.consume_expression(e, input_schema).await?,
+        )),
+        None => None,
+    };
+    Ok(Expr::Case(Case {
+        expr,
+        when_then_expr,
+        else_expr,
+    }))
+}