Merge pull request #7 from alamb/alamb/test_memory_limit

Add option to control predicate cache, documentation, `ArrowReaderMetrics` and tests

Author: XiangpengHao

parquet/examples/read_with_rowgroup.rs

@@ -17,6 +17,7 @@
 
 use arrow::util::pretty::print_batches;
 use bytes::{Buf, Bytes};
+use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics;
 use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, RowGroups, RowSelection};
 use parquet::arrow::async_reader::AsyncFileReader;
 use parquet::arrow::{parquet_to_arrow_field_levels, ProjectionMask};
@@ -38,7 +39,11 @@ async fn main() -> Result<()> {
     let metadata = file.get_metadata(None).await?;
 
     for rg in metadata.row_groups() {
-        let mut rowgroup = InMemoryRowGroup::create(rg.clone(), ProjectionMask::all());
+        let mut rowgroup = InMemoryRowGroup::create(
+            rg.clone(),
+            ProjectionMask::all(),
+            ArrowReaderMetrics::disabled(),
+        );
         rowgroup.async_fetch_data(&mut file, None).await?;
         let reader = rowgroup.build_reader(1024, None)?;
 
@@ -103,6 +108,7 @@ pub struct InMemoryRowGroup {
     pub metadata: RowGroupMetaData,
     mask: ProjectionMask,
     column_chunks: Vec<Option<Arc<ColumnChunkData>>>,
+    metrics: ArrowReaderMetrics,
 }
 
 impl RowGroups for InMemoryRowGroup {
@@ -132,13 +138,18 @@ impl RowGroups for InMemoryRowGroup {
 }
 
 impl InMemoryRowGroup {
-    pub fn create(metadata: RowGroupMetaData, mask: ProjectionMask) -> Self {
+    pub fn create(
+        metadata: RowGroupMetaData,
+        mask: ProjectionMask,
+        metrics: ArrowReaderMetrics,
+    ) -> Self {
         let column_chunks = metadata.columns().iter().map(|_| None).collect::<Vec<_>>();
 
         Self {
             metadata,
             mask,
             column_chunks,
+            metrics,
         }
     }
 
@@ -153,7 +164,13 @@ impl InMemoryRowGroup {
             None,
         )?;
 
-        ParquetRecordBatchReader::try_new_with_row_groups(&levels, self, batch_size, selection)
+        ParquetRecordBatchReader::try_new_with_row_groups(
+            &levels,
+            self,
+            batch_size,
+            selection,
+            &self.metrics,
+        )
     }
 
     /// fetch data from a reader in sync mode

parquet/src/arrow/array_reader/builder.rs

@@ -30,21 +30,14 @@ use crate::arrow::array_reader::{
     FixedSizeListArrayReader, ListArrayReader, MapArrayReader, NullArrayReader,
     PrimitiveArrayReader, RowGroups, StructArrayReader,
 };
+use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics;
 use crate::arrow::schema::{ParquetField, ParquetFieldType};
 use crate::arrow::ProjectionMask;
 use crate::basic::Type as PhysicalType;
 use crate::data_type::{BoolType, DoubleType, FloatType, Int32Type, Int64Type, Int96Type};
 use crate::errors::{ParquetError, Result};
 use crate::schema::types::{ColumnDescriptor, ColumnPath, Type};
 
-/// Builds [`ArrayReader`]s from parquet schema, projection mask, and RowGroups reader
-pub struct ArrayReaderBuilder<'a> {
-    /// Source of row group data
-    row_groups: &'a dyn RowGroups,
-    /// Optional cache options for the array reader
-    cache_options: Option<&'a CacheOptions<'a>>,
-}
-
 /// Builder for [`CacheOptions`]
 #[derive(Debug, Clone)]
 pub struct CacheOptionsBuilder<'a> {
@@ -90,11 +83,22 @@ pub struct CacheOptions<'a> {
     pub role: CacheRole,
 }
 
+/// Builds [`ArrayReader`]s from parquet schema, projection mask, and RowGroups reader
+pub struct ArrayReaderBuilder<'a> {
+    /// Source of row group data
+    row_groups: &'a dyn RowGroups,
+    /// Optional cache options for the array reader
+    cache_options: Option<&'a CacheOptions<'a>>,
+    /// metrics
+    metrics: &'a ArrowReaderMetrics,
+}
+
 impl<'a> ArrayReaderBuilder<'a> {
-    pub fn new(row_groups: &'a dyn RowGroups) -> Self {
+    pub fn new(row_groups: &'a dyn RowGroups, metrics: &'a ArrowReaderMetrics) -> Self {
         Self {
             row_groups,
             cache_options: None,
+            metrics,
         }
     }
 
@@ -143,6 +147,7 @@ impl<'a> ArrayReaderBuilder<'a> {
                         Arc::clone(&cache_options.cache),
                         col_idx,
                         cache_options.role,
+                        self.metrics.clone(), // cheap clone
                     ))))
                 } else {
                     Ok(Some(reader))
@@ -453,7 +458,8 @@ mod tests {
         )
         .unwrap();
 
-        let array_reader = ArrayReaderBuilder::new(&file_reader)
+        let metrics = ArrowReaderMetrics::disabled();
+        let array_reader = ArrayReaderBuilder::new(&file_reader, &metrics)
             .build_array_reader(fields.as_ref(), &mask)
             .unwrap();