Implement TPCH substrait integration teset, support tpch_1 (#10842)

* support tpch_1 consumer_producer_test

* refactor and optimize code

Author: Lordworms

datafusion/substrait/Cargo.toml

@@ -41,6 +41,7 @@ object_store = { workspace = true }
 pbjson-types = "0.6"
 prost = "0.12"
 substrait = { version = "0.34.0", features = ["serde"] }
+url = { workspace = true }
 
 [dev-dependencies]
 datafusion-functions-aggregate = { workspace = true }

datafusion/substrait/src/logical_plan/consumer.rs

@@ -22,6 +22,9 @@ use datafusion::arrow::datatypes::{
 use datafusion::common::{
     not_impl_err, substrait_datafusion_err, substrait_err, DFSchema, DFSchemaRef,
 };
+use substrait::proto::expression::literal::IntervalDayToSecond;
+use substrait::proto::read_rel::local_files::file_or_files::PathType::UriFile;
+use url::Url;
 
 use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
 use datafusion::execution::FunctionRegistry;
@@ -45,7 +48,7 @@ use datafusion::{
 use substrait::proto::exchange_rel::ExchangeKind;
 use substrait::proto::expression::literal::user_defined::Val;
 use substrait::proto::expression::subquery::SubqueryType;
-use substrait::proto::expression::{FieldReference, Literal, ScalarFunction};
+use substrait::proto::expression::{self, FieldReference, Literal, ScalarFunction};
 use substrait::proto::{
     aggregate_function::AggregationInvocation,
     expression::{
@@ -129,14 +132,7 @@ fn scalar_function_type_from_str(
     name: &str,
 ) -> Result<ScalarFunctionType> {
     let s = ctx.state();
-    let name = match name.rsplit_once(':') {
-        // Since 0.32.0, Substrait requires the function names to be in a compound format
-        // https://substrait.io/extensions/#function-signature-compound-names
-        // for example, `add:i8_i8`.
-        // On the consumer side, we don't really care about the signature though, just the name.
-        Some((name, _)) => name,
-        None => name,
-    };
+    let name = substrait_fun_name(name);
 
     if let Some(func) = s.scalar_functions().get(name) {
         return Ok(ScalarFunctionType::Udf(func.to_owned()));
@@ -153,6 +149,18 @@ fn scalar_function_type_from_str(
     not_impl_err!("Unsupported function name: {name:?}")
 }
 
+pub fn substrait_fun_name(name: &str) -> &str {
+    let name = match name.rsplit_once(':') {
+        // Since 0.32.0, Substrait requires the function names to be in a compound format
+        // https://substrait.io/extensions/#function-signature-compound-names
+        // for example, `add:i8_i8`.
+        // On the consumer side, we don't really care about the signature though, just the name.
+        Some((name, _)) => name,
+        None => name,
+    };
+    name
+}
+
 fn split_eq_and_noneq_join_predicate_with_nulls_equality(
     filter: &Expr,
 ) -> (Vec<(Column, Column)>, bool, Option<Expr>) {
@@ -239,6 +247,43 @@ pub async fn from_substrait_plan(
     }
 }
 
+/// parse projection
+pub fn extract_projection(
+    t: LogicalPlan,
+    projection: &::core::option::Option<expression::MaskExpression>,
+) -> Result<LogicalPlan> {
+    match projection {
+        Some(MaskExpression { select, .. }) => match &select.as_ref() {
+            Some(projection) => {
+                let column_indices: Vec<usize> = projection
+                    .struct_items
+                    .iter()
+                    .map(|item| item.field as usize)
+                    .collect();
+                match t {
+                    LogicalPlan::TableScan(mut scan) => {
+                        let fields = column_indices
+                            .iter()
+                            .map(|i| scan.projected_schema.qualified_field(*i))
+                            .map(|(qualifier, field)| {
+                                (qualifier.cloned(), Arc::new(field.clone()))
+                            })
+                            .collect();
+                        scan.projection = Some(column_indices);
+                        scan.projected_schema = DFSchemaRef::new(
+                            DFSchema::new_with_metadata(fields, HashMap::new())?,
+                        );
+                        Ok(LogicalPlan::TableScan(scan))
+                    }
+                    _ => plan_err!("unexpected plan for table"),
+                }
+            }
+            _ => Ok(t),
+        },
+        _ => Ok(t),
+    }
+}
+
 /// Convert Substrait Rel to DataFusion DataFrame
 #[async_recursion]
 pub async fn from_substrait_rel(
@@ -408,7 +453,6 @@ pub async fn from_substrait_rel(
                     };
                     aggr_expr.push(agg_func?.as_ref().clone());
                 }
-
                 input.aggregate(group_expr, aggr_expr)?.build()
             } else {
                 not_impl_err!("Aggregate without an input is not valid")
@@ -489,41 +533,7 @@ pub async fn from_substrait_rel(
                 };
                 let t = ctx.table(table_reference).await?;
                 let t = t.into_optimized_plan()?;
-                match &read.projection {
-                    Some(MaskExpression { select, .. }) => match &select.as_ref() {
-                        Some(projection) => {
-                            let column_indices: Vec<usize> = projection
-                                .struct_items
-                                .iter()
-                                .map(|item| item.field as usize)
-                                .collect();
-                            match &t {
-                                LogicalPlan::TableScan(scan) => {
-                                    let fields = column_indices
-                                        .iter()
-                                        .map(|i| {
-                                            scan.projected_schema.qualified_field(*i)
-                                        })
-                                        .map(|(qualifier, field)| {
-                                            (qualifier.cloned(), Arc::new(field.clone()))
-                                        })
-                                        .collect();
-                                    let mut scan = scan.clone();
-                                    scan.projection = Some(column_indices);
-                                    scan.projected_schema =
-                                        DFSchemaRef::new(DFSchema::new_with_metadata(
-                                            fields,
-                                            HashMap::new(),
-                                        )?);
-                                    Ok(LogicalPlan::TableScan(scan))
-                                }
-                                _ => plan_err!("unexpected plan for table"),
-                            }
-                        }
-                        _ => Ok(t),
-                    },
-                    _ => Ok(t),
-                }
+                extract_projection(t, &read.projection)
             }
             Some(ReadType::VirtualTable(vt)) => {
                 let base_schema = read.base_schema.as_ref().ok_or_else(|| {
@@ -569,7 +579,42 @@ pub async fn from_substrait_rel(
 
                 Ok(LogicalPlan::Values(Values { schema, values }))
             }
-            _ => not_impl_err!("Only NamedTable and VirtualTable reads are supported"),
+            Some(ReadType::LocalFiles(lf)) => {
+                fn extract_filename(name: &str) -> Option<String> {
+                    let corrected_url =
+                        if name.starts_with("file://") && !name.starts_with("file:///") {
+                            name.replacen("file://", "file:///", 1)
+                        } else {
+                            name.to_string()
+                        };
+
+                    Url::parse(&corrected_url).ok().and_then(|url| {
+                        let path = url.path();
+                        std::path::Path::new(path)
+                            .file_name()
+                            .map(|filename| filename.to_string_lossy().to_string())
+                    })
+                }
+
+                // we could use the file name to check the original table provider
+                // TODO: currently does not support multiple local files
+                let filename: Option<String> =
+                    lf.items.first().and_then(|x| match x.path_type.as_ref() {
+                        Some(UriFile(name)) => extract_filename(name),
+                        _ => None,
+                    });
+
+                if lf.items.len() > 1 || filename.is_none() {
+                    return not_impl_err!("Only single file reads are supported");
+                }
+                let name = filename.unwrap();
+                // directly use unwrap here since we could determine it is a valid one
+                let table_reference = TableReference::Bare { table: name.into() };
+                let t = ctx.table(table_reference).await?;
+                let t = t.into_optimized_plan()?;
+                extract_projection(t, &read.projection)
+            }
+            _ => not_impl_err!("Unsupported ReadType: {:?}", &read.as_ref().read_type),
         },
         Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
             Ok(set_op) => match set_op {
@@ -810,14 +855,22 @@ pub async fn from_substrait_agg_func(
             f.function_reference
         );
     };
-
+    // function_name.split(':').next().unwrap_or(function_name);
+    let function_name = substrait_fun_name((**function_name).as_str());
     // try udaf first, then built-in aggr fn.
     if let Ok(fun) = ctx.udaf(function_name) {
         Ok(Arc::new(Expr::AggregateFunction(
             expr::AggregateFunction::new_udf(fun, args, distinct, filter, order_by, None),
         )))
     } else if let Ok(fun) = aggregate_function::AggregateFunction::from_str(function_name)
     {
+        match &fun {
+            // deal with situation that count(*) got no arguments
+            aggregate_function::AggregateFunction::Count if args.is_empty() => {
+                args.push(Expr::Literal(ScalarValue::Int64(Some(1))));
+            }
+            _ => {}
+        }
         Ok(Arc::new(Expr::AggregateFunction(
             expr::AggregateFunction::new(fun, args, distinct, filter, order_by, None),
         )))
@@ -1261,6 +1314,8 @@ fn from_substrait_type(
             r#type::Kind::Struct(s) => Ok(DataType::Struct(from_substrait_struct_type(
                 s, dfs_names, name_idx,
             )?)),
+            r#type::Kind::Varchar(_) => Ok(DataType::Utf8),
+            r#type::Kind::FixedChar(_) => Ok(DataType::Utf8),
             _ => not_impl_err!("Unsupported Substrait type: {s_kind:?}"),
         },
         _ => not_impl_err!("`None` Substrait kind is not supported"),
@@ -1549,6 +1604,13 @@ fn from_substrait_literal(
         Some(LiteralType::Null(ntype)) => {
             from_substrait_null(ntype, dfs_names, name_idx)?
         }
+        Some(LiteralType::IntervalDayToSecond(IntervalDayToSecond {
+            days,
+            seconds,
+            microseconds,
+        })) => {
+            ScalarValue::new_interval_dt(*days, (seconds * 1000) + (microseconds / 1000))
+        }
         Some(LiteralType::UserDefined(user_defined)) => {
             match user_defined.type_reference {
                 INTERVAL_YEAR_MONTH_TYPE_REF => {

datafusion/substrait/tests/cases/consumer_integration.rs

@@ -0,0 +1,68 @@
+// 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.
+
+//! TPCH `substrait_consumer` tests
+//!
+//! This module tests that substrait plans as json encoded protobuf can be
+//! correctly read as DataFusion plans.
+//!
+//! The input data comes from  <https://github.com/substrait-io/consumer-testing/tree/main/substrait_consumer/tests/integration/queries/tpch_substrait_plans>
+
+#[cfg(test)]
+mod tests {
+    use datafusion::common::Result;
+    use datafusion::execution::options::CsvReadOptions;
+    use datafusion::prelude::SessionContext;
+    use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+    use std::fs::File;
+    use std::io::BufReader;
+    use substrait::proto::Plan;
+
+    #[tokio::test]
+    async fn tpch_test_1() -> Result<()> {
+        let ctx = create_context().await?;
+        let path = "tests/testdata/tpch_substrait_plans/query_1.json";
+        let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
+            File::open(path).expect("file not found"),
+        ))
+        .expect("failed to parse json");
+
+        let plan = from_substrait_plan(&ctx, &proto).await?;
+
+        assert!(
+            format!("{:?}", plan).eq_ignore_ascii_case(
+                "Sort: FILENAME_PLACEHOLDER_0.l_returnflag ASC NULLS LAST, FILENAME_PLACEHOLDER_0.l_linestatus ASC NULLS LAST\n  \
+                 Aggregate: groupBy=[[FILENAME_PLACEHOLDER_0.l_returnflag, FILENAME_PLACEHOLDER_0.l_linestatus]], aggr=[[SUM(FILENAME_PLACEHOLDER_0.l_quantity), SUM(FILENAME_PLACEHOLDER_0.l_extendedprice), SUM(FILENAME_PLACEHOLDER_0.l_extendedprice * Int32(1) - FILENAME_PLACEHOLDER_0.l_discount), SUM(FILENAME_PLACEHOLDER_0.l_extendedprice * Int32(1) - FILENAME_PLACEHOLDER_0.l_discount * Int32(1) + FILENAME_PLACEHOLDER_0.l_tax), AVG(FILENAME_PLACEHOLDER_0.l_quantity), AVG(FILENAME_PLACEHOLDER_0.l_extendedprice), AVG(FILENAME_PLACEHOLDER_0.l_discount), COUNT(Int64(1))]]\n    \
+                 Projection: FILENAME_PLACEHOLDER_0.l_returnflag, FILENAME_PLACEHOLDER_0.l_linestatus, FILENAME_PLACEHOLDER_0.l_quantity, FILENAME_PLACEHOLDER_0.l_extendedprice, FILENAME_PLACEHOLDER_0.l_extendedprice * (CAST(Int32(1) AS Decimal128(19, 0)) - FILENAME_PLACEHOLDER_0.l_discount), FILENAME_PLACEHOLDER_0.l_extendedprice * (CAST(Int32(1) AS Decimal128(19, 0)) - FILENAME_PLACEHOLDER_0.l_discount) * (CAST(Int32(1) AS Decimal128(19, 0)) + FILENAME_PLACEHOLDER_0.l_tax), FILENAME_PLACEHOLDER_0.l_discount\n      \
+                 Filter: FILENAME_PLACEHOLDER_0.l_shipdate <= Date32(\"1998-12-01\") - IntervalDayTime(\"IntervalDayTime { days: 120, milliseconds: 0 }\")\n        \
+                 TableScan: FILENAME_PLACEHOLDER_0 projection=[l_orderkey, l_partkey, l_suppkey, l_linenumber, l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate, l_commitdate, l_receiptdate, l_shipinstruct, l_shipmode, l_comment]"
+            )
+        );
+        Ok(())
+    }
+
+    async fn create_context() -> datafusion::common::Result<SessionContext> {
+        let ctx = SessionContext::new();
+        ctx.register_csv(
+            "FILENAME_PLACEHOLDER_0",
+            "tests/testdata/tpch/lineitem.csv",
+            CsvReadOptions::default(),
+        )
+        .await?;
+        Ok(ctx)
+    }
+}

datafusion/substrait/tests/cases/mod.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+mod consumer_integration;
 mod logical_plans;
 mod roundtrip_logical_plan;
 mod roundtrip_physical_plan;

datafusion/substrait/tests/testdata/tpch/lineitem.csv

@@ -0,0 +1,2 @@
+l_orderkey,l_partkey,l_suppkey,l_linenumber,l_quantity,l_extendedprice,l_discount,l_tax,l_returnflag,l_linestatus,l_shipdate,l_commitdate,l_receiptdate,l_shipinstruct,l_shipmode,l_comment
+1,1,1,1,17,21168.23,0.04,0.02,'N','O','1996-03-13','1996-02-12','1996-03-22','DELIVER IN PERSON','TRUCK','egular courts above the'

datafusion/substrait/tests/testdata/tpch_substrait_plans/README.md

@@ -0,0 +1,22 @@
+<!---
+  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.
+-->
+
+# Apache DataFusion Substrait consumer integration test
+
+these test json files come from [consumer-testing](https://github.com/substrait-io/consumer-testing/tree/main/substrait_consumer/tests/integration/queries/tpch_substrait_plans)