Passing deltalake test

Author: Brent Gardner

datafusion/core/src/datasource/custom.rs

@@ -80,13 +80,12 @@ pub trait TableProvider: Sync + Send {
 #[async_trait]
 pub trait TableProviderFactory: Sync + Send {
     /// Create a TableProvider given name and url
-    async fn create(&self, name: &str, url: &str) -> Result<Arc<dyn TableProvider>>;
+    async fn create(&self, url: &str) -> Result<Arc<dyn TableProvider>>;
 
     /// Create a TableProvider during execution with schema already known from planning
     fn with_schema(
         &self,
         schema: SchemaRef,
-        table_type: &str,
         url: &str,
     ) -> Result<Arc<dyn TableProvider>>;
 }

datafusion/core/src/datasource/datasource.rs

@@ -18,7 +18,6 @@
 //! DataFusion data sources
 
 #![allow(clippy::module_inception)]
-pub mod custom;
 pub mod datasource;
 pub mod default_table_source;
 pub mod empty;

datafusion/core/src/datasource/mod.rs

@@ -430,7 +430,7 @@ impl SessionContext {
                 ))
             })?;
         let table = (*factory)
-            .create(cmd.name.as_str(), cmd.location.as_str())
+            .create(cmd.location.as_str())
             .await?;
         self.register_table(cmd.name.as_str(), table)?;
         let plan = LogicalPlanBuilder::empty(false).build()?;

datafusion/core/src/execution/context.rs

@@ -122,11 +122,10 @@ message ViewTableScanNode {
 // Logical Plan to Scan a CustomTableProvider registered at runtime
 message CustomTableScanNode {
   string table_name = 1;
-  string path = 2;
-  string table_type = 3;
-  ProjectionColumns projection = 4;
-  datafusion.Schema schema = 5;
-  repeated datafusion.LogicalExprNode filters = 6;
+  ProjectionColumns projection = 2;
+  datafusion.Schema schema = 3;
+  repeated datafusion.LogicalExprNode filters = 4;
+  bytes custom_table_data = 5;
 }
 
 message ProjectionNode {

datafusion/proto/proto/datafusion.proto

@@ -16,6 +16,8 @@
 // under the License.
 
 //! Serialization / Deserialization to Bytes
+use std::sync::Arc;
+use arrow::datatypes::SchemaRef;
 use crate::logical_plan::{AsLogicalPlan, LogicalExtensionCodec};
 use crate::{from_proto::parse_expr, protobuf};
 use datafusion_common::{DataFusionError, Result};
@@ -24,6 +26,7 @@ use prost::{
     bytes::{Bytes, BytesMut},
     Message,
 };
+use datafusion::datasource::TableProvider;
 
 // Reexport Bytes which appears in the API
 use datafusion::execution::registry::FunctionRegistry;
@@ -180,6 +183,18 @@ impl LogicalExtensionCodec for DefaultExtensionCodec {
             "No extension codec provided".to_string(),
         ))
     }
+
+    fn try_decode_table_provider(&self, _buf: &[u8], _schema: SchemaRef, _ctx: &SessionContext) -> std::result::Result<Arc<dyn TableProvider>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "No extension codec provided".to_string(),
+        ))
+    }
+
+    fn try_encode_table_provider(&self, _node: Arc<dyn TableProvider>, _buf: &mut Vec<u8>) -> std::result::Result<(), DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "No extension codec provided".to_string(),
+        ))
+    }
 }
 
 #[cfg(test)]

datafusion/proto/src/bytes/mod.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::protobuf::{LogicalExprNode, ProjectionColumns};
+use crate::protobuf::{CustomTableScanNode, LogicalExprNode, ProjectionColumns};
 use crate::{
     from_proto::{self, parse_expr},
     protobuf::{
@@ -24,8 +24,7 @@ use crate::{
     },
     to_proto,
 };
-use arrow::datatypes::Schema;
-use datafusion::datasource::custom::CustomTable;
+use arrow::datatypes::{Schema, SchemaRef};
 use datafusion::{
     datasource::{
         file_format::{
@@ -37,7 +36,7 @@ use datafusion::{
     datasource::{provider_as_source, source_as_provider},
     prelude::SessionContext,
 };
-use datafusion_common::{Column, DataFusionError};
+use datafusion_common::{Column, context, DataFusionError};
 use datafusion_expr::{
     logical_plan::{
         Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView,
@@ -50,6 +49,8 @@ use prost::bytes::BufMut;
 use prost::Message;
 use std::fmt::Debug;
 use std::sync::Arc;
+use datafusion::datasource::TableProvider;
+use crate::protobuf::logical_plan_node::LogicalPlanType::CustomScan;
 
 fn byte_to_string(b: u8) -> Result<String, DataFusionError> {
     let b = &[b];
@@ -108,6 +109,19 @@ pub trait LogicalExtensionCodec: Debug + Send + Sync {
         node: &Extension,
         buf: &mut Vec<u8>,
     ) -> Result<(), DataFusionError>;
+
+    fn try_decode_table_provider(
+        &self,
+        buf: &[u8],
+        schema: SchemaRef,
+        ctx: &SessionContext,
+    ) -> Result<Arc<dyn TableProvider>, DataFusionError>;
+
+    fn try_encode_table_provider(
+        &self,
+        node: Arc<dyn TableProvider>,
+        buf: &mut Vec<u8>,
+    ) -> Result<(), DataFusionError>;
 }
 
 #[derive(Debug, Clone)]
@@ -134,6 +148,18 @@ impl LogicalExtensionCodec for DefaultLogicalExtensionCodec {
             "LogicalExtensionCodec is not provided".to_string(),
         ))
     }
+
+    fn try_decode_table_provider(&self, _buf: &[u8], _schema: SchemaRef, _ctx: &SessionContext) -> Result<Arc<dyn TableProvider>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "LogicalExtensionCodec is not provided".to_string(),
+        ))
+    }
+
+    fn try_encode_table_provider(&self, _node: Arc<dyn TableProvider>, _buf: &mut Vec<u8>) -> Result<(), DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "LogicalExtensionCodec is not provided".to_string(),
+        ))
+    }
 }
 
 #[macro_export]
@@ -414,6 +440,7 @@ impl AsLogicalPlan for LogicalPlanNode {
             }
             LogicalPlanType::CustomScan(scan) => {
                 let schema: Schema = convert_required!(scan.schema)?;
+                let schema = Arc::new(schema);
                 let mut projection = None;
                 if let Some(columns) = &scan.projection {
                     let column_indices = columns
@@ -429,22 +456,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                     .iter()
                     .map(|expr| parse_expr(expr, ctx))
                     .collect::<Result<Vec<_>, _>>()?;
-                let state = ctx.state.read();
-                let factory = state
-                    .runtime_env
-                    .table_factories
-                    .get(&scan.table_type)
-                    .ok_or_else(|| {
-                        DataFusionError::Plan(format!(
-                            "Planner unable to find factory for {}",
-                            scan.table_type
-                        ))
-                    })?;
-                let provider = (*factory).with_schema(
-                    Arc::new(schema),
-                    scan.table_type.as_str(),
-                    scan.path.as_str(),
-                )?;
+                let provider = extension_codec.try_decode_table_provider(&scan.custom_table_data, schema, ctx)?;
 
                 LogicalPlanBuilder::scan_with_filters(
                     &scan.table_name,
@@ -780,8 +792,8 @@ impl AsLogicalPlan for LogicalPlanNode {
                 projection,
                 ..
             }) => {
-                let source = source_as_provider(source)?;
-                let schema = source.schema();
+                let provider = source_as_provider(source)?;
+                let schema = provider.schema();
                 let source = source.as_any();
 
                 let projection = match projection {
@@ -811,14 +823,6 @@ impl AsLogicalPlan for LogicalPlanNode {
                         filters,
                         listing_table,
                     )
-                } else if let Some(custom_table) = source.downcast_ref::<CustomTable>() {
-                    Self::serialize_custom_table(
-                        table_name,
-                        projection,
-                        schema,
-                        filters,
-                        custom_table,
-                    )
                 } else if let Some(view_table) = source.downcast_ref::<ViewTable>() {
                     Ok(protobuf::LogicalPlanNode {
                         logical_plan_type: Some(LogicalPlanType::ViewScan(Box::new(
@@ -840,11 +844,20 @@ impl AsLogicalPlan for LogicalPlanNode {
                         ))),
                     })
                 } else {
-                    Err(DataFusionError::Internal(format!(
-                        "logical plan to_proto unsupported table provider {:?} {}",
-                        source,
-                        plan.display_indent()
-                    )))
+                    let mut bytes = vec![];
+                    extension_codec.try_encode_table_provider(provider, &mut bytes)
+                        .map_err(|e| context!("Error serializing custom table", e))?;
+                    let scan = CustomScan(CustomTableScanNode {
+                        table_name: table_name.clone(),
+                        projection,
+                        schema: Some(schema),
+                        filters,
+                        custom_table_data: bytes
+                    });
+                    let node = LogicalPlanNode {
+                        logical_plan_type: Some(scan)
+                    };
+                    Ok(node)
                 }
             }
             LogicalPlan::Projection(Projection {
@@ -1266,27 +1279,6 @@ impl AsLogicalPlan for LogicalPlanNode {
 }
 
 impl LogicalPlanNode {
-    fn serialize_custom_table(
-        table_name: &String,
-        projection: Option<ProjectionColumns>,
-        schema: protobuf::Schema,
-        filters: Vec<LogicalExprNode>,
-        custom_table: &CustomTable,
-    ) -> Result<LogicalPlanNode, DataFusionError> {
-        Ok(protobuf::LogicalPlanNode {
-            logical_plan_type: Some(LogicalPlanType::CustomScan(
-                protobuf::CustomTableScanNode {
-                    table_name: table_name.to_owned(),
-                    table_type: custom_table.get_table_type(),
-                    path: custom_table.get_path(),
-                    schema: Some(schema),
-                    projection,
-                    filters,
-                },
-            )),
-        })
-    }
-
     fn serialize_listing_table(
         table_name: &String,
         projection: Option<ProjectionColumns>,