feat(datafusion-proto): allow TableSource to be serialized

Currently, only instances of `TableProvider` are considered by
`LogicalExtensionCodec`, and are automatically wrapped in a
`DefaultTableSource` when deserializing. That doesn't work with custom
extensions.

Fixes #16749

Author: colinmarc

datafusion/catalog/src/default_table_source.rs

@@ -45,6 +45,26 @@ impl DefaultTableSource {
     pub fn new(table_provider: Arc<dyn TableProvider>) -> Self {
         Self { table_provider }
     }
+
+    /// Attempt to downcast a TableSource to DefaultTableSource and access the
+    /// TableProvider. This will only work with a TableSource created by DataFusion.
+    pub fn unwrap_provider<T: TableProvider + 'static>(
+        source: &Arc<dyn TableSource>,
+    ) -> datafusion_common::Result<&T> {
+        if let Some(source) = source
+            .as_ref()
+            .as_any()
+            .downcast_ref::<DefaultTableSource>()
+        {
+            if let Some(provider) =
+                source.table_provider.as_ref().as_any().downcast_ref::<T>()
+            {
+                return Ok(provider);
+            }
+        }
+
+        internal_err!("TableSource was not expected type")
+    }
 }
 
 impl TableSource for DefaultTableSource {
@@ -87,7 +107,7 @@ impl TableSource for DefaultTableSource {
     }
 }
 
-/// Wrap TableProvider in TableSource
+/// Wrap a TableProvider as a TableSource.
 pub fn provider_as_source(
     table_provider: Arc<dyn TableProvider>,
 ) -> Arc<dyn TableSource> {

datafusion/core/src/datasource/mod.rs

@@ -30,9 +30,7 @@ pub mod provider;
 mod view_test;
 
 // backwards compatibility
-pub use self::default_table_source::{
-    provider_as_source, source_as_provider, DefaultTableSource,
-};
+pub use self::default_table_source::{provider_as_source, DefaultTableSource};
 pub use self::memory::MemTable;
 pub use self::view::ViewTable;
 pub use crate::catalog::TableProvider;

datafusion/core/src/physical_planner.rs

@@ -24,7 +24,7 @@ use std::sync::Arc;
 use crate::datasource::file_format::file_type_to_format;
 use crate::datasource::listing::ListingTableUrl;
 use crate::datasource::physical_plan::FileSinkConfig;
-use crate::datasource::{source_as_provider, DefaultTableSource};
+use crate::datasource::DefaultTableSource;
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::{ExecutionProps, SessionState};
 use crate::logical_expr::utils::generate_sort_key;
@@ -60,6 +60,7 @@ use crate::schema_equivalence::schema_satisfied_by;
 use arrow::array::{builder::StringBuilder, RecordBatch};
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Schema, SchemaRef};
+use datafusion_catalog::default_table_source::source_as_provider;
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,

datafusion/proto/src/logical_plan/file_formats.rs

@@ -180,20 +180,20 @@ impl LogicalExtensionCodec for CsvLogicalExtensionCodec {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         _buf: &[u8],
         _table_ref: &TableReference,
         _schema: arrow::datatypes::SchemaRef,
         _ctx: &SessionContext,
-    ) -> datafusion_common::Result<Arc<dyn datafusion::datasource::TableProvider>> {
+    ) -> datafusion_common::Result<Arc<dyn datafusion_expr::TableSource>> {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         _table_ref: &TableReference,
-        _node: Arc<dyn datafusion::datasource::TableProvider>,
+        _node: Arc<dyn datafusion_expr::TableSource>,
         _buf: &mut Vec<u8>,
     ) -> datafusion_common::Result<()> {
         not_impl_err!("Method not implemented")
@@ -287,20 +287,20 @@ impl LogicalExtensionCodec for JsonLogicalExtensionCodec {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         _buf: &[u8],
         _table_ref: &TableReference,
         _schema: arrow::datatypes::SchemaRef,
         _ctx: &SessionContext,
-    ) -> datafusion_common::Result<Arc<dyn datafusion::datasource::TableProvider>> {
+    ) -> datafusion_common::Result<Arc<dyn datafusion_expr::TableSource>> {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         _table_ref: &TableReference,
-        _node: Arc<dyn datafusion::datasource::TableProvider>,
+        _node: Arc<dyn datafusion_expr::TableSource>,
         _buf: &mut Vec<u8>,
     ) -> datafusion_common::Result<()> {
         not_impl_err!("Method not implemented")
@@ -603,20 +603,20 @@ impl LogicalExtensionCodec for ParquetLogicalExtensionCodec {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         _buf: &[u8],
         _table_ref: &TableReference,
         _schema: arrow::datatypes::SchemaRef,
         _ctx: &SessionContext,
-    ) -> datafusion_common::Result<Arc<dyn datafusion::datasource::TableProvider>> {
+    ) -> datafusion_common::Result<Arc<dyn datafusion_expr::TableSource>> {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         _table_ref: &TableReference,
-        _node: Arc<dyn datafusion::datasource::TableProvider>,
+        _node: Arc<dyn datafusion_expr::TableSource>,
         _buf: &mut Vec<u8>,
     ) -> datafusion_common::Result<()> {
         not_impl_err!("Method not implemented")
@@ -689,20 +689,20 @@ impl LogicalExtensionCodec for ArrowLogicalExtensionCodec {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         _buf: &[u8],
         _table_ref: &TableReference,
         _schema: arrow::datatypes::SchemaRef,
         _ctx: &SessionContext,
-    ) -> datafusion_common::Result<Arc<dyn datafusion::datasource::TableProvider>> {
+    ) -> datafusion_common::Result<Arc<dyn datafusion_expr::TableSource>> {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         _table_ref: &TableReference,
-        _node: Arc<dyn datafusion::datasource::TableProvider>,
+        _node: Arc<dyn datafusion_expr::TableSource>,
         _buf: &mut Vec<u8>,
     ) -> datafusion_common::Result<()> {
         not_impl_err!("Method not implemented")
@@ -747,20 +747,20 @@ impl LogicalExtensionCodec for AvroLogicalExtensionCodec {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         _buf: &[u8],
         _table_ref: &TableReference,
         _schema: arrow::datatypes::SchemaRef,
         _cts: &SessionContext,
-    ) -> datafusion_common::Result<Arc<dyn datafusion::datasource::TableProvider>> {
+    ) -> datafusion_common::Result<Arc<dyn datafusion_expr::TableSource>> {
         not_impl_err!("Method not implemented")
     }
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         _table_ref: &TableReference,
-        _node: Arc<dyn datafusion::datasource::TableProvider>,
+        _node: Arc<dyn datafusion_expr::TableSource>,
         _buf: &mut Vec<u8>,
     ) -> datafusion_common::Result<()> {
         not_impl_err!("Method not implemented")

datafusion/proto/src/logical_plan/mod.rs

@@ -42,16 +42,16 @@ use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::file_format::{
     file_type_to_format, format_as_file_type, FileFormatFactory,
 };
+use datafusion::datasource::DefaultTableSource;
 use datafusion::{
+    datasource::provider_as_source,
     datasource::{
         file_format::{
             csv::CsvFormat, json::JsonFormat as OtherNdJsonFormat, FileFormat,
         },
         listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl},
         view::ViewTable,
-        TableProvider,
     },
-    datasource::{provider_as_source, source_as_provider},
     prelude::SessionContext,
 };
 use datafusion_common::file_options::file_type::FileType;
@@ -118,18 +118,18 @@ pub trait LogicalExtensionCodec: Debug + Send + Sync {
 
     fn try_encode(&self, node: &Extension, buf: &mut Vec<u8>) -> Result<()>;
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         buf: &[u8],
         table_ref: &TableReference,
         schema: SchemaRef,
         ctx: &SessionContext,
-    ) -> Result<Arc<dyn TableProvider>>;
+    ) -> Result<Arc<dyn TableSource>>;
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         table_ref: &TableReference,
-        node: Arc<dyn TableProvider>,
+        node: Arc<dyn TableSource>,
         buf: &mut Vec<u8>,
     ) -> Result<()>;
 
@@ -193,20 +193,20 @@ impl LogicalExtensionCodec for DefaultLogicalExtensionCodec {
         not_impl_err!("LogicalExtensionCodec is not provided")
     }
 
-    fn try_decode_table_provider(
+    fn try_decode_table_source(
         &self,
         _buf: &[u8],
         _table_ref: &TableReference,
         _schema: SchemaRef,
         _ctx: &SessionContext,
-    ) -> Result<Arc<dyn TableProvider>> {
+    ) -> Result<Arc<dyn TableSource>> {
         not_impl_err!("LogicalExtensionCodec is not provided")
     }
 
-    fn try_encode_table_provider(
+    fn try_encode_table_source(
         &self,
         _table_ref: &TableReference,
-        _node: Arc<dyn TableProvider>,
+        _node: Arc<dyn TableSource>,
         _buf: &mut Vec<u8>,
     ) -> Result<()> {
         not_impl_err!("LogicalExtensionCodec is not provided")
@@ -440,7 +440,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                         }
                         #[cfg_attr(not(feature = "avro"), allow(unused_variables))]
                         FileFormatType::Avro(..) => {
-                            #[cfg(feature = "avro")] 
+                            #[cfg(feature = "avro")]
                             {
                                 Arc::new(AvroFormat)
                             }
@@ -520,18 +520,15 @@ impl AsLogicalPlan for LogicalPlanNode {
                 let table_name =
                     from_table_reference(scan.table_name.as_ref(), "CustomScan")?;
 
-                let provider = extension_codec.try_decode_table_provider(
+                let source = extension_codec.try_decode_table_source(
                     &scan.custom_table_data,
                     &table_name,
                     schema,
                     ctx,
                 )?;
 
                 LogicalPlanBuilder::scan_with_filters(
-                    table_name,
-                    provider_as_source(provider),
-                    projection,
-                    filters,
+                    table_name, source, projection, filters,
                 )?
                 .build()
             }
@@ -1029,9 +1026,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                 projection,
                 ..
             }) => {
-                let provider = source_as_provider(source)?;
-                let schema = provider.schema();
-                let source = provider.as_any();
+                let schema = source.schema();
 
                 let projection = match projection {
                     None => None,
@@ -1049,7 +1044,9 @@ impl AsLogicalPlan for LogicalPlanNode {
                 let filters: Vec<protobuf::LogicalExprNode> =
                     serialize_exprs(filters, extension_codec)?;
 
-                if let Some(listing_table) = source.downcast_ref::<ListingTable>() {
+                if let Ok(listing_table) =
+                    DefaultTableSource::unwrap_provider::<ListingTable>(source)
+                {
                     let any = listing_table.options().format.as_any();
                     let file_format_type = {
                         let mut maybe_some_type = None;
@@ -1158,7 +1155,9 @@ impl AsLogicalPlan for LogicalPlanNode {
                             },
                         )),
                     })
-                } else if let Some(view_table) = source.downcast_ref::<ViewTable>() {
+                } else if let Ok(view_table) =
+                    DefaultTableSource::unwrap_provider::<ViewTable>(source)
+                {
                     let schema: protobuf::Schema = schema.as_ref().try_into()?;
                     Ok(LogicalPlanNode {
                         logical_plan_type: Some(LogicalPlanType::ViewScan(Box::new(
@@ -1179,7 +1178,8 @@ impl AsLogicalPlan for LogicalPlanNode {
                             },
                         ))),
                     })
-                } else if let Some(cte_work_table) = source.downcast_ref::<CteWorkTable>()
+                } else if let Ok(cte_work_table) =
+                    DefaultTableSource::unwrap_provider::<CteWorkTable>(source)
                 {
                     let name = cte_work_table.name().to_string();
                     let schema = cte_work_table.schema();
@@ -1197,7 +1197,11 @@ impl AsLogicalPlan for LogicalPlanNode {
                     let schema: protobuf::Schema = schema.as_ref().try_into()?;
                     let mut bytes = vec![];
                     extension_codec
-                        .try_encode_table_provider(table_name, provider, &mut bytes)
+                        .try_encode_table_source(
+                            table_name,
+                            Arc::clone(source),
+                            &mut bytes,
+                        )
                         .map_err(|e| context!("Error serializing custom table", e))?;
                     let scan = CustomScan(CustomTableScanNode {
                         table_name: Some(table_name.clone().into()),