Parquet datasource

Author: andygrove

rust/datafusion/src/datasource/mod.rs

@@ -18,6 +18,7 @@
 pub mod csv;
 pub mod datasource;
 pub mod memory;
+pub mod parquet;
 
 pub use self::csv::{CsvBatchIterator, CsvFile};
 pub use self::datasource::{RecordBatchIterator, ScanResult, Table};

rust/datafusion/src/datasource/parquet.rs

@@ -0,0 +1,310 @@
+// 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.
+
+//! Parquet Data source
+
+use std::cell::RefCell;
+use std::fs::File;
+use std::rc::Rc;
+use std::string::String;
+use std::sync::Arc;
+
+use arrow::array::Array;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+
+use parquet::basic;
+use parquet::column::reader::*;
+use parquet::data_type::ByteArray;
+use parquet::file::reader::*;
+use parquet::schema::types::Type;
+
+use crate::datasource::{RecordBatchIterator, Table};
+use crate::execution::error::{ExecutionError, Result};
+use arrow::builder::{BinaryBuilder, Float64Builder, Int32Builder};
+
+pub struct ParquetTable {
+    filename: String,
+    schema: Arc<Schema>,
+}
+
+impl ParquetTable {
+    pub fn new(filename: &str) -> Self {
+        let file = File::open(filename).unwrap();
+        let parquet_file = ParquetFile::open(file, None).unwrap();
+        let schema = parquet_file.schema.clone();
+        Self {
+            filename: filename.to_string(),
+            schema,
+        }
+    }
+}
+
+impl Table for ParquetTable {
+    fn schema(&self) -> &Arc<Schema> {
+        &self.schema
+    }
+
+    fn scan(
+        &self,
+        projection: &Option<Vec<usize>>,
+        _batch_size: usize,
+    ) -> Result<Rc<RefCell<RecordBatchIterator>>> {
+        let file = File::open(self.filename.clone()).unwrap();
+        let parquet_file = ParquetFile::open(file, projection.clone()).unwrap();
+        Ok(Rc::new(RefCell::new(parquet_file)))
+    }
+}
+
+pub struct ParquetFile {
+    reader: SerializedFileReader<File>,
+    row_group_index: usize,
+    schema: Arc<Schema>,
+    projection: Option<Vec<usize>>,
+    batch_size: usize,
+    current_row_group: Option<Box<RowGroupReader>>,
+    column_readers: Vec<ColumnReader>,
+}
+
+impl ParquetFile {
+    pub fn open(file: File, projection: Option<Vec<usize>>) -> Result<Self> {
+        let reader = SerializedFileReader::new(file).unwrap();
+
+        let metadata = reader.metadata();
+        let file_type = to_arrow(metadata.file_metadata().schema())?;
+
+        match file_type.data_type() {
+            DataType::Struct(fields) => {
+                let schema = Schema::new(fields.clone());
+                //println!("Parquet schema: {:?}", schema);
+                Ok(ParquetFile {
+                    reader: reader,
+                    row_group_index: 0,
+                    schema: Arc::new(schema),
+                    projection,
+                    batch_size: 64 * 1024,
+                    current_row_group: None,
+                    column_readers: vec![],
+                })
+            }
+            _ => Err(ExecutionError::General(
+                "Failed to read Parquet schema".to_string(),
+            )),
+        }
+    }
+
+    fn load_next_row_group(&mut self) {
+        if self.row_group_index < self.reader.num_row_groups() {
+            //println!("Loading row group {} of {}", self.row_group_index, self.reader.num_row_groups());
+            let reader = self.reader.get_row_group(self.row_group_index).unwrap();
+
+            self.column_readers = vec![];
+
+            match &self.projection {
+                None => {
+                    for i in 0..reader.num_columns() {
+                        self.column_readers
+                            .push(reader.get_column_reader(i).unwrap());
+                    }
+                }
+                Some(proj) => {
+                    for i in proj {
+                        //TODO validate index in bounds
+                        self.column_readers
+                            .push(reader.get_column_reader(*i).unwrap());
+                    }
+                }
+            }
+
+            self.current_row_group = Some(reader);
+            self.row_group_index += 1;
+        } else {
+            panic!()
+        }
+    }
+
+    fn load_batch(&mut self) -> Result<Option<RecordBatch>> {
+        match &self.current_row_group {
+            Some(reader) => {
+                let mut batch: Vec<Arc<Array>> = Vec::with_capacity(reader.num_columns());
+                let mut row_count = 0;
+                for i in 0..self.column_readers.len() {
+                    let array: Arc<Array> = match self.column_readers[i] {
+                        ColumnReader::Int32ColumnReader(ref mut r) => {
+                            let mut builder = Int32Builder::new(self.batch_size);
+                            let mut read_buffer: Vec<i32> =
+                                Vec::with_capacity(self.batch_size);
+                            match r.read_batch(
+                                self.batch_size,
+                                None,
+                                None,
+                                &mut read_buffer,
+                            ) {
+                                //TODO this isn't handling null values
+                                Ok((count, _)) => {
+                                    builder.append_slice(&read_buffer).unwrap();
+                                    row_count = count;
+                                    Arc::new(builder.finish())
+                                }
+                                _ => {
+                                    return Err(ExecutionError::NotImplemented(format!(
+                                        "Error reading parquet batch (column {})",
+                                        i
+                                    )));
+                                }
+                            }
+                        }
+                        ColumnReader::DoubleColumnReader(ref mut r) => {
+                            let mut builder = Float64Builder::new(self.batch_size);
+                            let mut read_buffer: Vec<f64> =
+                                Vec::with_capacity(self.batch_size);
+                            match r.read_batch(
+                                self.batch_size,
+                                None,
+                                None,
+                                &mut read_buffer,
+                            ) {
+                                //TODO this isn't handling null values
+                                Ok((count, _)) => {
+                                    builder.append_slice(&read_buffer).unwrap();
+                                    row_count = count;
+                                    Arc::new(builder.finish())
+                                }
+                                _ => {
+                                    return Err(ExecutionError::NotImplemented(format!(
+                                        "Error reading parquet batch (column {})",
+                                        i
+                                    )));
+                                }
+                            }
+                        }
+                        ColumnReader::ByteArrayColumnReader(ref mut r) => {
+                            let mut b: Vec<ByteArray> =
+                                Vec::with_capacity(self.batch_size);
+                            for _ in 0..self.batch_size {
+                                b.push(ByteArray::default());
+                            }
+                            match r.read_batch(self.batch_size, None, None, &mut b) {
+                                //TODO this isn't handling null values
+                                Ok((count, _)) => {
+                                    row_count = count;
+                                    //TODO this is horribly inefficient
+                                    let mut builder = BinaryBuilder::new(row_count);
+                                    for j in 0..row_count {
+                                        let foo = b[j].slice(0, b[j].len());
+                                        let bytes: &[u8] = foo.data();
+                                        let str =
+                                            String::from_utf8(bytes.to_vec()).unwrap();
+                                        builder.append_string(&str).unwrap();
+                                    }
+                                    Arc::new(builder.finish())
+                                }
+                                _ => {
+                                    return Err(ExecutionError::NotImplemented(format!(
+                                        "Error reading parquet batch (column {})",
+                                        i
+                                    )));
+                                }
+                            }
+                        }
+                        _ => {
+                            return Err(ExecutionError::NotImplemented(
+                                "unsupported column reader type".to_string(),
+                            ));
+                        }
+                    };
+
+                    batch.push(array);
+                }
+
+                //                println!("Loaded batch of {} rows", row_count);
+
+                if row_count == 0 {
+                    Ok(None)
+                } else {
+                    Ok(Some(RecordBatch::try_new(self.schema.clone(), batch)?))
+                }
+            }
+            _ => Ok(None),
+        }
+    }
+}
+
+fn to_arrow(t: &Type) -> Result<Field> {
+    match t {
+        Type::PrimitiveType {
+            basic_info,
+            physical_type,
+            ..
+        } => {
+            let arrow_type = match physical_type {
+                basic::Type::BOOLEAN => DataType::Boolean,
+                basic::Type::INT32 => DataType::Int32,
+                basic::Type::INT64 => DataType::Int64,
+                basic::Type::INT96 => DataType::Int64, //TODO ???
+                basic::Type::FLOAT => DataType::Float32,
+                basic::Type::DOUBLE => DataType::Float64,
+                basic::Type::BYTE_ARRAY => DataType::Utf8, /*match basic_info.logical_type() {
+                basic::LogicalType::UTF8 => DataType::Utf8,
+                _ => unimplemented!("No support for Parquet BYTE_ARRAY yet"),
+                }*/
+                basic::Type::FIXED_LEN_BYTE_ARRAY => {
+                    unimplemented!("No support for Parquet FIXED_LEN_BYTE_ARRAY yet")
+                }
+            };
+
+            Ok(Field::new(basic_info.name(), arrow_type, false))
+        }
+        Type::GroupType { basic_info, fields } => Ok(Field::new(
+            basic_info.name(),
+            DataType::Struct(
+                fields
+                    .iter()
+                    .map(|f| to_arrow(f))
+                    .collect::<Result<Vec<Field>>>()?,
+            ),
+            false,
+        )),
+    }
+}
+
+impl RecordBatchIterator for ParquetFile {
+    fn schema(&self) -> &Arc<Schema> {
+        &self.schema
+    }
+
+    fn next(&mut self) -> Result<Option<RecordBatch>> {
+        // advance the row group reader if necessary
+        if self.current_row_group.is_none() {
+            self.load_next_row_group();
+            self.load_batch()
+        } else {
+            match self.load_batch() {
+                Ok(Some(b)) => Ok(Some(b)),
+                Ok(None) => {
+                    if self.row_group_index < self.reader.num_row_groups() {
+                        self.load_next_row_group();
+                        self.load_batch()
+                    } else {
+                        Ok(None)
+                    }
+                }
+                Err(e) => Err(e),
+            }
+        }
+    }
+}