test

Author: andygrove

rust/datafusion/src/datasource/parquet.rs

@@ -144,10 +144,16 @@ impl ParquetFile {
                 let mut row_count = 0;
                 for i in 0..self.column_readers.len() {
                     let array: Arc<Array> = match self.column_readers[i] {
+                        ColumnReader::BoolColumnReader(ref mut r) => {
+                            return Err(ExecutionError::NotImplemented(
+                                "unsupported column reader type (BOOL)".to_string(),
+                            ));
+                        }
                         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,
@@ -156,6 +162,8 @@ impl ParquetFile {
                             ) {
                                 //TODO this isn't handling null values
                                 Ok((count, _)) => {
+                                    println!("Read {} rows", count);
+
                                     builder.append_slice(&read_buffer).unwrap();
                                     row_count = count;
                                     Arc::new(builder.finish())
@@ -168,6 +176,21 @@ impl ParquetFile {
                                 }
                             }
                         }
+                        ColumnReader::Int64ColumnReader(ref mut r) => {
+                            return Err(ExecutionError::NotImplemented(
+                                "unsupported column reader type (INT64)".to_string(),
+                            ));
+                        }
+                        ColumnReader::Int96ColumnReader(ref mut r) => {
+                            return Err(ExecutionError::NotImplemented(
+                                "unsupported column reader type (INT96)".to_string(),
+                            ));
+                        }
+                        ColumnReader::FloatColumnReader(ref mut r) => {
+                            return Err(ExecutionError::NotImplemented(
+                                "unsupported column reader type (FLOAT)".to_string(),
+                            ));
+                        }
                         ColumnReader::DoubleColumnReader(ref mut r) => {
                             let mut builder = Float64Builder::new(self.batch_size);
                             let mut read_buffer: Vec<f64> =
@@ -192,6 +215,12 @@ impl ParquetFile {
                                 }
                             }
                         }
+                        ColumnReader::FixedLenByteArrayColumnReader(ref mut r) => {
+                            return Err(ExecutionError::NotImplemented(
+                                "unsupported column reader type (FixedLenByteArray)"
+                                    .to_string(),
+                            ));
+                        }
                         ColumnReader::ByteArrayColumnReader(ref mut r) => {
                             let mut b: Vec<ByteArray> =
                                 Vec::with_capacity(self.batch_size);
@@ -221,17 +250,13 @@ impl ParquetFile {
                                 }
                             }
                         }
-                        _ => {
-                            return Err(ExecutionError::NotImplemented(
-                                "unsupported column reader type".to_string(),
-                            ));
-                        }
                     };
 
+                    println!("Adding array to batch");
                     batch.push(array);
                 }
 
-                //                println!("Loaded batch of {} rows", row_count);
+                println!("Loaded batch of {} rows", row_count);
 
                 if row_count == 0 {
                     Ok(None)
@@ -308,3 +333,27 @@ impl RecordBatchIterator for ParquetFile {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::env;
+
+    #[test]
+    fn read_parquet_file() {
+        let testdata = env::var("PARQUET_TEST_DATA").unwrap();
+        let filename = format!("{}/alltypes_plain.parquet", testdata);
+
+        let table = ParquetTable::new(&filename);
+
+        println!("{:?}", table.schema());
+
+        let projection = Some(vec![0]);
+        let scan = table.scan(&projection, 1024).unwrap();
+        let mut it = scan.borrow_mut();
+        let batch = it.next().unwrap().unwrap();
+
+        assert_eq!(1, batch.num_columns());
+        assert_eq!(1, batch.num_rows());
+    }
+}