registry

Author: yjshen

datafusion/src/datasource/datasource2.rs

@@ -43,28 +43,23 @@ pub struct FilePartition {
 }
 
 #[derive(Debug, Clone)]
-pub struct ParquetSourceDesc {
-    partition_files: Vec<PartitionedFile>,
-    schema: SchemaRef,
+pub struct SourceDescriptor {
+    pub partition_files: Vec<PartitionedFile>,
+    pub schema: SchemaRef,
 }
 
-pub trait DataSource2<R>: Send + Sync
-where
-    R: ChunkReader + 'static,
-{
+pub trait DataSource2<R: ChunkReader + 'static>: Send + Sync {
     fn list_partitions(&self, max_concurrency: usize) -> Result<Arc<FilePartition>>;
 
     fn schema(&self) -> Result<Arc<Schema>>;
 
     fn get_read_for_file(&self, partitioned_file: PartitionedFile) -> Result<R>;
 
     fn statistics(&self) -> &Statistics;
+}
 
-    fn new_chunk_reader(file_path: &str) -> Result<R>;
-
-    fn get_all_files(root_path: &str) -> Result<Vec<String>>;
-
-    fn get_source_desc(root_path: &str) -> Result<ParquetSourceDesc> {
+pub trait SourceDescBuilder<R: ChunkReader + 'static> {
+    fn get_source_desc(root_path: &str) -> Result<SourceDescriptor> {
         let filenames = Self::get_all_files(root_path)?;
         if filenames.is_empty() {
             return Err(DataFusionError::Plan(format!(
@@ -90,22 +85,31 @@ where
                     // schema merging and this is a limitation.
                     // See https://issues.apache.org/jira/browse/ARROW-11017
                     return Err(DataFusionError::Plan(format!(
-                        "The Parquet file {} have different schema from the first file and DataFusion does \
+                        "The file {} have different schema from the first file and DataFusion does \
                         not yet support schema merging",
                         file_path
                     )));
                 }
                 Ok(pf)
             }).collect::<Result<Vec<PartitionedFile>>>();
 
-        Ok(ParquetSourceDesc {
+        Ok(SourceDescriptor {
             partition_files: partitioned_files?,
             schema: Arc::new(schemas.pop().unwrap()),
         })
     }
 
-    fn get_file_meta(file_path: &str) -> Result<PartitionedFile> {
-        let chunk_reader = Self::new_chunk_reader(file_path)?;
+    fn get_all_files(root_path: &str) -> Result<Vec<String>>;
+
+    fn get_file_meta(file_path: &str) -> Result<PartitionedFile>;
+
+    fn reader_for_file_meta(file_path: &str) -> Result<R>;
+}
+
+pub trait ParquetSourceDescBuilder: SourceDescBuilder {
+
+    fn get_file_meta(file_path: &str) {
+        let chunk_reader = Self::reader_for_file_meta(file_path)?;
         let file_reader = Arc::new(SerializedFileReader::new(chunk_reader)?);
         let arrow_reader = ParquetFileArrowReader::new(file_reader);
         let file_path = file_path.to_string();

datafusion/src/datasource/local/mod.rs

@@ -19,9 +19,16 @@ pub mod parquet_source;
 
 use crate::error::DataFusionError;
 use crate::error::Result;
+use super::datasource2::DataSource2;
 use std::fs;
 use std::fs::metadata;
 
+struct LocalFSHander {}
+
+impl DataSource2 for LocalFSHander {
+    
+}
+
 /// Recursively build a list of files in a directory with a given extension with an accumulator list
 fn list_all_files(dir: &str, filenames: &mut Vec<String>, ext: &str) -> Result<()> {
     let metadata = metadata(dir)?;

datafusion/src/datasource/local/parquet_source.rs

@@ -15,13 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::datasource::datasource2::DataSource2;
-use crate::datasource::datasource2::ParquetSourceDesc;
+use crate::datasource::datasource2::SourceDescBuilder;
+use crate::datasource::datasource2::SourceDescriptor;
 use crate::error::Result;
 use std::fs::File;
 
 struct LocalParquetSource {
-    desc: ParquetSourceDesc,
+    desc: SourceDescriptor,
 }
 
 impl LocalParquetSource {
@@ -32,8 +32,8 @@ impl LocalParquetSource {
     }
 }
 
-impl DataSource2<File> for LocalParquetSource {
-    fn new_chunk_reader(file_path: &str) -> Result<File> {
+impl SourceDescBuilder<File> for LocalParquetSource {
+    fn reader_for_file_meta(file_path: &str) -> Result<File> {
         Ok(File::open(file_path)?)
     }
 

datafusion/src/datasource/mod.rs

@@ -25,6 +25,7 @@ pub mod json;
 pub mod memory;
 pub mod parquet;
 pub mod local;
+pub mod protocol_registry;
 
 pub use self::csv::{CsvFile, CsvReadOptions};
 pub use self::datasource::{TableProvider, TableType};

datafusion/src/datasource/parquet.rs

@@ -47,7 +47,7 @@ impl<R: ChunkReader + 'static> ParquetTable<R> {
         max_concurrency: usize,
     ) -> Result<Self> {
         Ok(Self {
-            source,
+            source: Box::new(source),
             max_concurrency,
             enable_pruning: true,
         })

datafusion/src/datasource/protocol_registry.rs

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+use super::datasource2::DataSource2;
+
+pub trait ProtocolHander<R: ChunkReader + 'static>: Sync + Send {
+    /// Returns the protocol handler as [`Any`](std::any::Any)
+    /// so that it can be downcast to a specific implementation.
+    fn as_any(&self) -> &dyn Any;
+
+    fn source(&self,
+    
+    ) -> Result<Arc<dyn DataSource2>>;
+
+    fn list_all_files(&self, root_path: &str) -> Result<Vec<String>>;
+
+    fn get_reader(&self, file_path: &str) -> Result<R>;
+}
+
+struct LocalFSHander {
+
+}
+
+impl ProtocolHander for LocalFSHander {
+
+}
+
+pub struct ProtocolRegistry {
+    pub protocol_handlers: RwLock<HashMap<String, Arc<dyn ProtocolHander>>>,
+}
+
+impl ProtocolRegistry {
+    pub fn new() -> Self {
+        Self {
+            protocol_handlers: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Adds a new handler to this registry.
+    /// If a handler of the same prefix existed before, it is replaced in the registry and returned.
+    pub fn register_handler(
+        &self,
+        prefix: &str,
+        handler: Arc<dyn ProtocolHander>,
+    ) -> Option<Arc<dyn ProtocolHander>>  {
+        let mut handler = self.protocol_handlers.write().unwrap();
+        handler.insert(prefix.to_string(), handler)
+    }
+
+    pub fn handler(&self, prefix: &str) -> Option<Arc<dyn ProtocolHander>> {
+        let handler = self.protocol_handlers.read().unwrap();
+        handler.get(prefix).cloned()
+    }
+}

datafusion/src/execution/context.rs

@@ -156,6 +156,9 @@ impl ExecutionContext {
                 .register_catalog(config.default_catalog.clone(), default_catalog);
         }
 
+        let protocol_registry = ProtocolRegistry::new();
+        protocol_registry.register_handler("file", )
+
         Self {
             state: Arc::new(Mutex::new(ExecutionContextState {
                 catalog_list,
@@ -164,6 +167,7 @@ impl ExecutionContext {
                 aggregate_functions: HashMap::new(),
                 config,
                 execution_props: ExecutionProps::new(),
+                protocol_registry,
             })),
         }
     }
@@ -358,6 +362,20 @@ impl ExecutionContext {
         state.catalog_list.register_catalog(name, catalog)
     }
 
+    pub fn register_protocol_handler(
+        &self,
+        prefix: &str,
+        handler: Arc<dyn ProtocolHander>,
+    ) -> Option<Arc<dyn ProtocolHander>> {
+        let prefix = prefix.to_string();
+
+        self.state
+            .lock()
+            .unwrap()
+            .protocol_registry
+            .register_handler(prefix, handler)
+    }
+
     /// Retrieves a `CatalogProvider` instance by name
     pub fn catalog(&self, name: &str) -> Option<Arc<dyn CatalogProvider>> {
         self.state.lock().unwrap().catalog_list.catalog(name)
@@ -840,6 +858,8 @@ pub struct ExecutionContextState {
     pub config: ExecutionConfig,
     /// Execution properties
     pub execution_props: ExecutionProps,
+    /// Protocol handlers
+    pub protocol_registry: ProtocolRegistry, 
 }
 
 impl ExecutionProps {
@@ -867,6 +887,7 @@ impl ExecutionContextState {
             aggregate_functions: HashMap::new(),
             config: ExecutionConfig::new(),
             execution_props: ExecutionProps::new(),
+            protocol_registry: ProtocolRegistry::new(),
         }
     }