validate and adjust Substrait NamedTable schemas (#12223)

datafusion/substrait/src/logical_plan/consumer.rs

@@ -657,16 +657,34 @@ pub async fn from_substrait_rel(
                         table: nt.names[2].clone().into(),
                     },
                 };
-                let t = ctx.table(table_reference).await?;
+                let t = ctx.table(table_reference.clone()).await?;
                 let t = t.into_optimized_plan()?;
+                let datafusion_schema = t.schema();
+
+                let named_struct = read.base_schema.as_ref().ok_or_else(|| {
+                    substrait_datafusion_err!("No base schema provided for Named Scan")
+                })?;
+                let substrait_schema = from_substrait_named_struct(
+                    named_struct,
+                    extensions,
+                    Some(table_reference),
+                )?;
+
+                if !validate_substrait_schema(datafusion_schema, &substrait_schema) {
+                    return Err(substrait_datafusion_err!(
+                        "Schema mismatch in ReadRel: substrait: {:?}, DataFusion: {:?}",
+                        substrait_schema,
+                        datafusion_schema
+                    ));
+                };
                 extract_projection(t, &read.projection)
             }
             Some(ReadType::VirtualTable(vt)) => {
                 let base_schema = read.base_schema.as_ref().ok_or_else(|| {
                     substrait_datafusion_err!("No base schema provided for Virtual Table")
                 })?;
 
-                let schema = from_substrait_named_struct(base_schema, extensions)?;
+                let schema = from_substrait_named_struct(base_schema, extensions, None)?;
 
                 if vt.values.is_empty() {
                     return Ok(LogicalPlan::EmptyRelation(EmptyRelation {
@@ -850,6 +868,31 @@ pub async fn from_substrait_rel(
     }
 }
 
+/// Validates that the given Substrait schema matches the given DataFusion schema
+/// Returns true if the two schemas have the same qualified named fields with the same data types.
+/// Returns false otherwise.
+/// Ignores case when comparing field names
+///
+/// This code is equivalent to [DFSchema::equivalent_names_and_types] except that the field name
+/// checking is case-insensitive
+fn validate_substrait_schema(
+    datafusion_schema: &DFSchema,
+    substrait_schema: &DFSchema,
+) -> bool {
+    if datafusion_schema.fields().len() != substrait_schema.fields().len() {
+        return false;
+    }
+    let datafusion_fields = datafusion_schema.iter();
+    let substrait_fields = substrait_schema.iter();
+    datafusion_fields
+        .zip(substrait_fields)
+        .all(|((q1, f1), (q2, f2))| {
+            q1 == q2
+                && f1.name().to_lowercase() == f2.name().to_lowercase()
+                && f1.data_type().equals_datatype(f2.data_type())
+        })
+}
+
 /// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise
 /// conflict with the columns from the other.  
 /// Substrait doesn't currently allow specifying aliases, neither for columns nor for tables. For
@@ -1586,9 +1629,12 @@ fn next_struct_field_name(
     }
 }
 
-fn from_substrait_named_struct(
+/// Convert Substrait NamedStruct to DataFusion DFSchemaRef
+pub fn from_substrait_named_struct(
     base_schema: &NamedStruct,
     extensions: &Extensions,
+    // optional qualifier to apply to every field in the schema
+    field_qualifier: Option<TableReference>,
 ) -> Result<DFSchemaRef> {
     let mut name_idx = 0;
     let fields = from_substrait_struct_type(
@@ -1601,12 +1647,17 @@ fn from_substrait_named_struct(
     );
     if name_idx != base_schema.names.len() {
         return substrait_err!(
-                                "Names list must match exactly to nested schema, but found {} uses for {} names",
-                                name_idx,
-                                base_schema.names.len()
-                            );
+            "Names list must match exactly to nested schema, but found {} uses for {} names",
+            name_idx,
+            base_schema.names.len()
+        );
+    }
+    let mut df_schema = DFSchema::try_from(Schema::new(fields?))?;
+    match field_qualifier {
+        None => (),
+        Some(fq) => df_schema = df_schema.replace_qualifier(fq),
     }
-    Ok(DFSchemaRef::new(DFSchema::try_from(Schema::new(fields?))?))
+    Ok(DFSchemaRef::new(df_schema))
 }
 
 fn from_substrait_bound(

datafusion/substrait/src/logical_plan/producer.rs

@@ -43,8 +43,8 @@ use crate::variation_const::{
 use datafusion::arrow::array::{Array, GenericListArray, OffsetSizeTrait};
 use datafusion::common::{
     exec_err, internal_err, not_impl_err, plan_err, substrait_datafusion_err,
+    substrait_err, DFSchemaRef, ToDFSchema,
 };
-use datafusion::common::{substrait_err, DFSchemaRef};
 #[allow(unused_imports)]
 use datafusion::logical_expr::expr::{
     Alias, BinaryExpr, Case, Cast, GroupingSet, InList, InSubquery, Sort, WindowFunction,
@@ -140,19 +140,13 @@ pub fn to_substrait_rel(
                 maintain_singular_struct: false,
             });
 
+            let table_schema = scan.source.schema().to_dfschema_ref()?;
+            let base_schema = to_substrait_named_struct(&table_schema, extensions)?;
+
             Ok(Box::new(Rel {
                 rel_type: Some(RelType::Read(Box::new(ReadRel {
                     common: None,
-                    base_schema: Some(NamedStruct {
-                        names: scan
-                            .source
-                            .schema()
-                            .fields()
-                            .iter()
-                            .map(|f| f.name().to_owned())
-                            .collect(),
-                        r#struct: None,
-                    }),
+                    base_schema: Some(base_schema),
                     filter: None,
                     best_effort_filter: None,
                     projection,

datafusion/substrait/tests/cases/function_test.rs

@@ -19,40 +19,33 @@
 
 #[cfg(test)]
 mod tests {
+    use crate::utils::test::TestSchemaCollector;
     use datafusion::common::Result;
-    use datafusion::prelude::{CsvReadOptions, 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 contains_function_test() -> Result<()> {
-        let ctx = create_context().await?;
-
         let path = "tests/testdata/contains_plan.substrait.json";
-        let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
+
+        let proto_plan = 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?;
+        let ctx = TestSchemaCollector::generate_context_from_plan(&proto_plan);
+        let plan = from_substrait_plan(&ctx, &proto_plan).await?;
 
         let plan_str = format!("{}", plan);
 
         assert_eq!(
             plan_str,
-            "Projection: nation.b AS n_name\
-            \n  Filter: contains(nation.b, Utf8(\"IA\"))\
-            \n    TableScan: nation projection=[a, b, c, d, e, f]"
+            "Projection: nation.n_name\
+            \n  Filter: contains(nation.n_name, Utf8(\"IA\"))\
+            \n    TableScan: nation projection=[n_nationkey, n_name, n_regionkey, n_comment]"
         );
         Ok(())
     }
-
-    async fn create_context() -> datafusion::common::Result<SessionContext> {
-        let ctx = SessionContext::new();
-        ctx.register_csv("nation", "tests/testdata/data.csv", CsvReadOptions::new())
-            .await?;
-        Ok(ctx)
-    }
 }

datafusion/substrait/tests/cases/logical_plans.rs

@@ -19,13 +19,10 @@
 
 #[cfg(test)]
 mod tests {
+    use crate::utils::test::{read_json, TestSchemaCollector};
     use datafusion::common::Result;
     use datafusion::dataframe::DataFrame;
-    use datafusion::prelude::{CsvReadOptions, 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 scalar_function_compound_signature() -> Result<()> {
@@ -35,18 +32,17 @@ mod tests {
         // we don't yet produce such plans.
         // Once we start producing plans with compound signatures, this test can be replaced by the roundtrip tests.
 
-        let ctx = create_context().await?;
-
         // File generated with substrait-java's Isthmus:
-        // ./isthmus-cli/build/graal/isthmus "select not d from data" -c "create table data (d boolean)"
-        let proto = read_json("tests/testdata/test_plans/select_not_bool.substrait.json");
-
-        let plan = from_substrait_plan(&ctx, &proto).await?;
+        // ./isthmus-cli/build/graal/isthmus --create "create table data (d boolean)" "select not d from data"
+        let proto_plan =
+            read_json("tests/testdata/test_plans/select_not_bool.substrait.json");
+        let ctx = TestSchemaCollector::generate_context_from_plan(&proto_plan);
+        let plan = from_substrait_plan(&ctx, &proto_plan).await?;
 
         assert_eq!(
             format!("{}", plan),
-            "Projection: NOT DATA.a AS EXPR$0\
-            \n  TableScan: DATA projection=[a, b, c, d, e, f]"
+            "Projection: NOT DATA.D AS EXPR$0\
+            \n  TableScan: DATA projection=[D]"
         );
         Ok(())
     }
@@ -61,19 +57,18 @@ mod tests {
         // we don't yet produce such plans.
         // Once we start producing plans with compound signatures, this test can be replaced by the roundtrip tests.
 
-        let ctx = create_context().await?;
-
         // File generated with substrait-java's Isthmus:
-        // ./isthmus-cli/build/graal/isthmus "select sum(d) OVER (PARTITION BY part ORDER BY ord ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING) AS lead_expr from data" -c "create table data (d int, part int, ord int)"
-        let proto = read_json("tests/testdata/test_plans/select_window.substrait.json");
-
-        let plan = from_substrait_plan(&ctx, &proto).await?;
+        // ./isthmus-cli/build/graal/isthmus --create "create table data (d int, part int, ord int)" "select sum(d) OVER (PARTITION BY part ORDER BY ord ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING) AS lead_expr from data"
+        let proto_plan =
+            read_json("tests/testdata/test_plans/select_window.substrait.json");
+        let ctx = TestSchemaCollector::generate_context_from_plan(&proto_plan);
+        let plan = from_substrait_plan(&ctx, &proto_plan).await?;
 
         assert_eq!(
             format!("{}", plan),
-            "Projection: sum(DATA.a) PARTITION BY [DATA.b] ORDER BY [DATA.c ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING AS LEAD_EXPR\
-            \n  WindowAggr: windowExpr=[[sum(DATA.a) PARTITION BY [DATA.b] ORDER BY [DATA.c ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING]]\
-            \n    TableScan: DATA projection=[a, b, c, d, e, f]"
+            "Projection: sum(DATA.D) PARTITION BY [DATA.PART] ORDER BY [DATA.ORD ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING AS LEAD_EXPR\
+            \n  WindowAggr: windowExpr=[[sum(DATA.D) PARTITION BY [DATA.PART] ORDER BY [DATA.ORD ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING]]\
+            \n    TableScan: DATA projection=[D, PART, ORD]"
         );
         Ok(())
     }
@@ -83,11 +78,10 @@ mod tests {
         // DataFusion's Substrait consumer treats all lists as nullable, even if the Substrait plan specifies them as non-nullable.
         // That's because implementing the non-nullability consistently is non-trivial.
         // This test confirms that reading a plan with non-nullable lists works as expected.
-        let ctx = create_context().await?;
-        let proto =
+        let proto_plan =
             read_json("tests/testdata/test_plans/non_nullable_lists.substrait.json");
-
-        let plan = from_substrait_plan(&ctx, &proto).await?;
+        let ctx = TestSchemaCollector::generate_context_from_plan(&proto_plan);
+        let plan = from_substrait_plan(&ctx, &proto_plan).await?;
 
         assert_eq!(format!("{}", &plan), "Values: (List([1, 2]))");
 
@@ -96,18 +90,4 @@ mod tests {
 
         Ok(())
     }
-
-    fn read_json(path: &str) -> Plan {
-        serde_json::from_reader::<_, Plan>(BufReader::new(
-            File::open(path).expect("file not found"),
-        ))
-        .expect("failed to parse json")
-    }
-
-    async fn create_context() -> datafusion::common::Result<SessionContext> {
-        let ctx = SessionContext::new();
-        ctx.register_csv("DATA", "tests/testdata/data.csv", CsvReadOptions::new())
-            .await?;
-        Ok(ctx)
-    }
 }

datafusion/substrait/tests/cases/mod.rs

@@ -21,3 +21,4 @@ mod logical_plans;
 mod roundtrip_logical_plan;
 mod roundtrip_physical_plan;
 mod serialize;
+mod substrait_validations;

datafusion/substrait/tests/cases/substrait_validations.rs

@@ -0,0 +1,101 @@
+// 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.
+
+#[cfg(test)]
+mod tests {
+    use crate::utils::test::read_json;
+    use datafusion::arrow::datatypes::{DataType, Field};
+    use datafusion::catalog_common::TableReference;
+    use datafusion::common::{DFSchema, Result};
+    use datafusion::datasource::empty::EmptyTable;
+    use datafusion::prelude::SessionContext;
+    use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+    use std::collections::HashMap;
+    use std::sync::Arc;
+
+    fn generate_context_with_table(
+        table_name: &str,
+        field_data_type_pairs: Vec<(&str, DataType)>,
+    ) -> Result<SessionContext> {
+        let table_ref = TableReference::bare(table_name);
+        let fields: Vec<(Option<TableReference>, Arc<Field>)> = field_data_type_pairs
+            .into_iter()
+            .map(|pair| {
+                let (field_name, data_type) = pair;
+                (
+                    Some(table_ref.clone()),
+                    Arc::new(Field::new(field_name, data_type, false)),
+                )
+            })
+            .collect();
+
+        let df_schema = DFSchema::new_with_metadata(fields, HashMap::default())?;
+
+        let ctx = SessionContext::new();
+        ctx.register_table(
+            table_ref,
+            Arc::new(EmptyTable::new(df_schema.inner().clone())),
+        )?;
+        Ok(ctx)
+    }
+
+    #[tokio::test]
+    async fn substrait_schema_validation_ignores_field_name_case() -> Result<()> {
+        let proto_plan =
+            read_json("tests/testdata/test_plans/simple_select.substrait.json");
+
+        let ctx = generate_context_with_table("DATA", vec![("a", DataType::Int32)])?;
+        from_substrait_plan(&ctx, &proto_plan).await?;
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn reject_plans_with_mismatched_number_of_fields() -> Result<()> {
+        let proto_plan =
+            read_json("tests/testdata/test_plans/simple_select.substrait.json");
+
+        let ctx = generate_context_with_table(
+            "DATA",
+            vec![("a", DataType::Int32), ("b", DataType::Int32)],
+        )?;
+        let res = from_substrait_plan(&ctx, &proto_plan).await;
+        assert!(res.is_err());
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn reject_plans_with_mismatched_field_names() -> Result<()> {
+        let proto_plan =
+            read_json("tests/testdata/test_plans/simple_select.substrait.json");
+
+        let ctx = generate_context_with_table("DATA", vec![("b", DataType::Date32)])?;
+        let res = from_substrait_plan(&ctx, &proto_plan).await;
+        assert!(res.is_err());
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn reject_plans_with_incompatible_field_types() -> Result<()> {
+        let proto_plan =
+            read_json("tests/testdata/test_plans/simple_select.substrait.json");
+
+        let ctx = generate_context_with_table("DATA", vec![("a", DataType::Date32)])?;
+        let res = from_substrait_plan(&ctx, &proto_plan).await;
+        assert!(res.is_err());
+        Ok(())
+    }
+}