Improved flexibility of reading parquet (#820)

examples/parquet_read_async.rs

@@ -39,7 +39,7 @@ async fn main() -> Result<()> {
         // A row group is consumed in two steps: the first step is to read the (compressed)
         // columns into memory, which is IO-bounded.
         let column_chunks =
-            read::read_columns_async(factory, row_group, schema.fields.clone(), None).await?;
+            read::read_columns_many_async(factory, row_group, schema.fields.clone(), None).await?;
 
         // the second step is to iterate over the columns in chunks.
         // this operation is CPU-bounded and should be sent to a separate thread pool (e.g. `tokio_rayon`) to not block

src/io/parquet/read/file.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use crate::array::Array;
 use crate::chunk::Chunk;
 use crate::datatypes::Schema;
-use crate::io::parquet::read::read_columns;
+use crate::io::parquet::read::read_columns_many;
 use crate::{
     datatypes::Field,
     error::{ArrowError, Result},
@@ -233,7 +233,7 @@ impl<R: Read + Seek> RowGroupReader<R> {
         }
         self.current_group += 1;
 
-        let column_chunks = read_columns(
+        let column_chunks = read_columns_many(
             &mut self.reader,
             row_group,
             self.schema.fields.clone(),

src/io/parquet/read/row_group.rs

@@ -100,24 +100,35 @@ pub(super) fn get_field_columns<'a>(
 }
 
 /// Reads all columns that are part of the parquet field `field_name`
-pub(super) fn _read_columns<'a, R: Read + Seek>(
+/// # Implementation
+/// This operation is IO-bounded `O(C)` where C is the number of columns associated to
+/// the field (one for non-nested types)
+pub fn read_columns<'a, R: Read + Seek>(
     reader: &mut R,
     columns: &'a [ColumnChunkMetaData],
     field_name: &str,
 ) -> Result<Vec<(&'a ColumnChunkMetaData, Vec<u8>)>> {
     get_field_columns(columns, field_name)
         .into_iter()
-        .map(|meta| {
-            let (start, len) = meta.byte_range();
-            reader.seek(std::io::SeekFrom::Start(start))?;
-            let mut chunk = vec![0; len as usize];
-            reader.read_exact(&mut chunk)?;
-            Ok((meta, chunk))
-        })
+        .map(|meta| _read_single_column(reader, meta))
         .collect()
 }
 
-async fn _read_single_column<'b, R, F>(
+fn _read_single_column<'a, R>(
+    reader: &mut R,
+    meta: &'a ColumnChunkMetaData,
+) -> Result<(&'a ColumnChunkMetaData, Vec<u8>)>
+where
+    R: Read + Seek,
+{
+    let (start, len) = meta.byte_range();
+    reader.seek(std::io::SeekFrom::Start(start))?;
+    let mut chunk = vec![0; len as usize];
+    reader.read_exact(&mut chunk)?;
+    Ok((meta, chunk))
+}
+
+async fn _read_single_column_async<'b, R, F>(
     factory: F,
     meta: &ColumnChunkMetaData,
 ) -> Result<(&ColumnChunkMetaData, Vec<u8>)>
@@ -134,7 +145,13 @@ where
 }
 
 /// Reads all columns that are part of the parquet field `field_name`
-async fn _read_columns_async<
+/// # Implementation
+/// This operation is IO-bounded `O(C)` where C is the number of columns associated to
+/// the field (one for non-nested types)
+///
+/// It does so asynchronously via a single `join_all` over all the necessary columns for
+/// `field_name`.
+pub async fn read_columns_async<
     'a,
     'b,
     R: AsyncRead + AsyncSeek + Send + Unpin,
@@ -146,54 +163,53 @@ async fn _read_columns_async<
 ) -> Result<Vec<(&'a ColumnChunkMetaData, Vec<u8>)>> {
     let futures = get_field_columns(columns, field_name)
         .into_iter()
-        .map(|meta| async { _read_single_column(factory.clone(), meta).await });
+        .map(|meta| async { _read_single_column_async(factory.clone(), meta).await });
 
     try_join_all(futures).await
 }
 
-fn to_deserializers<'a>(
-    field_columns: Vec<Vec<(&ColumnChunkMetaData, Vec<u8>)>>,
-    fields: Vec<Field>,
-    row_group: &RowGroupMetaData,
+/// Converts a vector of columns associated with the parquet field whose name is [`Field`]
+/// to an iterator of [`Array`], [`ArrayIter`] of chunk size `chunk_size`.
+pub fn to_deserializer<'a>(
+    columns: Vec<(&ColumnChunkMetaData, Vec<u8>)>,
+    field: Field,
+    num_rows: usize,
     chunk_size: Option<usize>,
-) -> Result<Vec<ArrayIter<'a>>> {
-    let chunk_size = chunk_size
-        .unwrap_or(usize::MAX)
-        .min(row_group.num_rows() as usize);
+) -> Result<ArrayIter<'a>> {
+    let chunk_size = chunk_size.unwrap_or(usize::MAX).min(num_rows);
 
-    field_columns
+    let (columns, types): (Vec<_>, Vec<_>) = columns
         .into_iter()
-        .map(|columns| {
-            let (columns, types): (Vec<_>, Vec<_>) = columns
-                .into_iter()
-                .map(|(column_meta, chunk)| {
-                    let pages = PageIterator::new(
-                        std::io::Cursor::new(chunk),
-                        column_meta.num_values(),
-                        column_meta.compression(),
-                        column_meta.descriptor().clone(),
-                        Arc::new(|_, _| true),
-                        vec![],
-                    );
-                    (
-                        BasicDecompressor::new(pages, vec![]),
-                        column_meta.descriptor().type_(),
-                    )
-                })
-                .unzip();
-            (columns, types)
+        .map(|(column_meta, chunk)| {
+            let pages = PageIterator::new(
+                std::io::Cursor::new(chunk),
+                column_meta.num_values(),
+                column_meta.compression(),
+                column_meta.descriptor().clone(),
+                Arc::new(|_, _| true),
+                vec![],
+            );
+            (
+                BasicDecompressor::new(pages, vec![]),
+                column_meta.descriptor().type_(),
+            )
         })
-        .zip(fields.into_iter())
-        .map(|((columns, types), field)| column_iter_to_arrays(columns, types, field, chunk_size))
-        .collect()
+        .unzip();
+
+    column_iter_to_arrays(columns, types, field, chunk_size)
 }
 
-/// Returns a vector of iterators of [`Array`] corresponding to the top level parquet fields whose
-/// name matches `fields`'s names.
+/// Returns a vector of iterators of [`Array`] ([`ArrayIter`]) corresponding to the top
+/// level parquet fields whose name matches `fields`'s names.
 ///
+/// # Implementation
 /// This operation is IO-bounded `O(C)` where C is the number of columns in the row group -
 /// it reads all the columns to memory from the row group associated to the requested fields.
-pub fn read_columns<'a, R: Read + Seek>(
+///
+/// This operation is single-threaded. For readers with stronger invariants
+/// (e.g. implement [`Clone`]) you can use [`read_columns`] to read multiple columns at once
+/// and convert them to [`ArrayIter`] via [`to_deserializer`].
+pub fn read_columns_many<'a, R: Read + Seek>(
     reader: &mut R,
     row_group: &RowGroupMetaData,
     fields: Vec<Field>,
@@ -203,10 +219,16 @@ pub fn read_columns<'a, R: Read + Seek>(
     // This operation is IO-bounded `O(C)` where C is the number of columns in the row group
     let field_columns = fields
         .iter()
-        .map(|field| _read_columns(reader, row_group.columns(), &field.name))
+        .map(|field| read_columns(reader, row_group.columns(), &field.name))
         .collect::<Result<Vec<_>>>()?;
 
-    to_deserializers(field_columns, fields, row_group, chunk_size)
+    field_columns
+        .into_iter()
+        .zip(fields.into_iter())
+        .map(|(columns, field)| {
+            to_deserializer(columns, field, row_group.num_rows() as usize, chunk_size)
+        })
+        .collect()
 }
 
 /// Returns a vector of iterators of [`Array`] corresponding to the top level parquet fields whose
@@ -216,8 +238,10 @@ pub fn read_columns<'a, R: Read + Seek>(
 /// it reads all the columns to memory from the row group associated to the requested fields.
 ///
 /// # Implementation
-/// This function concurrently reads all columns.
-pub async fn read_columns_async<
+/// This operation is IO-bounded `O(C)` where C is the number of columns in the row group -
+/// it reads all the columns to memory from the row group associated to the requested fields.
+/// It does so asynchronously via `join_all`
+pub async fn read_columns_many_async<
     'a,
     'b,
     R: AsyncRead + AsyncSeek + Send + Unpin,
@@ -230,9 +254,15 @@ pub async fn read_columns_async<
 ) -> Result<Vec<ArrayIter<'a>>> {
     let futures = fields
         .iter()
-        .map(|field| _read_columns_async(factory.clone(), row_group.columns(), &field.name));
+        .map(|field| read_columns_async(factory.clone(), row_group.columns(), &field.name));
 
     let field_columns = try_join_all(futures).await?;
 
-    to_deserializers(field_columns, fields, row_group, chunk_size)
+    field_columns
+        .into_iter()
+        .zip(fields.into_iter())
+        .map(|(columns, field)| {
+            to_deserializer(columns, field, row_group.num_rows() as usize, chunk_size)
+        })
+        .collect()
 }