feat: push down filter.

Author: RinChanNOWWW

Cargo.lock

@@ -142,3 +142,4 @@ rpath = false
 # If there are dependencies that need patching, they can be listed below.
 # For example:
 # arrow-format = { git = "https://github.com/datafuse-extras/arrow-format", rev = "78dacc1" }
+parquet2 = { git = "https://github.com/jorgecarleitao/parquet2", rev = "fb08b72" }

Cargo.toml

@@ -15,12 +15,15 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::sync::Mutex;
 
+use common_arrow::arrow::bitmap::Bitmap;
 use common_arrow::arrow::datatypes::Field;
 use common_arrow::arrow::io::parquet::read::column_iter_to_arrays;
 use common_arrow::arrow::io::parquet::read::ArrayIter;
 use common_arrow::arrow::io::parquet::read::RowGroupDeserializer;
 use common_arrow::parquet::metadata::ColumnDescriptor;
+use common_arrow::parquet::page::CompressedPage;
 use common_arrow::parquet::read::BasicDecompressor;
 use common_arrow::parquet::read::PageMetaData;
 use common_arrow::parquet::read::PageReader;
@@ -29,6 +32,7 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_storage::ColumnLeaf;
 
+use super::filter::FilterState;
 use crate::parquet_part::ColumnMeta;
 use crate::parquet_part::ParquetRowGroupPart;
 use crate::ParquetReader;
@@ -38,6 +42,7 @@ impl ParquetReader {
         &self,
         part: &ParquetRowGroupPart,
         chunks: Vec<(usize, Vec<u8>)>,
+        filter: Option<Bitmap>,
     ) -> Result<DataBlock> {
         let mut chunk_map: HashMap<usize, Vec<u8>> = chunks.into_iter().collect();
         let mut columns_array_iter = Vec::with_capacity(self.projected_schema.num_fields());
@@ -62,12 +67,18 @@ impl ParquetReader {
                 metas.push((column_meta, descriptor));
                 chunks.push(column_chunk);
             }
-            columns_array_iter.push(Self::to_array_iter(
-                metas,
-                chunks,
-                part.num_rows,
-                column_leaf.field.clone(),
-            )?);
+            let array_iter = if let Some(ref bitmap) = filter {
+                Self::to_array_iter_with_filter(
+                    metas,
+                    chunks,
+                    part.num_rows,
+                    column_leaf.field.clone(),
+                    bitmap.clone(),
+                )?
+            } else {
+                Self::to_array_iter(metas, chunks, part.num_rows, column_leaf.field.clone())?
+            };
+            columns_array_iter.push(array_iter);
         }
 
         let mut deserializer = RowGroupDeserializer::new(columns_array_iter, part.num_rows, None);
@@ -115,6 +126,75 @@ impl ParquetReader {
         )?)
     }
 
+    /// Almost the same as `to_array_iter`, but with a filter.
+    fn to_array_iter_with_filter(
+        metas: Vec<(&ColumnMeta, &ColumnDescriptor)>,
+        chunks: Vec<Vec<u8>>,
+        rows: usize,
+        field: Field,
+        filter: Bitmap,
+    ) -> Result<ArrayIter<'static>> {
+        let (columns, types) = metas
+            .iter()
+            .zip(chunks.into_iter())
+            .map(|(&(meta, descriptor), chunk)| {
+                let filter_state = Arc::new(Mutex::new(FilterState::new(filter.clone())));
+                let iter_filter_state = filter_state.clone();
+
+                let pages = PageReader::new_with_page_meta(
+                    std::io::Cursor::new(chunk),
+                    PageMetaData {
+                        column_start: meta.offset,
+                        num_values: meta.length as i64,
+                        compression: meta.compression.into(),
+                        descriptor: descriptor.descriptor.clone(),
+                    },
+                    Arc::new(move |_, header| {
+                        // If the bitmap for current page is all unset, skip it.
+                        let mut state = filter_state.lock().unwrap();
+                        let num_rows = header.num_values();
+                        let all_unset = state.range_all_unset(num_rows);
+                        if all_unset {
+                            // skip this page.
+                            state.advance(num_rows);
+                        }
+                        !all_unset
+                    }),
+                    vec![],
+                    usize::MAX,
+                )
+                .map(move |page| {
+                    page.map(|page| match page {
+                        CompressedPage::Data(mut page) => {
+                            let num_rows = page.num_values();
+                            let mut state = iter_filter_state.lock().unwrap();
+                            if state.range_all_unset(num_rows) {
+                                page.select_rows(vec![]);
+                            } else if !state.range_all_set(num_rows) {
+                                page.select_rows(state.convert_to_intervals(num_rows));
+                            };
+                            state.advance(num_rows);
+                            CompressedPage::Data(page)
+                        }
+                        CompressedPage::Dict(_) => page, // do nothing
+                    })
+                });
+                (
+                    BasicDecompressor::new(pages, vec![]),
+                    &descriptor.descriptor.primitive_type,
+                )
+            })
+            .unzip();
+
+        Ok(column_iter_to_arrays(
+            columns,
+            types,
+            field,
+            Some(rows - filter.unset_bits()),
+            rows,
+        )?)
+    }
+
     fn try_next_block(&self, deserializer: &mut RowGroupDeserializer) -> Result<DataBlock> {
         match deserializer.next() {
             None => Err(ErrorCode::Internal(

src/query/storages/parquet/src/parquet_reader/deserialize.rs

@@ -0,0 +1,72 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_arrow::arrow::bitmap::Bitmap;
+use common_arrow::parquet::indexes::Interval;
+
+/// A wrapper of [`Bitmap`] with a position mark. It is used to filter rows when reading a parquet file.
+///
+/// If a whole page is filtered, there is no need to decompress the page.
+/// If a row is filtered, there is no need to decode the row.
+pub struct FilterState {
+    bitmap: Bitmap,
+    pos: usize,
+}
+
+impl FilterState {
+    pub fn new(bitmap: Bitmap) -> Self {
+        Self { bitmap, pos: 0 }
+    }
+
+    #[inline]
+    pub fn advance(&mut self, num: usize) {
+        self.pos += num;
+    }
+
+    /// Return true if [`self.pos`, `self.pos + num`) are set.
+    #[inline]
+    pub fn range_all_set(&self, num: usize) -> bool {
+        self.bitmap.null_count_range(self.pos, num) == 0
+    }
+
+    /// Return true if [`self.pos`, `self.pos + num`) are unset.
+    #[inline]
+    pub fn range_all_unset(&self, num: usize) -> bool {
+        self.bitmap.null_count_range(self.pos, num) == num
+    }
+
+    /// Convert the valditiy of [`self.pos`, `self.pos + num`) to [`Interval`]s.
+    pub fn convert_to_intervals(&self, num_rows: usize) -> Vec<Interval> {
+        let mut res = vec![];
+        let mut started = false;
+        let mut start = 0;
+        for (i, v) in self.bitmap.iter().skip(self.pos).take(num_rows).enumerate() {
+            if v {
+                if !started {
+                    start = i;
+                    started = true;
+                }
+            } else if started {
+                res.push(Interval::new(start, i - start));
+                started = false;
+            }
+        }
+
+        if started {
+            res.push(Interval::new(start, num_rows - start));
+        }
+
+        res
+    }
+}

src/query/storages/parquet/src/parquet_reader/filter.rs

@@ -13,8 +13,10 @@
 // limitations under the License.
 
 mod deserialize;
+mod filter;
 mod meta;
 mod read;
+
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
@@ -38,12 +40,14 @@ pub struct ParquetReader {
     /// The indices of columns need to read by this reader.
     ///
     /// Use [`HashSet`] to avoid duplicate indices.
-    /// Duplicate indices will exist when there are nested types.
+    /// Duplicate indices will exist when there are nested types or
+    /// select a same field multiple times.
     ///
     /// For example:
     ///
     /// ```sql
     /// select a, a.b, a.c from t;
+    /// select a, b, a from t;
     /// ```
     columns_to_read: HashSet<usize>,
     /// The schema of the [`common_datablocks::DataBlock`] this reader produces.

src/query/storages/parquet/src/parquet_reader/mod.rs

@@ -20,8 +20,10 @@ use common_base::base::ProgressValues;
 use common_catalog::plan::PartInfoPtr;
 use common_catalog::table_context::TableContext;
 use common_datablocks::DataBlock;
+use common_datavalues::BooleanColumn;
 use common_datavalues::ColumnRef;
 use common_datavalues::DataSchemaRef;
+use common_datavalues::Series;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_functions::scalars::FunctionContext;
@@ -93,7 +95,7 @@ impl ParquetSource {
     fn do_prewhere_filter(&mut self, part: PartInfoPtr, chunks: Vec<IndexedChunk>) -> Result<()> {
         let rg_part = ParquetRowGroupPart::from_part(&part)?;
         // deserialize prewhere data block first
-        let data_block = self.prewhere_reader.deserialize(rg_part, chunks)?;
+        let data_block = self.prewhere_reader.deserialize(rg_part, chunks, None)?;
         if let Some(filter) = self.prewhere_filter.as_ref() {
             // do filter
             let res = filter
@@ -128,6 +130,7 @@ impl ParquetSource {
                 self.state =
                     Generated(self.ctx.try_get_part(), block.resort(self.output_schema())?);
             } else {
+                let data_block = DataBlock::filter_block(data_block, &filter)?;
                 self.state = State::ReadDataRemain(part, PrewhereData { data_block, filter });
             }
             Ok(())
@@ -153,7 +156,28 @@ impl ParquetSource {
             let block = if chunks.is_empty() {
                 prewhere_blocks
             } else if let Some(remain_reader) = self.remain_reader.as_ref() {
-                let remain_block = remain_reader.deserialize(rg_part, chunks)?;
+                // filter is already converted to non-null boolean column
+                let remain_block = if filter.is_const() && filter.get_bool(0)? {
+                    // don't need filter
+                    remain_reader.deserialize(rg_part, chunks, None)?
+                } else {
+                    let boolean_col = Series::check_get::<BooleanColumn>(&filter)?;
+                    let bitmap = boolean_col.values();
+                    if bitmap.unset_bits() == 0 {
+                        // don't need filter
+                        remain_reader.deserialize(rg_part, chunks, None)?
+                    } else {
+                        remain_reader.deserialize(rg_part, chunks, Some(bitmap.clone()))?
+                    }
+                };
+                assert!(
+                    prewhere_blocks.num_rows() == remain_block.num_rows(),
+                    "prewhere and remain blocks should have same row number. (prewhere: {}, remain: {})",
+                    prewhere_blocks.num_rows(),
+                    remain_block.num_rows()
+                );
+
+                // Combine two blocks.
                 for (col, field) in remain_block
                     .columns()
                     .iter()
@@ -171,11 +195,11 @@ impl ParquetSource {
                 bytes: block.memory_size(),
             };
             self.scan_progress.incr(&progress_values);
-            DataBlock::filter_block(block, &filter)?
+            block
         } else {
             // There is only prewhere reader.
             assert!(self.remain_reader.is_none());
-            let block = self.prewhere_reader.deserialize(rg_part, chunks)?;
+            let block = self.prewhere_reader.deserialize(rg_part, chunks, None)?;
             let progress_values = ProgressValues {
                 rows: block.num_rows(),
                 bytes: block.memory_size(),