Add roundtrip test

Author: Brent Gardner

datafusion/core/src/datasource/custom.rs

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+//! A TableProvider that wraps tables registered through TableProviderFactories so they can be reified during serde
+
 use crate::datasource::TableProvider;
 use crate::execution::context::SessionState;
 use crate::physical_plan::ExecutionPlan;

datafusion/core/src/test_util.rs

@@ -17,13 +17,20 @@
 
 //! Utility functions to make testing DataFusion based crates easier
 
+use std::any::Any;
 use std::collections::BTreeMap;
 use std::{env, error::Error, path::PathBuf, sync::Arc};
 
-use crate::datasource::{empty::EmptyTable, provider_as_source};
+use crate::datasource::custom::CustomTable;
+use crate::datasource::datasource::TableProviderFactory;
+use crate::datasource::{empty::EmptyTable, provider_as_source, TableProvider};
+use crate::execution::context::SessionState;
 use crate::logical_expr::{LogicalPlanBuilder, UNNAMED_TABLE};
+use crate::physical_plan::ExecutionPlan;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use async_trait::async_trait;
 use datafusion_common::DataFusionError;
+use datafusion_expr::{Expr, TableType};
 
 /// Compares formatted output of a record batch with an expected
 /// vector of strings, with the result of pretty formatting record
@@ -317,6 +324,71 @@ pub fn aggr_test_schema_with_missing_col() -> SchemaRef {
     Arc::new(schema)
 }
 
+/// TableFactory for tests
+pub struct TestTableFactory {}
+
+#[async_trait]
+impl TableProviderFactory for TestTableFactory {
+    async fn create(
+        &self,
+        _name: &str,
+        url: &str,
+    ) -> datafusion_common::Result<Arc<dyn TableProvider>> {
+        Ok(Arc::new(CustomTable::new(
+            "deltatable",
+            url,
+            Arc::new(TestTableProvider {}),
+        )))
+    }
+
+    fn with_schema(
+        &self,
+        _schema: SchemaRef,
+        table_type: &str,
+        url: &str,
+    ) -> datafusion_common::Result<Arc<dyn TableProvider>> {
+        Ok(Arc::new(CustomTable::new(
+            table_type,
+            url,
+            Arc::new(TestTableProvider {}),
+        )))
+    }
+}
+
+/// TableProvider for testing purposes
+pub struct TestTableProvider {}
+
+impl TestTableProvider {}
+
+#[async_trait]
+impl TableProvider for TestTableProvider {
+    fn as_any(&self) -> &dyn Any {
+        unimplemented!("TestTableProvider is a stub for testing.")
+    }
+
+    fn schema(&self) -> SchemaRef {
+        let schema = Schema::new(vec![
+            Field::new("a", DataType::Int64, true),
+            Field::new("b", DataType::Decimal128(15, 2), true),
+        ]);
+        Arc::new(schema)
+    }
+
+    fn table_type(&self) -> TableType {
+        unimplemented!("TestTableProvider is a stub for testing.")
+    }
+
+    async fn scan(
+        &self,
+        _ctx: &SessionState,
+        _projection: &Option<Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
+        unimplemented!("TestTableProvider is a stub for testing.")
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion/core/tests/sql/create_drop.rs

@@ -15,15 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use async_trait::async_trait;
-use std::any::Any;
 use std::collections::HashMap;
 use std::io::Write;
 
 use datafusion::datasource::datasource::TableProviderFactory;
-use datafusion::execution::context::SessionState;
 use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
-use datafusion_expr::TableType;
+use datafusion::test_util::TestTableFactory;
 use tempfile::TempDir;
 
 use super::*;
@@ -369,53 +366,6 @@ async fn create_pipe_delimited_csv_table() -> Result<()> {
     Ok(())
 }
 
-struct TestTableProvider {}
-
-impl TestTableProvider {}
-
-#[async_trait]
-impl TableProvider for TestTableProvider {
-    fn as_any(&self) -> &dyn Any {
-        unimplemented!("TestTableProvider is a stub for testing.")
-    }
-
-    fn schema(&self) -> SchemaRef {
-        unimplemented!("TestTableProvider is a stub for testing.")
-    }
-
-    fn table_type(&self) -> TableType {
-        unimplemented!("TestTableProvider is a stub for testing.")
-    }
-
-    async fn scan(
-        &self,
-        _ctx: &SessionState,
-        _projection: &Option<Vec<usize>>,
-        _filters: &[Expr],
-        _limit: Option<usize>,
-    ) -> Result<Arc<dyn ExecutionPlan>> {
-        unimplemented!("TestTableProvider is a stub for testing.")
-    }
-}
-
-struct TestTableFactory {}
-
-#[async_trait]
-impl TableProviderFactory for TestTableFactory {
-    async fn create(&self, _name: &str, _url: &str) -> Result<Arc<dyn TableProvider>> {
-        Ok(Arc::new(TestTableProvider {}))
-    }
-
-    fn with_schema(
-        &self,
-        _schema: SchemaRef,
-        _table_type: &str,
-        _url: &str,
-    ) -> Result<Arc<dyn TableProvider>> {
-        Ok(Arc::new(TestTableProvider {}))
-    }
-}
-
 #[tokio::test]
 async fn create_custom_table() -> Result<()> {
     let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =

datafusion/proto/src/lib.rs

@@ -59,8 +59,13 @@ mod roundtrip_tests {
             TimeUnit, UnionMode,
         },
     };
+    use datafusion::datasource::datasource::TableProviderFactory;
+    use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
     use datafusion::physical_plan::functions::make_scalar_function;
-    use datafusion::prelude::{create_udf, CsvReadOptions, SessionContext};
+    use datafusion::prelude::{
+        create_udf, CsvReadOptions, SessionConfig, SessionContext,
+    };
+    use datafusion::test_util::TestTableFactory;
     use datafusion_common::{DFSchemaRef, DataFusionError, ScalarValue};
     use datafusion_expr::create_udaf;
     use datafusion_expr::expr::{Between, BinaryExpr, Case, GroupingSet, Like};
@@ -71,6 +76,7 @@ mod roundtrip_tests {
     };
     use prost::Message;
     use std::any::Any;
+    use std::collections::HashMap;
     use std::fmt;
     use std::fmt::Debug;
     use std::fmt::Formatter;
@@ -129,6 +135,26 @@ mod roundtrip_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn roundtrip_custom_tables() -> Result<(), DataFusionError> {
+        let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =
+            HashMap::new();
+        table_factories.insert("deltatable".to_string(), Arc::new(TestTableFactory {}));
+        let cfg = RuntimeConfig::new().with_table_factories(table_factories);
+        let env = RuntimeEnv::new(cfg).unwrap();
+        let ses = SessionConfig::new();
+        let ctx = SessionContext::with_config_rt(ses, Arc::new(env));
+
+        let sql = "CREATE EXTERNAL TABLE dt STORED AS DELTATABLE LOCATION 's3://bucket/schema/table';";
+        ctx.sql(sql).await.unwrap();
+
+        let scan = ctx.table("dt")?.to_logical_plan()?;
+        let bytes = logical_plan_to_bytes(&scan)?;
+        let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx)?;
+        assert_eq!(format!("{:?}", scan), format!("{:?}", logical_round_trip));
+        Ok(())
+    }
+
     #[tokio::test]
     async fn roundtrip_logical_plan_aggregation() -> Result<(), DataFusionError> {
         let ctx = SessionContext::new();