Factor out Substrait consumers into separate files

Author: gabotechs

datafusion/substrait/src/logical_plan/consumer.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 super::{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/agg_func.rs

@@ -0,0 +1,120 @@
+// 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 super::grouping::from_substrait_grouping;
+use super::SubstraitConsumer;
+use super::{from_substrait_agg_func, from_substrait_sorts};
+use datafusion::common::not_impl_err;
+use datafusion::logical_expr::{Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder};
+use substrait::proto::aggregate_function::AggregationInvocation;
+use substrait::proto::AggregateRel;
+
+pub async fn from_aggregate_rel(
+    consumer: &impl SubstraitConsumer,
+    agg: &AggregateRel,
+) -> datafusion::common::Result<LogicalPlan> {
+    if let Some(input) = agg.input.as_ref() {
+        let input = LogicalPlanBuilder::from(consumer.consume_rel(input).await?);
+        let mut ref_group_exprs = vec![];
+
+        for e in &agg.grouping_expressions {
+            let x = consumer.consume_expression(e, input.schema()).await?;
+            ref_group_exprs.push(x);
+        }
+
+        let mut group_exprs = vec![];
+        let mut aggr_exprs = vec![];
+
+        match agg.groupings.len() {
+            1 => {
+                group_exprs.extend_from_slice(
+                    &from_substrait_grouping(
+                        consumer,
+                        &agg.groupings[0],
+                        &ref_group_exprs,
+                        input.schema(),
+                    )
+                    .await?,
+                );
+            }
+            _ => {
+                let mut grouping_sets = vec![];
+                for grouping in &agg.groupings {
+                    let grouping_set = from_substrait_grouping(
+                        consumer,
+                        grouping,
+                        &ref_group_exprs,
+                        input.schema(),
+                    )
+                    .await?;
+                    grouping_sets.push(grouping_set);
+                }
+                // Single-element grouping expression of type Expr::GroupingSet.
+                // Note that GroupingSet::Rollup would become GroupingSet::GroupingSets, when
+                // parsed by the producer and consumer, since Substrait does not have a type dedicated
+                // to ROLLUP. Only vector of Groupings (grouping sets) is available.
+                group_exprs
+                    .push(Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)));
+            }
+        };
+
+        for m in &agg.measures {
+            let filter = match &m.filter {
+                Some(fil) => Some(Box::new(
+                    consumer.consume_expression(fil, input.schema()).await?,
+                )),
+                None => None,
+            };
+            let agg_func = match &m.measure {
+                Some(f) => {
+                    let distinct = match f.invocation {
+                        _ if f.invocation == AggregationInvocation::Distinct as i32 => {
+                            true
+                        }
+                        _ if f.invocation == AggregationInvocation::All as i32 => false,
+                        _ => false,
+                    };
+                    let order_by = if !f.sorts.is_empty() {
+                        Some(
+                            from_substrait_sorts(consumer, &f.sorts, input.schema())
+                                .await?,
+                        )
+                    } else {
+                        None
+                    };
+
+                    from_substrait_agg_func(
+                        consumer,
+                        f,
+                        input.schema(),
+                        filter,
+                        order_by,
+                        distinct,
+                    )
+                    .await
+                }
+                None => {
+                    not_impl_err!("Aggregate without aggregate function is not supported")
+                }
+            };
+            aggr_exprs.push(agg_func?.as_ref().clone());
+        }
+        input.aggregate(group_exprs, aggr_exprs)?.build()
+    } else {
+        not_impl_err!("Aggregate without an input is not valid")
+    }
+}

datafusion/substrait/src/logical_plan/consumer/aggregate_rel.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 datafusion::common::{plan_err, substrait_err, ScalarValue};
+use datafusion::logical_expr::WindowFrameBound;
+use substrait::proto::expression::{
+    window_function::bound as SubstraitBound, window_function::bound::Kind as BoundKind,
+    window_function::Bound,
+};
+
+pub(super) fn from_substrait_bound(
+    bound: &Option<Bound>,
+    is_lower: bool,
+) -> datafusion::common::Result<WindowFrameBound> {
+    match bound {
+        Some(b) => match &b.kind {
+            Some(k) => match k {
+                BoundKind::CurrentRow(SubstraitBound::CurrentRow {}) => {
+                    Ok(WindowFrameBound::CurrentRow)
+                }
+                BoundKind::Preceding(SubstraitBound::Preceding { offset }) => {
+                    if *offset <= 0 {
+                        return plan_err!("Preceding bound must be positive");
+                    }
+                    Ok(WindowFrameBound::Preceding(ScalarValue::UInt64(Some(
+                        *offset as u64,
+                    ))))
+                }
+                BoundKind::Following(SubstraitBound::Following { offset }) => {
+                    if *offset <= 0 {
+                        return plan_err!("Following bound must be positive");
+                    }
+                    Ok(WindowFrameBound::Following(ScalarValue::UInt64(Some(
+                        *offset as u64,
+                    ))))
+                }
+                BoundKind::Unbounded(SubstraitBound::Unbounded {}) => {
+                    if is_lower {
+                        Ok(WindowFrameBound::Preceding(ScalarValue::Null))
+                    } else {
+                        Ok(WindowFrameBound::Following(ScalarValue::Null))
+                    }
+                }
+            },
+            None => substrait_err!("WindowFunction missing Substrait Bound kind"),
+        },
+        None => {
+            if is_lower {
+                Ok(WindowFrameBound::Preceding(ScalarValue::Null))
+            } else {
+                Ok(WindowFrameBound::Following(ScalarValue::Null))
+            }
+        }
+    }
+}

datafusion/substrait/src/logical_plan/consumer/bound.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 super::r#type::from_substrait_type_without_names;
+use super::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/cast.rs

@@ -0,0 +1,35 @@
+// 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 super::utils::requalify_sides_if_needed;
+use super::SubstraitConsumer;
+use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder};
+use substrait::proto::CrossRel;
+
+pub async fn from_cross_rel(
+    consumer: &impl SubstraitConsumer,
+    cross: &CrossRel,
+) -> datafusion::common::Result<LogicalPlan> {
+    let left = LogicalPlanBuilder::from(
+        consumer.consume_rel(cross.left.as_ref().unwrap()).await?,
+    );
+    let right = LogicalPlanBuilder::from(
+        consumer.consume_rel(cross.right.as_ref().unwrap()).await?,
+    );
+    let (left, right) = requalify_sides_if_needed(left, right)?;
+    left.cross_join(right.build()?)?.build()
+}