Can compile and run test

Failing on scheduler due to no factories

Tests pass

Back to "no object store available for delta-rs://home-bgardner-workspace"

Switch back to git refs

Author: Brent Gardner

datafusion/core/src/datasource/custom.rs

@@ -0,0 +1,65 @@
+use crate::datasource::TableProvider;
+use crate::execution::context::SessionState;
+use crate::physical_plan::ExecutionPlan;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion_expr::{Expr, TableType};
+use std::any::Any;
+
+use std::sync::Arc;
+
+/// Exists to pass table_type, path, and options from parsing to serialization
+pub struct CustomTable {
+    /// File/Table type used for factory to create TableProvider
+    table_type: String,
+    /// Path used for factory to create TableProvider
+    path: String,
+    /// TableProvider that was created
+    provider: Arc<dyn TableProvider>,
+}
+
+impl CustomTable {
+    /// Creates a new CustomTable
+    pub fn new(table_type: &str, path: &str, provider: Arc<dyn TableProvider>) -> Self {
+        Self {
+            table_type: table_type.to_string(),
+            path: path.to_string(),
+            provider,
+        }
+    }
+
+    /// Returns the table type
+    pub fn get_table_type(&self) -> String {
+        self.table_type.clone()
+    }
+
+    /// Returns the table path
+    pub fn get_path(&self) -> String {
+        self.path.clone()
+    }
+}
+
+#[async_trait]
+impl TableProvider for CustomTable {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        self.provider.schema()
+    }
+
+    fn table_type(&self) -> TableType {
+        self.provider.table_type()
+    }
+
+    async fn scan(
+        &self,
+        ctx: &SessionState,
+        projection: &Option<Vec<usize>>,
+        filters: &[Expr],
+        limit: Option<usize>,
+    ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
+        self.provider.scan(ctx, projection, filters, limit).await
+    }
+}

datafusion/core/src/datasource/datasource.rs

@@ -81,4 +81,12 @@ pub trait TableProvider: Sync + Send {
 pub trait TableProviderFactory: Sync + Send {
     /// Create a TableProvider given name and url
     async fn create(&self, name: &str, 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/mod.rs

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

datafusion/core/src/execution/context.rs

@@ -418,10 +418,11 @@ impl SessionContext {
         cmd: &CreateExternalTable,
     ) -> Result<Arc<DataFrame>> {
         let state = self.state.read().clone();
+        let file_type = cmd.file_type.to_lowercase();
         let factory = &state
             .runtime_env
             .table_factories
-            .get(&cmd.file_type)
+            .get(file_type.as_str())
             .ok_or_else(|| {
                 DataFusionError::Execution(format!(
                     "Unable to find factory for {}",

datafusion/proto/proto/datafusion.proto

@@ -70,6 +70,7 @@ message LogicalPlanNode {
     CreateViewNode create_view = 22;
     DistinctNode distinct = 23;
     ViewTableScanNode view_scan = 24;
+    CustomTableScanNode custom_scan = 25;
   }
 }
 
@@ -118,6 +119,16 @@ message ViewTableScanNode {
   string definition = 5;
 }
 
+// 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;
+}
+
 message ProjectionNode {
   LogicalPlanNode input = 1;
   repeated datafusion.LogicalExprNode expr = 2;

datafusion/proto/src/logical_plan.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::protobuf::{LogicalExprNode, ProjectionColumns};
 use crate::{
     from_proto::{self, parse_expr},
     protobuf::{
@@ -24,6 +25,7 @@ use crate::{
     to_proto,
 };
 use arrow::datatypes::Schema;
+use datafusion::datasource::custom::CustomTable;
 use datafusion::{
     datasource::{
         file_format::{
@@ -410,6 +412,48 @@ impl AsLogicalPlan for LogicalPlanNode {
                 )?
                 .build()
             }
+            LogicalPlanType::CustomScan(scan) => {
+                let schema: Schema = convert_required!(scan.schema)?;
+                let mut projection = None;
+                if let Some(columns) = &scan.projection {
+                    let column_indices = columns
+                        .columns
+                        .iter()
+                        .map(|name| schema.index_of(name))
+                        .collect::<Result<Vec<usize>, _>>()?;
+                    projection = Some(column_indices);
+                }
+
+                let filters = scan
+                    .filters
+                    .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(),
+                )?;
+
+                LogicalPlanBuilder::scan_with_filters(
+                    &scan.table_name,
+                    provider_as_source(provider),
+                    projection,
+                    filters,
+                )?
+                .build()
+            }
             LogicalPlanType::Sort(sort) => {
                 let input: LogicalPlan =
                     into_logical_plan!(sort.input, ctx, extension_codec)?;
@@ -760,55 +804,21 @@ impl AsLogicalPlan for LogicalPlanNode {
                     .collect::<Result<Vec<_>, _>>()?;
 
                 if let Some(listing_table) = source.downcast_ref::<ListingTable>() {
-                    let any = listing_table.options().format.as_any();
-                    let file_format_type = if let Some(parquet) =
-                        any.downcast_ref::<ParquetFormat>()
-                    {
-                        FileFormatType::Parquet(protobuf::ParquetFormat {
-                            enable_pruning: parquet.enable_pruning(),
-                        })
-                    } else if let Some(csv) = any.downcast_ref::<CsvFormat>() {
-                        FileFormatType::Csv(protobuf::CsvFormat {
-                            delimiter: byte_to_string(csv.delimiter())?,
-                            has_header: csv.has_header(),
-                        })
-                    } else if any.is::<AvroFormat>() {
-                        FileFormatType::Avro(protobuf::AvroFormat {})
-                    } else {
-                        return Err(proto_error(format!(
-                            "Error converting file format, {:?} is invalid as a datafusion foramt.",
-                            listing_table.options().format
-                        )));
-                    };
-                    Ok(protobuf::LogicalPlanNode {
-                        logical_plan_type: Some(LogicalPlanType::ListingScan(
-                            protobuf::ListingTableScanNode {
-                                file_format_type: Some(file_format_type),
-                                table_name: table_name.to_owned(),
-                                collect_stat: listing_table.options().collect_stat,
-                                file_extension: listing_table
-                                    .options()
-                                    .file_extension
-                                    .clone(),
-                                table_partition_cols: listing_table
-                                    .options()
-                                    .table_partition_cols
-                                    .clone(),
-                                paths: listing_table
-                                    .table_paths()
-                                    .iter()
-                                    .map(|x| x.to_string())
-                                    .collect(),
-                                schema: Some(schema),
-                                projection,
-                                filters,
-                                target_partitions: listing_table
-                                    .options()
-                                    .target_partitions
-                                    as u32,
-                            },
-                        )),
-                    })
+                    Self::serialize_listing_table(
+                        table_name,
+                        projection,
+                        schema,
+                        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(
@@ -831,8 +841,9 @@ impl AsLogicalPlan for LogicalPlanNode {
                     })
                 } else {
                     Err(DataFusionError::Internal(format!(
-                        "logical plan to_proto unsupported table provider {:?}",
-                        source
+                        "logical plan to_proto unsupported table provider {:?} {}",
+                        source,
+                        plan.display_indent()
                     )))
                 }
             }
@@ -1253,3 +1264,77 @@ 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>,
+        schema: protobuf::Schema,
+        filters: Vec<LogicalExprNode>,
+        listing_table: &ListingTable,
+    ) -> Result<LogicalPlanNode, DataFusionError> {
+        let any = listing_table.options().format.as_any();
+        let file_format_type = if let Some(parquet) = any.downcast_ref::<ParquetFormat>()
+        {
+            FileFormatType::Parquet(protobuf::ParquetFormat {
+                enable_pruning: parquet.enable_pruning(),
+            })
+        } else if let Some(csv) = any.downcast_ref::<CsvFormat>() {
+            FileFormatType::Csv(protobuf::CsvFormat {
+                delimiter: byte_to_string(csv.delimiter())?,
+                has_header: csv.has_header(),
+            })
+        } else if any.is::<AvroFormat>() {
+            FileFormatType::Avro(protobuf::AvroFormat {})
+        } else {
+            return Err(proto_error(format!(
+                "Error converting file format, {:?} is invalid as a datafusion foramt.",
+                listing_table.options().format
+            )));
+        };
+        Ok(protobuf::LogicalPlanNode {
+            logical_plan_type: Some(LogicalPlanType::ListingScan(
+                protobuf::ListingTableScanNode {
+                    file_format_type: Some(file_format_type),
+                    table_name: table_name.to_owned(),
+                    collect_stat: listing_table.options().collect_stat,
+                    file_extension: listing_table.options().file_extension.clone(),
+                    table_partition_cols: listing_table
+                        .options()
+                        .table_partition_cols
+                        .clone(),
+                    paths: listing_table
+                        .table_paths()
+                        .iter()
+                        .map(|x| x.to_string())
+                        .collect(),
+                    schema: Some(schema),
+                    projection,
+                    filters,
+                    target_partitions: listing_table.options().target_partitions as u32,
+                },
+            )),
+        })
+    }
+}