Row-based read support

src/bin/dump-schema.rs

@@ -26,7 +26,7 @@ use parquet::file::reader::{FileReader, SerializedFileReader};
 use parquet::schema::printer::{print_parquet_metadata, print_file_metadata};
 
 fn main() {
-  let args: Vec<_> = env::args().collect();
+  let args: Vec<String> = env::args().collect();
   if args.len() != 2 && args.len() != 3 {
     println!("Usage: dump-schema <file-path> <verbose>");
     process::exit(1);

src/bin/read-file.rs

@@ -0,0 +1,43 @@
+extern crate parquet;
+
+use std::env;
+use std::fs::File;
+use std::path::Path;
+use std::process;
+
+use parquet::file::reader::{FileReader, SerializedFileReader};
+
+fn main() {
+  let args: Vec<String> = env::args().collect();
+  if args.len() != 2 && args.len() != 3 {
+    println!("Usage: read-file <file-path> <num-records>");
+    process::exit(1);
+  }
+
+  let mut num_records: Option<usize> = None;
+  if args.len() == 3 {
+    match args[2].parse() {
+      Ok(value) => num_records = Some(value),
+      Err(e) => panic!("Error when reading value for <num-records>, {}", e)
+    }
+  }
+
+  let path = Path::new(&args[1]);
+  let file = File::open(&path).unwrap();
+  let parquet_reader = SerializedFileReader::new(file).unwrap();
+
+  // Use full schema as projected schema
+  let mut iter = parquet_reader.get_row_iter(None).unwrap();
+
+  let mut start = 0;
+  let end = num_records.unwrap_or(0);
+  let all_records = num_records.is_none();
+
+  while all_records || start < end {
+    match iter.next() {
+      Some(row) => println!("{}", row),
+      None => break,
+    }
+    start += 1;
+  }
+}

src/file/metadata.rs

@@ -92,6 +92,10 @@ impl FileMetaData {
   pub fn schema_descr(&self) -> &SchemaDescriptor {
     &self.schema_descr
   }
+
+  pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
+    self.schema_descr.clone()
+  }
 }
 
 pub type RowGroupMetaDataPtr = Rc<RowGroupMetaData>;
@@ -192,6 +196,11 @@ impl ColumnChunkMetaData {
     self.column_descr.as_ref()
   }
 
+  /// Reference counted clone of descriptor for this column
+  pub fn column_descr_ptr(&self) -> ColumnDescPtr {
+    self.column_descr.clone()
+  }
+
   /// All encodings used for this column
   pub fn encodings(&self) -> &Vec<Encoding> {
     &self.encodings

src/file/reader.rs

@@ -26,10 +26,11 @@ use byteorder::{LittleEndian, ByteOrder};
 use thrift::protocol::TCompactInputProtocol;
 use parquet_thrift::parquet::FileMetaData as TFileMetaData;
 use parquet_thrift::parquet::{PageType, PageHeader};
-use schema::types::{self, SchemaDescriptor};
+use schema::types::{self, Type as SchemaType, SchemaDescriptor};
 use column::page::{Page, PageReader};
 use column::reader::{ColumnReader, ColumnReaderImpl};
 use compression::{Codec, create_codec};
+use record::reader::{FileRowIter, RowIter, TreeBuilder};
 use util::io::FileChunk;
 use util::memory::ByteBufferPtr;
 
@@ -50,6 +51,11 @@ pub trait FileReader {
   /// the same as this. Otherwise, the row group metadata stored in the row group reader
   /// may outlive the file reader.
   fn get_row_group(&self, i: usize) -> Result<Box<RowGroupReader>>;
+
+  /// Get full iterator of `Row` from a file (over all row groups).
+  /// Projected schema can be a subset of or equal to the file schema, when it is None,
+  /// full file schema is assumed.
+  fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<FileRowIter>;
 }
 
 /// Parquet row group reader API. With this, user can get metadata information about the
@@ -66,6 +72,11 @@ pub trait RowGroupReader {
 
   /// Get value reader for the `i`th column chunk
   fn get_column_reader(&self, i: usize) -> Result<ColumnReader>;
+
+  /// Get row iterator for this row group, using schema descriptor.
+  /// Schema descriptor must be a valid subset of the file schema or be an original
+  /// descriptor, see `FileReader::get_row_iter` for more information.
+  fn get_row_iter(&self, proj_descr: Rc<SchemaDescriptor>) -> RowIter;
 }
 
 
@@ -174,6 +185,27 @@ impl FileReader for SerializedFileReader {
     let f = self.buf.get_ref().try_clone()?;
     Ok(Box::new(SerializedRowGroupReader::new(f, row_group_metadata)))
   }
+
+  fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<FileRowIter> {
+    let file_metadata = self.metadata().file_metadata();
+
+    // Resolve projected schema (either full file schema or a subset)
+    let proj_descr = match projection {
+      Some(projection) => {
+        // check if projection is part of file schema
+        let root_schema = file_metadata.schema_descr().root_schema();
+        if !root_schema.check_contains(&projection) {
+          return Err(general_err!("Root schema does not contain projection"));
+        }
+        Rc::new(SchemaDescriptor::new(Rc::new(projection)))
+      },
+      None => {
+        file_metadata.schema_descr_ptr()
+      }
+    };
+
+    Ok(FileRowIter::new(proj_descr, self))
+  }
 }
 
 /// A serialized impl for row group reader
@@ -183,7 +215,7 @@ pub struct SerializedRowGroupReader {
 }
 
 impl SerializedRowGroupReader {
-  pub fn new(file: File, metadata: RowGroupMetaDataPtr ) -> Self {
+  pub fn new(file: File, metadata: RowGroupMetaDataPtr) -> Self {
     let buf = BufReader::new(file);
     Self { buf, metadata }
   }
@@ -237,6 +269,11 @@ impl RowGroupReader for SerializedRowGroupReader {
     };
     Ok(col_reader)
   }
+
+  fn get_row_iter(&self, proj_descr: Rc<SchemaDescriptor>) -> RowIter {
+    // Build new tree of readers and return it as a row iterator
+    TreeBuilder::new().as_row_iter(proj_descr, self)
+  }
 }
 
 
@@ -382,9 +419,7 @@ impl PageReader for SerializedPageReader {
 #[cfg(test)]
 mod tests {
   use super::*;
-  use std::fs;
-  use std::env;
-  use std::io::Write;
+  use util::test_common::{get_temp_file, get_test_file};
 
   #[test]
   fn test_file_reader_metadata_size_smaller_than_footer() {
@@ -577,33 +612,4 @@ mod tests {
     }
     assert_eq!(page_count, 2);
   }
-
-  fn get_test_file(file_name: &str) -> fs::File {
-    let mut path_buf = env::current_dir().unwrap();
-    path_buf.push("data");
-    path_buf.push(file_name);
-    let file = File::open(path_buf.as_path());
-    assert!(file.is_ok());
-    file.unwrap()
-  }
-
-  fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
-    // build tmp path to a file in "target/debug/testdata"
-    let mut path_buf = env::current_dir().unwrap();
-    path_buf.push("target");
-    path_buf.push("debug");
-    path_buf.push("testdata");
-    fs::create_dir_all(&path_buf).unwrap();
-    path_buf.push(file_name);
-
-    // write file content
-    let mut tmp_file = File::create(path_buf.as_path()).unwrap();
-    tmp_file.write_all(content).unwrap();
-    tmp_file.sync_all().unwrap();
-
-    // read file and return file handle
-    let file = File::open(path_buf.as_path());
-    assert!(file.is_ok());
-    file.unwrap()
-  }
 }

src/lib.rs

@@ -47,6 +47,7 @@ mod parquet_thrift;
 #[macro_use]
 pub mod util;
 pub mod column;
+pub mod record;
 pub mod compression;
 
 pub mod schema;