Merge pull request #9199 from RinChanNOWWW/read-parquet-filter-push-down

feat: push down filter to parquet reader.

Author: BohuTANG

Cargo.lock

@@ -144,3 +144,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

@@ -16,6 +16,7 @@
 
 mod parquet_read;
 mod parquet_write;
+pub mod schema_projection;
 
 pub use arrow;
 pub use arrow_format;

src/common/arrow/src/lib.rs

@@ -0,0 +1,50 @@
+// 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 std::collections::BTreeMap;
+
+use arrow::datatypes::DataType;
+use arrow::datatypes::Field;
+use arrow::datatypes::Schema;
+
+/// Project a [`Schema`] by picking the fields at the given indices.
+pub fn project(schema: &Schema, indices: &[usize]) -> Schema {
+    let fields = indices
+        .iter()
+        .map(|idx| schema.fields[*idx].clone())
+        .collect::<Vec<_>>();
+    Schema::with_metadata(fields.into(), schema.metadata.clone())
+}
+
+/// Project a [`Schema`] with inner columns by path.
+pub fn inner_project(schema: &Schema, path_indices: &BTreeMap<usize, Vec<usize>>) -> Schema {
+    let paths: Vec<Vec<usize>> = path_indices.values().cloned().collect();
+    let fields = paths
+        .iter()
+        .map(|path| traverse_paths(&schema.fields, path))
+        .collect::<Vec<_>>();
+    Schema::with_metadata(fields.into(), schema.metadata.clone())
+}
+
+fn traverse_paths(fields: &[Field], path: &[usize]) -> Field {
+    assert!(!path.is_empty());
+    let field = &fields[path[0]];
+    if path.len() == 1 {
+        return field.clone();
+    }
+    if let DataType::Struct(inner_fields) = field.data_type() {
+        return traverse_paths(inner_fields, &path[1..]);
+    }
+    unreachable!("Unable to get field paths. Fields: {:?}", fields);
+}

src/common/arrow/src/schema_projection.rs

@@ -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,9 +42,10 @@ 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());
+        let mut columns_array_iter = Vec::with_capacity(self.projected_arrow_schema.fields.len());
 
         let column_leaves = &self.projected_column_leaves.column_leaves;
         let mut cnt_map = Self::build_projection_count_map(column_leaves);
@@ -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,13 +126,82 @@ 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(
                 "deserializer from row group: fail to get a chunk",
             )),
             Some(Err(cause)) => Err(ErrorCode::from(cause)),
-            Some(Ok(chunk)) => DataBlock::from_chunk(&self.projected_schema, &chunk),
+            Some(Ok(chunk)) => DataBlock::from_chunk(&self.output_schema, &chunk),
         }
     }
 

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

@@ -14,9 +14,9 @@
 
 use std::fs::File;
 
+use common_arrow::arrow::datatypes::Schema as ArrowSchema;
 use common_arrow::arrow::io::parquet::read as pread;
 use common_arrow::parquet::metadata::FileMetaData;
-use common_datavalues::DataSchema;
 use common_exception::ErrorCode;
 use common_exception::Result;
 
@@ -36,11 +36,11 @@ impl ParquetReader {
     }
 
     #[inline]
-    pub fn infer_schema(meta: &FileMetaData) -> Result<DataSchema> {
+    pub fn infer_schema(meta: &FileMetaData) -> Result<ArrowSchema> {
         let mut arrow_schema = pread::infer_schema(meta)?;
         arrow_schema.fields.iter_mut().for_each(|f| {
             f.name = f.name.to_lowercase();
         });
-        Ok(DataSchema::from(arrow_schema))
+        Ok(arrow_schema)
     }
 }