Add filter pushdown example (apache#1191)

parquet/examples/parquet_filter_pushdown.rs

@@ -0,0 +1,269 @@
+//! Generates a large parquet file containing dictionary encoded data and demonstrates how
+//! the page index, and the record skipping API can dramatically improve performance
+
+use arrow::array::{
+    ArrayRef, Float64Builder, Int32Builder, StringBuilder, StringDictionaryBuilder,
+};
+use arrow::compute::SlicesIterator;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use arrow::util::pretty::pretty_format_columns;
+use bytes::Bytes;
+use parquet::arrow::arrow_reader::{ArrowReaderOptions, RowSelection};
+use parquet::arrow::{ArrowReader, ArrowWriter, ParquetFileArrowReader, ProjectionMask};
+use parquet::file::properties::{WriterProperties, WriterVersion};
+use std::cmp::Ordering;
+use std::ops::Range;
+use std::sync::Arc;
+use std::time::Instant;
+
+const NUM_ROW_GROUPS: usize = 2;
+const NUM_KEYS: usize = 1024;
+const ROWS_PER_ROW_GROUP: usize = 1024 * 1024;
+const ROWS_PER_FILE: usize = ROWS_PER_ROW_GROUP * NUM_ROW_GROUPS;
+
+fn generate_batch() -> RecordBatch {
+    let string_dict_t =
+        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("dict1", string_dict_t.clone(), true),
+        Field::new("dict2", string_dict_t, true),
+        Field::new("f64_values", DataType::Float64, true),
+    ]));
+
+    let mut dict1 = StringDictionaryBuilder::new(
+        Int32Builder::new(ROWS_PER_FILE),
+        StringBuilder::new(1024),
+    );
+    let mut dict2 = StringDictionaryBuilder::new(
+        Int32Builder::new(ROWS_PER_FILE),
+        StringBuilder::new(1024),
+    );
+    let mut values = Float64Builder::new(ROWS_PER_FILE);
+    let dict: Vec<_> = (0..NUM_KEYS).map(|key| format!("key{}", key)).collect();
+
+    // ~1 runs of each dictionary key
+    let dict1_divisor = ROWS_PER_FILE / NUM_KEYS;
+
+    // ~8 runs for each dictionary key
+    let dict2_divisor = dict1_divisor / 8;
+
+    for i in 0..ROWS_PER_FILE {
+        dict1
+            .append(&dict[(i / dict1_divisor) % dict.len()])
+            .unwrap();
+        dict2
+            .append(&dict[(i / dict2_divisor) % dict.len()])
+            .unwrap();
+
+        values.append_value(i as f64);
+    }
+
+    RecordBatch::try_new(
+        schema,
+        vec![
+            Arc::new(dict1.finish()),
+            Arc::new(dict2.finish()),
+            Arc::new(values.finish()),
+        ],
+    )
+    .unwrap()
+}
+
+fn generate_parquet() -> Vec<u8> {
+    let mut out = Vec::with_capacity(1024);
+
+    let data = generate_batch();
+
+    let props = WriterProperties::builder()
+        .set_writer_version(WriterVersion::PARQUET_2_0)
+        .set_max_row_group_size(ROWS_PER_ROW_GROUP)
+        .build();
+    let mut writer = ArrowWriter::try_new(&mut out, data.schema(), Some(props)).unwrap();
+
+    writer.write(&data).unwrap();
+
+    let metadata = writer.close().unwrap();
+    assert_eq!(metadata.row_groups.len(), 2);
+    assert!(metadata.row_groups[0].columns[0]
+        .column_index_length
+        .is_some());
+    out
+}
+
+fn evaluate_basic(file: Bytes) -> Vec<ArrayRef> {
+    let mut reader = ParquetFileArrowReader::try_new(file).unwrap();
+
+    reader
+        .get_record_reader(1024)
+        .unwrap()
+        .map(|result| {
+            let batch = result.unwrap();
+
+            let filter_a =
+                arrow::compute::eq_dyn_utf8_scalar(&batch.columns()[0], "key0").unwrap();
+            let filter_b =
+                arrow::compute::eq_dyn_utf8_scalar(&batch.columns()[1], "key1").unwrap();
+
+            let combined = arrow::compute::and(&filter_a, &filter_b).unwrap();
+            arrow::compute::filter(&batch.column(2), &combined).unwrap()
+        })
+        .collect()
+}
+
+fn selection_from_ranges(
+    ranges: Vec<Range<usize>>,
+    total_rows: usize,
+) -> Vec<RowSelection> {
+    let mut selection: Vec<RowSelection> = Vec::with_capacity(ranges.len() * 2);
+    let mut last_end = 0;
+    for range in ranges {
+        let len = range.end - range.start;
+
+        match range.start.cmp(&last_end) {
+            Ordering::Equal => match selection.last_mut() {
+                Some(last) => last.row_count += len,
+                None => selection.push(RowSelection::select(len)),
+            },
+            Ordering::Greater => {
+                selection.push(RowSelection::skip(range.start - last_end));
+                selection.push(RowSelection::select(len))
+            }
+            Ordering::Less => panic!("out of order"),
+        }
+        last_end = range.end;
+    }
+
+    if last_end != total_rows {
+        selection.push(RowSelection::skip(total_rows - last_end))
+    }
+
+    selection
+}
+
+fn evaluate_selection(
+    mut reader: ParquetFileArrowReader,
+    column: usize,
+    key: &str,
+) -> Vec<RowSelection> {
+    let mask = ProjectionMask::leaves(reader.parquet_schema(), [column]);
+
+    let mut range_offset = 0;
+    let mut ranges = vec![];
+    for result in reader.get_record_reader_by_columns(mask, 1024).unwrap() {
+        let batch = result.unwrap();
+        let filter =
+            arrow::compute::eq_dyn_utf8_scalar(&batch.columns()[0], key).unwrap();
+
+        let valid = SlicesIterator::new(&filter)
+            .map(|(start, end)| start + range_offset..end + range_offset);
+        ranges.extend(valid);
+        range_offset += batch.num_rows();
+    }
+
+    selection_from_ranges(ranges, range_offset)
+}
+
+// Combine a selection where `second` was computed with `first` applied
+fn combine_selection(
+    first: &[RowSelection],
+    second: &[RowSelection],
+) -> Vec<RowSelection> {
+    let mut selection = vec![];
+    let mut first = first.iter().cloned().peekable();
+    let mut second = second.iter().cloned().peekable();
+
+    let mut to_skip = 0;
+    while let (Some(a), Some(b)) = (first.peek_mut(), second.peek_mut()) {
+        if a.row_count == 0 {
+            first.next().unwrap();
+            continue;
+        }
+
+        if b.row_count == 0 {
+            second.next().unwrap();
+            continue;
+        }
+
+        if a.skip {
+            // Records were skipped when producing second
+            to_skip += a.row_count;
+            first.next().unwrap();
+            continue;
+        }
+
+        let skip = b.skip;
+        let to_process = a.row_count.min(b.row_count);
+
+        a.row_count -= to_process;
+        b.row_count -= to_process;
+
+        match skip {
+            true => to_skip += to_process,
+            false => {
+                if to_skip != 0 {
+                    selection.push(RowSelection::skip(to_skip));
+                    to_skip = 0;
+                }
+                selection.push(RowSelection::select(to_process))
+            }
+        }
+    }
+    selection
+}
+
+fn evaluate_pushdown(file: Bytes) -> Vec<ArrayRef> {
+    // TODO: This could also make use of the page index
+
+    let reader = ParquetFileArrowReader::try_new(file.clone()).unwrap();
+    let s1 = evaluate_selection(reader, 0, "key0");
+
+    // Perhaps we need a way to keep the provide a selection to an existing reader?
+    let options = ArrowReaderOptions::default().with_row_selection(s1.clone());
+    let reader =
+        ParquetFileArrowReader::try_new_with_options(file.clone(), options).unwrap();
+
+    let s2 = evaluate_selection(reader, 1, "key1");
+    let s3 = combine_selection(&s1, &s2);
+
+    let total_rows = s3
+        .iter()
+        .filter_map(|x| (!x.skip).then(|| x.row_count))
+        .sum::<usize>();
+
+    let options = ArrowReaderOptions::default().with_row_selection(s3);
+    let mut reader = ParquetFileArrowReader::try_new_with_options(file, options).unwrap();
+    let mask = ProjectionMask::leaves(reader.parquet_schema(), [2]);
+
+    reader
+        .get_record_reader_by_columns(mask, total_rows)
+        .unwrap()
+        .map(|r| r.unwrap().columns()[0].clone())
+        .collect()
+}
+
+fn main() {
+    let data: Bytes = generate_parquet().into();
+    let t0 = Instant::now();
+    let basic = evaluate_basic(data.clone());
+    let t1 = Instant::now();
+    let complex = evaluate_pushdown(data);
+    let t2 = Instant::now();
+
+    let basic = pretty_format_columns("f64_values", &basic)
+        .unwrap()
+        .to_string();
+
+    let complex = pretty_format_columns("f64_values", &complex)
+        .unwrap()
+        .to_string();
+
+    println!(
+        "Simple strategy took {}s vs {}s",
+        (t1 - t0).as_secs_f64(),
+        (t2 - t1).as_secs_f64()
+    );
+
+    assert_eq!(basic, complex);
+}

parquet/src/arrow/arrow_reader.rs

@@ -75,7 +75,7 @@ pub trait ArrowReader {
 /// [`RowSelection`] allows selecting or skipping a provided number of rows
 /// when scanning the parquet file
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct RowSelection {
+pub struct RowSelection {
     /// The number of rows
     pub row_count: usize,
 
@@ -127,12 +127,7 @@ impl ArrowReaderOptions {
     }
 
     /// Scan rows from the parquet file according to the provided `selection`
-    ///
-    /// TODO: Make public once row selection fully implemented (#1792)
-    pub(crate) fn with_row_selection(
-        self,
-        selection: impl Into<Vec<RowSelection>>,
-    ) -> Self {
+    pub fn with_row_selection(self, selection: impl Into<Vec<RowSelection>>) -> Self {
         Self {
             selection: Some(selection.into()),
             ..self
@@ -359,9 +354,7 @@ impl ParquetRecordBatchReader {
     /// Create a new [`ParquetRecordBatchReader`] that will read at most `batch_size` rows at
     /// a time from [`ArrayReader`] based on the configured `selection`. If `selection` is `None`
     /// all rows will be returned
-    ///
-    /// TODO: Make public once row selection fully implemented (#1792)
-    pub(crate) fn new(
+    pub fn new(
         batch_size: usize,
         array_reader: Box<dyn ArrayReader>,
         selection: Option<VecDeque<RowSelection>>,

parquet/src/file/page_index/index_reader.rs

@@ -35,7 +35,7 @@ pub fn read_columns_indexes<R: ChunkReader>(
     let length = lengths.iter().sum::<usize>();
 
     //read all need data into buffer
-    let mut reader = reader.get_read(offset, reader.len() as usize)?;
+    let mut reader = reader.get_read(offset, length)?;
     let mut data = vec![0; length];
     reader.read_exact(&mut data)?;
 
@@ -65,7 +65,7 @@ pub fn read_pages_locations<R: ChunkReader>(
     let (offset, total_length) = get_location_offset_and_total_length(chunks)?;
 
     //read all need data into buffer
-    let mut reader = reader.get_read(offset, reader.len() as usize)?;
+    let mut reader = reader.get_read(offset, total_length)?;
     let mut data = vec![0; total_length];
     reader.read_exact(&mut data)?;