only read subset of rows

Author: alamb

parquet/benches/arrow_reader_clickbench.rs

@@ -28,9 +28,12 @@
 use arrow::compute::kernels::cmp::neq;
 use arrow::record_batch::RecordBatch;
 use arrow_array::StringViewArray;
+use arrow_schema::{DataType, Fields, Schema};
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use futures::StreamExt;
-use parquet::arrow::arrow_reader::{ArrowPredicate, ArrowPredicateFn, ArrowReaderMetadata, ArrowReaderOptions, RowFilter};
+use parquet::arrow::arrow_reader::{
+    ArrowPredicate, ArrowPredicateFn, ArrowReaderMetadata, ArrowReaderOptions, RowFilter,
+};
 use parquet::arrow::{ArrowSchemaConverter, ParquetRecordBatchStreamBuilder, ProjectionMask};
 use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 use parquet::schema::types::SchemaDescriptor;
@@ -154,6 +157,25 @@ enum QueryPattern {
 }
 
 impl QueryPattern {
+    /// Returns the number of (passing) rows to read from the dataset
+    /// 
+    /// There are 100M rows in the ClickBench dataset, and it is not necessary
+    /// to read and filter all `100M` rows to get a representative set in this
+    /// benchmark.
+    pub(crate) fn row_count(&self) -> usize {
+        match self {
+            QueryPattern::SearchPhrase1 => 100_000,
+            QueryPattern::SearchPhrase2 => 100_000,
+            QueryPattern::SearchPhrase3 => 100_000,
+            QueryPattern::SearchPhrase4 => 100_000,
+            QueryPattern::SearchPhrase5 => 100_000,
+
+            QueryPattern::AdvEngine => 100_000,
+            QueryPattern::MobilePhoneModel1 => 100_000,
+            QueryPattern::MobilePhoneModel2 => 100_000,
+        }
+    }
+    
     pub(crate) fn row_filter(&self, schema: &SchemaDescriptor) -> RowFilter {
         let filter_columns = match self {
             QueryPattern::SearchPhrase1 => {
@@ -177,24 +199,11 @@ impl QueryPattern {
                     move |record_batch: RecordBatch| {
                         let search_phrase = record_batch.column(0);
 
-                        // TODO make it utf8view
-                        println!("Search phrase type: {:?}", search_phrase.data_type());
-
                         let search_phrase = search_phrase
                             .as_any()
                             .downcast_ref::<StringViewArray>()
                             .unwrap();
-                        let res = neq(search_phrase, &empty_string);
-
-                        match &res {
-                            Ok(res) => {
-                                println!("Predicate result is: {:?}", res);
-                            }
-                            Err(e) => {
-                                println!("Predicate error is: {:?}", e);
-                            }
-                        }
-                        res
+                        neq(search_phrase, &empty_string)
                     },
                 )));
             }
@@ -314,10 +323,10 @@ fn benchmark_filters_and_projections(c: &mut Criterion) {
         let hits_path = &files.hits_path;
         let reader_metadata = load_metadata(hits_path);
         let schema_descr = reader_metadata.metadata().file_metadata().schema_descr();
-        
+
         // Determine the correct selection
         let projection_mask = pattern.projection_mask(schema_descr);
-        
+
         group.bench_function(bench_id, |b| {
             b.iter(|| {
                 rt_captured.block_on(async {
@@ -327,18 +336,28 @@ fn benchmark_filters_and_projections(c: &mut Criterion) {
 
                     // Create the arrow predicate (does not implement clone, so must create it each time)
                     let row_filter = pattern.row_filter(schema_descr);
-                    
-                    let mut stream = ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_file, reader_metadata.clone())
-                        .with_batch_size(8192)
-                        .with_projection(projection_mask.clone())
-                        .with_row_filter(row_filter)
-                        .build()
-                        .unwrap();
 
+                    let mut stream = ParquetRecordBatchStreamBuilder::new_with_metadata(
+                        parquet_file,
+                        reader_metadata.clone(),
+                    )
+                    .with_batch_size(8192)
+                    .with_projection(projection_mask.clone())
+                    .with_row_filter(row_filter)
+                    .build()
+                    .unwrap();
+
+                    // read up to the row count limit rows from batches to see
+                    // how fast they can be read;
+                    let mut remain = pattern.row_count();
+                    
                     while let Some(b) = stream.next().await {
-                        // consume the batches to see how fast they can be read
-                        // checking for errors
-                        b.unwrap();
+                        let b = b.unwrap();
+                        let num_rows = b.num_rows();
+                        if remain <= num_rows {
+                            break;
+                        }
+                        remain -= num_rows;
                     }
                 })
             })
@@ -349,20 +368,37 @@ fn benchmark_filters_and_projections(c: &mut Criterion) {
 /// Loads Parquet metadata from the given path, including page indexes
 fn load_metadata(path: &PathBuf) -> ArrowReaderMetadata {
     let file = std::fs::File::open(path).unwrap();
-    let options = ArrowReaderOptions::new()
-        .with_page_index(true);
-    let orig_metadata = ArrowReaderMetadata::load(&file, options.clone()).expect("parquet-metadata loading failed");
-    
+    let options = ArrowReaderOptions::new().with_page_index(true);
+    let orig_metadata =
+        ArrowReaderMetadata::load(&file, options.clone()).expect("parquet-metadata loading failed");
+
     // Update the arrow schema so that it reads View types for binary and utf8 columns
-    let arrow_schema = orig_metadata.schema();
-    
-    
-    let new_arrow_schema = arrow_schema.clone();
+    let new_fields = orig_metadata
+        .schema()
+        .fields()
+        .iter()
+        .map(|f| {
+            //println!("Converting field: {:?}", f);
+            // Read UTF8 fields as Utf8View
+            //
+            // The clickbench_partitioned dataset has textual fields listed as
+            // binary for some historical reason so translate Binary to Utf8View as
+            // well
+            if matches!(f.data_type(), DataType::Utf8 | DataType::Binary) {
+                let new_field = f.as_ref().clone().with_data_type(DataType::Utf8View);
+                //println!("  set to Utf8View");
+                Arc::new(new_field)
+            } else {
+                // otherwise just clone the field
+                Arc::clone(f)
+            }
+        })
+        .collect::<Vec<_>>();
+
+    let new_arrow_schema = Arc::new(Schema::new(new_fields));
+
     let new_options = options.with_schema(new_arrow_schema);
-    ArrowReaderMetadata::try_new(
-        Arc::clone(orig_metadata.metadata()),
-        new_options,
-    ).unwrap()
+    ArrowReaderMetadata::try_new(Arc::clone(orig_metadata.metadata()), new_options).unwrap()
 }
 
 criterion_group!(benches, benchmark_filters_and_projections,);