Refactor substrait producer into multiple files (#16089)

Author: gabotechs

datafusion/substrait/src/logical_plan/producer.rs

@@ -0,0 +1,99 @@
+// 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::producer::SubstraitProducer;
+use datafusion::common::DFSchemaRef;
+use datafusion::logical_expr::expr;
+use datafusion::logical_expr::expr::AggregateFunctionParams;
+use substrait::proto::aggregate_function::AggregationInvocation;
+use substrait::proto::aggregate_rel::Measure;
+use substrait::proto::function_argument::ArgType;
+use substrait::proto::sort_field::{SortDirection, SortKind};
+use substrait::proto::{
+    AggregateFunction, AggregationPhase, FunctionArgument, SortField,
+};
+
+pub fn from_aggregate_function(
+    producer: &mut impl SubstraitProducer,
+    agg_fn: &expr::AggregateFunction,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<Measure> {
+    let expr::AggregateFunction {
+        func,
+        params:
+            AggregateFunctionParams {
+                args,
+                distinct,
+                filter,
+                order_by,
+                null_treatment: _null_treatment,
+            },
+    } = agg_fn;
+    let sorts = if let Some(order_by) = order_by {
+        order_by
+            .iter()
+            .map(|expr| to_substrait_sort_field(producer, expr, schema))
+            .collect::<datafusion::common::Result<Vec<_>>>()?
+    } else {
+        vec![]
+    };
+    let mut arguments: Vec<FunctionArgument> = vec![];
+    for arg in args {
+        arguments.push(FunctionArgument {
+            arg_type: Some(ArgType::Value(producer.handle_expr(arg, schema)?)),
+        });
+    }
+    let function_anchor = producer.register_function(func.name().to_string());
+    #[allow(deprecated)]
+    Ok(Measure {
+        measure: Some(AggregateFunction {
+            function_reference: function_anchor,
+            arguments,
+            sorts,
+            output_type: None,
+            invocation: match distinct {
+                true => AggregationInvocation::Distinct as i32,
+                false => AggregationInvocation::All as i32,
+            },
+            phase: AggregationPhase::Unspecified as i32,
+            args: vec![],
+            options: vec![],
+        }),
+        filter: match filter {
+            Some(f) => Some(producer.handle_expr(f, schema)?),
+            None => None,
+        },
+    })
+}
+
+/// Converts sort expression to corresponding substrait `SortField`
+fn to_substrait_sort_field(
+    producer: &mut impl SubstraitProducer,
+    sort: &expr::Sort,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<SortField> {
+    let sort_kind = match (sort.asc, sort.nulls_first) {
+        (true, true) => SortDirection::AscNullsFirst,
+        (true, false) => SortDirection::AscNullsLast,
+        (false, true) => SortDirection::DescNullsFirst,
+        (false, false) => SortDirection::DescNullsLast,
+    };
+    Ok(SortField {
+        expr: Some(producer.handle_expr(&sort.expr, schema)?),
+        sort_kind: Some(SortKind::Direction(sort_kind.into())),
+    })
+}

datafusion/substrait/src/logical_plan/producer/expr/aggregate_function.rs

@@ -0,0 +1,154 @@
+// 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::producer::{to_substrait_type, SubstraitProducer};
+use crate::variation_const::DEFAULT_TYPE_VARIATION_REF;
+use datafusion::common::{DFSchemaRef, ScalarValue};
+use datafusion::logical_expr::{Cast, Expr, TryCast};
+use substrait::proto::expression::cast::FailureBehavior;
+use substrait::proto::expression::literal::LiteralType;
+use substrait::proto::expression::{Literal, RexType};
+use substrait::proto::Expression;
+
+pub fn from_cast(
+    producer: &mut impl SubstraitProducer,
+    cast: &Cast,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<Expression> {
+    let Cast { expr, data_type } = cast;
+    // since substrait Null must be typed, so if we see a cast(null, dt), we make it a typed null
+    if let Expr::Literal(lit) = expr.as_ref() {
+        // only the untyped(a null scalar value) null literal need this special handling
+        // since all other kind of nulls are already typed and can be handled by substrait
+        // e.g. null::<Int32Type> or null::<Utf8Type>
+        if matches!(lit, ScalarValue::Null) {
+            let lit = Literal {
+                nullable: true,
+                type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
+                literal_type: Some(LiteralType::Null(to_substrait_type(
+                    data_type, true,
+                )?)),
+            };
+            return Ok(Expression {
+                rex_type: Some(RexType::Literal(lit)),
+            });
+        }
+    }
+    Ok(Expression {
+        rex_type: Some(RexType::Cast(Box::new(
+            substrait::proto::expression::Cast {
+                r#type: Some(to_substrait_type(data_type, true)?),
+                input: Some(Box::new(producer.handle_expr(expr, schema)?)),
+                failure_behavior: FailureBehavior::ThrowException.into(),
+            },
+        ))),
+    })
+}
+
+pub fn from_try_cast(
+    producer: &mut impl SubstraitProducer,
+    cast: &TryCast,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<Expression> {
+    let TryCast { expr, data_type } = cast;
+    Ok(Expression {
+        rex_type: Some(RexType::Cast(Box::new(
+            substrait::proto::expression::Cast {
+                r#type: Some(to_substrait_type(data_type, true)?),
+                input: Some(Box::new(producer.handle_expr(expr, schema)?)),
+                failure_behavior: FailureBehavior::ReturnNull.into(),
+            },
+        ))),
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::logical_plan::producer::to_substrait_extended_expr;
+    use datafusion::arrow::datatypes::{DataType, Field};
+    use datafusion::common::DFSchema;
+    use datafusion::execution::SessionStateBuilder;
+    use datafusion::logical_expr::ExprSchemable;
+    use substrait::proto::expression_reference::ExprType;
+
+    #[tokio::test]
+    async fn fold_cast_null() {
+        let state = SessionStateBuilder::default().build();
+        let empty_schema = DFSchemaRef::new(DFSchema::empty());
+        let field = Field::new("out", DataType::Int32, false);
+
+        let expr = Expr::Literal(ScalarValue::Null)
+            .cast_to(&DataType::Int32, &empty_schema)
+            .unwrap();
+
+        let typed_null =
+            to_substrait_extended_expr(&[(&expr, &field)], &empty_schema, &state)
+                .unwrap();
+
+        if let ExprType::Expression(expr) =
+            typed_null.referred_expr[0].expr_type.as_ref().unwrap()
+        {
+            let lit = Literal {
+                nullable: true,
+                type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
+                literal_type: Some(LiteralType::Null(
+                    to_substrait_type(&DataType::Int32, true).unwrap(),
+                )),
+            };
+            let expected = Expression {
+                rex_type: Some(RexType::Literal(lit)),
+            };
+            assert_eq!(*expr, expected);
+        } else {
+            panic!("Expected expression type");
+        }
+
+        // a typed null should not be folded
+        let expr = Expr::Literal(ScalarValue::Int64(None))
+            .cast_to(&DataType::Int32, &empty_schema)
+            .unwrap();
+
+        let typed_null =
+            to_substrait_extended_expr(&[(&expr, &field)], &empty_schema, &state)
+                .unwrap();
+
+        if let ExprType::Expression(expr) =
+            typed_null.referred_expr[0].expr_type.as_ref().unwrap()
+        {
+            let cast_expr = substrait::proto::expression::Cast {
+                r#type: Some(to_substrait_type(&DataType::Int32, true).unwrap()),
+                input: Some(Box::new(Expression {
+                    rex_type: Some(RexType::Literal(Literal {
+                        nullable: true,
+                        type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
+                        literal_type: Some(LiteralType::Null(
+                            to_substrait_type(&DataType::Int64, true).unwrap(),
+                        )),
+                    })),
+                })),
+                failure_behavior: FailureBehavior::ThrowException as i32,
+            };
+            let expected = Expression {
+                rex_type: Some(RexType::Cast(Box::new(cast_expr))),
+            };
+            assert_eq!(*expr, expected);
+        } else {
+            panic!("Expected expression type");
+        }
+    }
+}

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

@@ -0,0 +1,103 @@
+// 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::{substrait_err, Column, DFSchemaRef};
+use datafusion::logical_expr::Expr;
+use substrait::proto::expression::field_reference::{
+    ReferenceType, RootReference, RootType,
+};
+use substrait::proto::expression::{
+    reference_segment, FieldReference, ReferenceSegment, RexType,
+};
+use substrait::proto::Expression;
+
+pub fn from_column(
+    col: &Column,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<Expression> {
+    let index = schema.index_of_column(col)?;
+    substrait_field_ref(index)
+}
+
+pub(crate) fn substrait_field_ref(
+    index: usize,
+) -> datafusion::common::Result<Expression> {
+    Ok(Expression {
+        rex_type: Some(RexType::Selection(Box::new(FieldReference {
+            reference_type: Some(ReferenceType::DirectReference(ReferenceSegment {
+                reference_type: Some(reference_segment::ReferenceType::StructField(
+                    Box::new(reference_segment::StructField {
+                        field: index as i32,
+                        child: None,
+                    }),
+                )),
+            })),
+            root_type: Some(RootType::RootReference(RootReference {})),
+        }))),
+    })
+}
+
+/// Try to convert an [Expr] to a [FieldReference].
+/// Returns `Err` if the [Expr] is not a [Expr::Column].
+pub(crate) fn try_to_substrait_field_reference(
+    expr: &Expr,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<FieldReference> {
+    match expr {
+        Expr::Column(col) => {
+            let index = schema.index_of_column(col)?;
+            Ok(FieldReference {
+                reference_type: Some(ReferenceType::DirectReference(ReferenceSegment {
+                    reference_type: Some(reference_segment::ReferenceType::StructField(
+                        Box::new(reference_segment::StructField {
+                            field: index as i32,
+                            child: None,
+                        }),
+                    )),
+                })),
+                root_type: Some(RootType::RootReference(RootReference {})),
+            })
+        }
+        _ => substrait_err!("Expect a `Column` expr, but found {expr:?}"),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::common::Result;
+
+    #[test]
+    fn to_field_reference() -> Result<()> {
+        let expression = substrait_field_ref(2)?;
+
+        match &expression.rex_type {
+            Some(RexType::Selection(field_ref)) => {
+                assert_eq!(
+                    field_ref
+                        .root_type
+                        .clone()
+                        .expect("root type should be set"),
+                    RootType::RootReference(RootReference {})
+                );
+            }
+
+            _ => panic!("Should not be anything other than field reference"),
+        }
+        Ok(())
+    }
+}