[WIP] Arrow (file) datasource

ballista/rust/client/src/context.rs

@@ -17,6 +17,7 @@
 
 //! Distributed execution context.
 
+use datafusion::execution::options::ArrowReadOptions;
 use parking_lot::Mutex;
 use sqlparser::ast::Statement;
 use std::collections::HashMap;
@@ -186,6 +187,30 @@ impl BallistaContext {
         Ok(df)
     }
 
+    /// Create a DataFrame representing a Arrow table scan
+    /// TODO fetch schema from scheduler instead of resolving locally
+    pub async fn read_arrow(
+        &self,
+        path: &str,
+        options: ArrowReadOptions<'_>,
+    ) -> Result<Arc<dyn DataFrame>> {
+        // convert to absolute path because the executor likely has a different working directory
+        let path = PathBuf::from(path);
+        let path = fs::canonicalize(&path)?;
+
+        // use local DataFusion context for now but later this might call the scheduler
+        let mut ctx = {
+            let guard = self.state.lock();
+            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
+                &guard.scheduler_host,
+                guard.scheduler_port,
+                guard.config(),
+            )
+        };
+        let df = ctx.read_arrow(path.to_str().unwrap(), options).await?;
+        Ok(df)
+    }
+
     /// Create a DataFrame representing a CSV table scan
     /// TODO fetch schema from scheduler instead of resolving locally
     pub async fn read_csv(
@@ -244,6 +269,20 @@ impl BallistaContext {
         }
     }
 
+    pub async fn register_arrow(
+        &self,
+        name: &str,
+        path: &str,
+        options: ArrowReadOptions<'_>,
+    ) -> Result<()> {
+        match self.read_arrow(path, options).await?.to_logical_plan() {
+            LogicalPlan::TableScan(TableScan { source, .. }) => {
+                self.register_table(name, source)
+            }
+            _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())),
+        }
+    }
+
     pub async fn register_avro(
         &self,
         name: &str,

ballista/rust/core/proto/ballista.proto

@@ -376,6 +376,7 @@ enum FileType {
   Parquet = 1;
   CSV = 2;
   Avro = 3;
+  Arrow = 4;
 }
 
 message AnalyzeNode {

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -1006,6 +1006,7 @@ impl Into<datafusion::sql::parser::FileType> for protobuf::FileType {
             protobuf::FileType::Parquet => FileType::Parquet,
             protobuf::FileType::Csv => FileType::CSV,
             protobuf::FileType::Avro => FileType::Avro,
+            protobuf::FileType::Arrow => FileType::Arrow,
         }
     }
 }

ballista/rust/core/src/serde/logical_plan/mod.rs

@@ -739,6 +739,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                     FileType::Parquet => protobuf::FileType::Parquet,
                     FileType::CSV => protobuf::FileType::Csv,
                     FileType::Avro => protobuf::FileType::Avro,
+                    FileType::Arrow => protobuf::FileType::Arrow,
                 };
 
                 Ok(protobuf::LogicalPlanNode {

datafusion/src/datasource/file_format/arrow_file.rs

@@ -0,0 +1,116 @@
+// 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.
+
+//! Apache Arrow format abstractions
+
+use std::any::Any;
+use std::io::{Read, Seek};
+use std::sync::Arc;
+
+use arrow::datatypes::Schema;
+use arrow::ipc::reader::FileReader;
+use arrow::{self, datatypes::SchemaRef};
+use async_trait::async_trait;
+use futures::StreamExt;
+
+use super::FileFormat;
+use crate::datasource::object_store::{ObjectReader, ObjectReaderStream};
+use crate::error::Result;
+use crate::logical_plan::Expr;
+use crate::physical_plan::file_format::{ArrowExec, FileScanConfig};
+use crate::physical_plan::ExecutionPlan;
+use crate::physical_plan::Statistics;
+
+/// The default file extension of arrow files
+pub const DEFAULT_ARROW_EXTENSION: &str = ".arrow";
+/// Arrow `FileFormat` implementation.
+#[derive(Default, Debug)]
+pub struct ArrowFormat;
+
+#[async_trait]
+impl FileFormat for ArrowFormat {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    async fn infer_schema(&self, mut readers: ObjectReaderStream) -> Result<SchemaRef> {
+        let mut schemas = vec![];
+        while let Some(obj_reader) = readers.next().await {
+            let mut reader = obj_reader?.sync_reader()?;
+            let schema = read_arrow_schema_from_reader(&mut reader)?;
+            schemas.push(schema.as_ref().clone());
+        }
+        let merged_schema = Schema::try_merge(schemas)?;
+        Ok(Arc::new(merged_schema))
+    }
+
+    async fn infer_stats(&self, _reader: Arc<dyn ObjectReader>) -> Result<Statistics> {
+        Ok(Statistics::default())
+    }
+
+    async fn create_physical_plan(
+        &self,
+        conf: FileScanConfig,
+        _filters: &[Expr],
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let exec = ArrowExec::new(conf);
+        Ok(Arc::new(exec))
+    }
+}
+
+fn read_arrow_schema_from_reader<R: Read + Seek>(reader: R) -> Result<SchemaRef> {
+    let reader = FileReader::try_new(reader)?;
+    Ok(reader.schema())
+}
+
+#[cfg(test)]
+mod tests {
+    use arrow::datatypes::DataType;
+
+    use crate::datasource::{
+        file_format::FileFormat, object_store::local::local_object_reader_stream,
+    };
+
+    use super::ArrowFormat;
+
+    #[tokio::test]
+    async fn test_schema() {
+        let filename = "tests/example.arrow";
+        let format = ArrowFormat {};
+        let file_schema = format
+            .infer_schema(local_object_reader_stream(vec![filename.to_owned()]))
+            .await
+            .expect("Schema inference");
+        assert_eq!(
+            vec!["f0", "f1", "f2"],
+            file_schema
+                .fields()
+                .iter()
+                .map(|x| x.name().clone())
+                .collect::<Vec<String>>()
+        );
+
+        assert_eq!(
+            vec![DataType::Int64, DataType::Utf8, DataType::Boolean],
+            file_schema
+                .fields()
+                .iter()
+                .map(|x| x.data_type().clone())
+                .collect::<Vec<DataType>>()
+        );
+    }
+}

datafusion/src/datasource/file_format/mod.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 //! Module containing helper methods for the various file formats
-
+pub mod arrow_file;
 pub mod avro;
 pub mod csv;
 pub mod json;

datafusion/src/datasource/file_format/parquet.rs

@@ -18,7 +18,6 @@
 //! Parquet format abstractions
 
 use std::any::Any;
-use std::io::Read;
 use std::sync::Arc;
 
 use arrow::datatypes::Schema;
@@ -40,6 +39,7 @@ use crate::arrow::array::{
     BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array,
 };
 use crate::arrow::datatypes::{DataType, Field};
+use crate::datasource::object_store::ReadSeek;
 use crate::datasource::object_store::{ObjectReader, ObjectReaderStream};
 use crate::datasource::{create_max_min_accs, get_col_stats};
 use crate::error::DataFusionError;
@@ -346,7 +346,7 @@ impl Length for ChunkObjectReader {
 }
 
 impl ChunkReader for ChunkObjectReader {
-    type T = Box<dyn Read + Send + Sync>;
+    type T = Box<dyn ReadSeek + Send + Sync>;
 
     fn get_read(&self, start: u64, length: usize) -> ParquetResult<Self::T> {
         self.0

datafusion/src/datasource/listing/table.rs

@@ -27,7 +27,7 @@ use crate::{
     datasource::file_format::avro::AvroFormat,
     datasource::file_format::csv::CsvFormat,
     datasource::file_format::json::JsonFormat,
-    datasource::file_format::parquet::ParquetFormat,
+    datasource::file_format::{arrow_file::ArrowFormat, parquet::ParquetFormat},
     error::{DataFusionError, Result},
     logical_plan::Expr,
     physical_plan::{
@@ -95,6 +95,7 @@ impl ListingTableConfig {
             "csv" => Ok(Arc::new(CsvFormat::default())),
             "json" => Ok(Arc::new(JsonFormat::default())),
             "parquet" => Ok(Arc::new(ParquetFormat::default())),
+            "arrow" => Ok(Arc::new(ArrowFormat::default())),
             _ => Err(DataFusionError::Internal(format!(
                 "Unable to infer file type from suffix {}",
                 suffix

datafusion/src/datasource/object_store/local.rs

@@ -18,7 +18,7 @@
 //! Object store that represents the Local File System.
 
 use std::fs::{self, File, Metadata};
-use std::io::{BufReader, Read, Seek, SeekFrom};
+use std::io::{BufReader, Seek, SeekFrom};
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -30,7 +30,7 @@ use crate::datasource::object_store::{
 use crate::datasource::PartitionedFile;
 use crate::error::{DataFusionError, Result};
 
-use super::{ObjectReaderStream, SizedFile};
+use super::{ObjectReaderStream, ReadSeek, SizedFile};
 
 #[derive(Debug)]
 /// Local File System as Object Store.
@@ -85,15 +85,14 @@ impl ObjectReader for LocalFileReader {
     fn sync_chunk_reader(
         &self,
         start: u64,
-        length: usize,
-    ) -> Result<Box<dyn Read + Send + Sync>> {
+        _length: usize,
+    ) -> Result<Box<dyn ReadSeek + Send + Sync>> {
         // A new file descriptor is opened for each chunk reader.
         // This okay because chunks are usually fairly large.
         let mut file = File::open(&self.file.path)?;
         file.seek(SeekFrom::Start(start))?;
 
-        let file = BufReader::new(file.take(length as u64));
-
+        let file = BufReader::new(file);
         Ok(Box::new(file))
     }
 

datafusion/src/datasource/object_store/mod.rs

@@ -22,7 +22,7 @@ pub mod local;
 use parking_lot::RwLock;
 use std::collections::HashMap;
 use std::fmt::{self, Debug};
-use std::io::Read;
+use std::io::{BufReader, Read, Seek};
 use std::pin::Pin;
 use std::sync::Arc;
 
@@ -34,6 +34,10 @@ use local::LocalFileSystem;
 
 use crate::error::{DataFusionError, Result};
 
+/// Trait Combining Read and Seek
+pub trait ReadSeek: Read + Seek {}
+impl<R: Seek + Read> ReadSeek for BufReader<R> {}
+
 /// Object Reader for one file in an object store.
 ///
 /// Note that the dynamic dispatch on the reader might
@@ -49,10 +53,10 @@ pub trait ObjectReader: Send + Sync {
         &self,
         start: u64,
         length: usize,
-    ) -> Result<Box<dyn Read + Send + Sync>>;
+    ) -> Result<Box<dyn ReadSeek + Send + Sync>>;
 
     /// Get reader for the entire file
-    fn sync_reader(&self) -> Result<Box<dyn Read + Send + Sync>> {
+    fn sync_reader(&self) -> Result<Box<dyn ReadSeek + Send + Sync>> {
         self.sync_chunk_reader(0, self.length() as usize)
     }