fix: parse table name into TableReference on converting substrait read (

#5716) * fix: parse table name into TableReference on converting substrait read Signed-off-by: Ruihang Xia <[email protected]> * add tests Signed-off-by: Ruihang Xia <[email protected]> * move functionality into TableReference Signed-off-by: Ruihang Xia <[email protected]> --------- Signed-off-by: Ruihang Xia <[email protected]>
apache · Mar 25, 2023 · 7b67d28 · 7b67d28
1 parent 9e8d928
commit 7b67d28
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 4 deletions.
diff --git a/datafusion/common/src/table_reference.rs b/datafusion/common/src/table_reference.rs
@@ -317,6 +317,25 @@ impl<'a> TableReference<'a> {
             _ => Self::Bare { table: s.into() },
         }
     }
+
+    /// Decompose a [`TableReference`] to separate parts. The result vector contains
+    /// at most three elements in the following sequence:
+    /// ```no_rust
+    /// [<catalog>, <schema>, table]
+    /// ```
+    pub fn to_vec(&self) -> Vec<String> {
+        match self {
+            TableReference::Bare { table } => vec![table.to_string()],
+            TableReference::Partial { schema, table } => {
+                vec![schema.to_string(), table.to_string()]
+            }
+            TableReference::Full {
+                catalog,
+                schema,
+                table,
+            } => vec![catalog.to_string(), schema.to_string(), table.to_string()],
+        }
+    }
 }
 
 /// Parse a `String` into a OwnedTableReference as a multipart SQL identifier.
@@ -408,4 +427,26 @@ mod tests {
         let actual = TableReference::from("TABLE()");
         assert_eq!(expected, actual);
     }
+
+    #[test]
+    fn test_table_reference_to_vector() {
+        let table_reference = TableReference::parse_str("table");
+        assert_eq!(vec!["table".to_string()], table_reference.to_vec());
+
+        let table_reference = TableReference::parse_str("schema.table");
+        assert_eq!(
+            vec!["schema".to_string(), "table".to_string()],
+            table_reference.to_vec()
+        );
+
+        let table_reference = TableReference::parse_str("catalog.schema.table");
+        assert_eq!(
+            vec![
+                "catalog".to_string(),
+                "schema".to_string(),
+                "table".to_string()
+            ],
+            table_reference.to_vec()
+        );
+    }
 }
diff --git a/datafusion/substrait/src/logical_plan/producer.rs b/datafusion/substrait/src/logical_plan/producer.rs
@@ -130,7 +130,7 @@ pub fn to_substrait_rel(
                         }),
                         advanced_extension: None,
                         read_type: Some(ReadType::NamedTable(NamedTable {
-                            names: vec![scan.table_name.to_string()],
+                            names: scan.table_name.to_vec(),
                             advanced_extension: None,
                         })),
                     }))),

diff --git a/datafusion/substrait/tests/roundtrip_logical_plan.rs b/datafusion/substrait/tests/roundtrip_logical_plan.rs
@@ -255,6 +255,16 @@ mod tests {
         roundtrip("SELECT RANK() OVER (PARTITION BY a ORDER BY b), d, SUM(b) OVER (PARTITION BY a) FROM data;").await
     }
 
+    #[tokio::test]
+    async fn qualified_schema_table_reference() -> Result<()> {
+        roundtrip("SELECT * FROM public.data;").await
+    }
+
+    #[tokio::test]
+    async fn qualified_catalog_schema_table_reference() -> Result<()> {
+        roundtrip("SELECT * FROM datafusion.public.data;").await
+    }
+
     async fn assert_expected_plan(sql: &str, expected_plan_str: &str) -> Result<()> {
         let mut ctx = create_context().await?;
         let df = ctx.sql(sql).await?;
@@ -306,7 +316,6 @@ mod tests {
         Ok(())
     }
 
-    #[allow(deprecated)]
     async fn roundtrip(sql: &str) -> Result<()> {
         let mut ctx = create_context().await?;
         let df = ctx.sql(sql).await?;

diff --git a/datafusion/substrait/tests/serialize.rs b/datafusion/substrait/tests/serialize.rs
@@ -40,8 +40,6 @@ mod tests {
         let proto = serializer::deserialize(path).await?;
         // Check plan equality
         let plan = from_substrait_plan(&mut ctx, &proto).await?;
-        // #[allow(deprecated)]
-        // let plan = ctx.optimize(&plan)?;
         let plan_str_ref = format!("{plan_ref:?}");
         let plan_str = format!("{plan:?}");
         assert_eq!(plan_str_ref, plan_str);