Custom table provider factories

Author: Brent Gardner

.github/workflows/rust.yml

@@ -280,7 +280,9 @@ jobs:
           rustup default stable
           rustup component add rustfmt
       - name: Run
-        run: ci/scripts/rust_fmt.sh
+        run: |
+          echo '' > datafusion/proto/src/generated/datafusion.rs
+          ci/scripts/rust_fmt.sh
 
   coverage:
     name: coverage

datafusion/core/Cargo.toml

@@ -93,6 +93,7 @@ url = "2.2"
 uuid = { version = "1.0", features = ["v4"] }
 
 [dev-dependencies]
+async-trait = "0.1.53"
 criterion = "0.3"
 csv = "1.1.6"
 ctor = "0.1.22"

datafusion/core/src/datasource/datasource.rs

@@ -72,3 +72,12 @@ pub trait TableProvider: Sync + Send {
         Ok(TableProviderFilterPushDown::Unsupported)
     }
 }
+
+/// A factory which creates [`TableProvider`]s at runtime given a URL.
+/// 
+/// For example, this can be used to create a table "on the fly" 
+/// from a directory of files only when that name is referenced.  
+pub trait TableProviderFactory: Sync + Send {
+    /// Create a TableProvider given name and url
+    fn create(&self, name: &str, url: &str) -> Arc<dyn TableProvider>;
+}

datafusion/core/src/logical_plan/mod.rs

@@ -45,10 +45,9 @@ pub use datafusion_expr::{
             build_join_schema, union_with_alias, LogicalPlanBuilder, UNNAMED_TABLE,
         },
         CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateMemoryTable,
-        CreateView, CrossJoin, DropTable, EmptyRelation, FileType, JoinConstraint,
-        JoinType, Limit, LogicalPlan, Partitioning, PlanType, PlanVisitor, Repartition,
-        StringifiedPlan, Subquery, TableScan, ToStringifiedPlan, Union,
-        UserDefinedLogicalNode, Values,
+        CreateView, CrossJoin, DropTable, EmptyRelation, JoinConstraint, JoinType, Limit,
+        LogicalPlan, Partitioning, PlanType, PlanVisitor, Repartition, StringifiedPlan,
+        Subquery, TableScan, ToStringifiedPlan, Union, UserDefinedLogicalNode, Values,
     },
     lower, lpad, ltrim, max, md5, min, not_exists, not_in_subquery, now, nullif,
     octet_length, or, power, random, regexp_match, regexp_replace, repeat, replace,

datafusion/core/src/logical_plan/plan.rs

@@ -23,8 +23,8 @@ pub use datafusion_expr::{
         display::{GraphvizVisitor, IndentVisitor},
         Aggregate, Analyze, CreateCatalog, CreateCatalogSchema, CreateExternalTable,
         CreateMemoryTable, CreateView, CrossJoin, Distinct, DropTable, EmptyRelation,
-        Explain, Extension, FileType, Filter, Join, JoinConstraint, JoinType, Limit,
-        LogicalPlan, Partitioning, PlanType, PlanVisitor, Projection, Repartition, Sort,
+        Explain, Extension, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan,
+        Partitioning, PlanType, PlanVisitor, Projection, Repartition, Sort,
         StringifiedPlan, Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union,
         UserDefinedLogicalNode, Values, Window,
     },

datafusion/core/tests/sql/create_drop.rs

@@ -15,8 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use async_trait::async_trait;
+use std::any::Any;
 use std::io::Write;
 
+use datafusion::datasource::datasource::TableProviderFactory;
+use datafusion::execution::context::SessionState;
+use datafusion_expr::TableType;
 use tempfile::TempDir;
 
 use super::*;
@@ -360,6 +365,74 @@ 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 {}
+
+impl TableProviderFactory for TestTableFactory {
+    fn create(&self, _name: &str, _path: &str) -> Arc<dyn TableProvider> {
+        Arc::new(TestTableProvider {})
+    }
+}
+
+#[tokio::test]
+async fn create_custom_table() -> Result<()> {
+    let mut ctx = SessionContext::new();
+    ctx.register_table_factory("DELTATABLE", Arc::new(TestTableFactory {}));
+
+    let sql = "CREATE EXTERNAL TABLE dt STORED AS DELTATABLE LOCATION 's3://bucket/schema/table';";
+    ctx.sql(sql).await.unwrap();
+
+    let cat = ctx.catalog("datafusion").unwrap();
+    let schema = cat.schema("public").unwrap();
+    let exists = schema.table_exist("dt");
+    assert!(exists, "Table should have been created!");
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn create_bad_custom_table() {
+    let ctx = SessionContext::new();
+
+    let sql = "CREATE EXTERNAL TABLE dt STORED AS DELTATABLE LOCATION 's3://bucket/schema/table';";
+    let res = ctx.sql(sql).await;
+    match res {
+        Ok(_) => panic!("Registration of tables without factories should fail"),
+        Err(e) => {
+            assert!(e.to_string().contains("Unable to find factory for"), "Registration of tables without factories should throw correct error")
+            }
+        }
+    }
+}
+
 #[tokio::test]
 async fn create_csv_table_empty_file() -> Result<()> {
     let ctx =

datafusion/core/tests/sql/timestamp.rs

@@ -1398,6 +1398,7 @@ async fn timestamp_sub_interval_days() -> Result<()> {
 }
 
 #[tokio::test]
+#[ignore] // https://github.com/apache/arrow-datafusion/issues/3327
 async fn timestamp_add_interval_months() -> Result<()> {
     let ctx = SessionContext::new();
 

datafusion/expr/src/logical_plan/mod.rs

@@ -24,10 +24,9 @@ pub use builder::{table_scan, LogicalPlanBuilder};
 pub use plan::{
     Aggregate, Analyze, CreateCatalog, CreateCatalogSchema, CreateExternalTable,
     CreateMemoryTable, CreateView, CrossJoin, Distinct, DropTable, DropView,
-    EmptyRelation, Explain, Extension, FileType, Filter, Join, JoinConstraint, JoinType,
-    Limit, LogicalPlan, Partitioning, PlanType, PlanVisitor, Projection, Repartition,
-    Sort, StringifiedPlan, Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union,
-    Values, Window,
+    EmptyRelation, Explain, Extension, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan,
+    Partitioning, PlanType, PlanVisitor, Projection, Repartition, Sort, StringifiedPlan,
+    Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union, Values, Window,
 };
 
 pub use display::display_schema;

datafusion/expr/src/logical_plan/plan.rs

@@ -1236,19 +1236,6 @@ pub struct CreateView {
     pub definition: Option<String>,
 }
 
-/// Types of files to parse as DataFrames
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum FileType {
-    /// Newline-delimited JSON
-    NdJson,
-    /// Apache Parquet columnar storage
-    Parquet,
-    /// Comma separated values
-    CSV,
-    /// Avro binary records
-    Avro,
-}
-
 /// Creates an external table.
 #[derive(Clone)]
 pub struct CreateExternalTable {
@@ -1259,7 +1246,7 @@ pub struct CreateExternalTable {
     /// The physical location
     pub location: String,
     /// The file type of physical file
-    pub file_type: FileType,
+    pub file_type: String,
     /// Whether the CSV file contains a header
     pub has_header: bool,
     /// Delimiter for CSV

datafusion/proto/build.rs

@@ -30,8 +30,10 @@ fn main() -> Result<(), String> {
 
 #[cfg(feature = "json")]
 fn build() -> Result<(), String> {
-    let descriptor_path = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap())
-        .join("proto_descriptor.bin");
+    use std::io::Write;
+
+    let out = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
+    let descriptor_path = out.join("proto_descriptor.bin");
 
     prost_build::Config::new()
         .file_descriptor_set_path(&descriptor_path)
@@ -47,12 +49,24 @@ fn build() -> Result<(), String> {
         .build(&[".datafusion"])
         .map_err(|e| format!("pbjson compilation failed: {}", e))?;
 
+    // .serde.rs is not a valid package name, so append to datafusion.rs so we can treat it normally
+    let proto = std::fs::read_to_string(out.join("datafusion.rs")).unwrap();
+    let json = std::fs::read_to_string(out.join("datafusion.serde.rs")).unwrap();
+    let mut file = std::fs::OpenOptions::new()
+        .write(true)
+        .create(true)
+        .open("src/generated/datafusion.rs")
+        .unwrap();
+    file.write(proto.as_str().as_ref()).unwrap();
+    file.write(json.as_str().as_ref()).unwrap();
+
     Ok(())
 }
 
 #[cfg(not(feature = "json"))]
 fn build() -> Result<(), String> {
     prost_build::Config::new()
+        .out_dir("src/generated")
         .compile_protos(&["proto/datafusion.proto"], &["proto"])
         .map_err(|e| format!("protobuf compilation failed: {}", e))
 }