Merge pull request #9228 from RinChanNOWWW/prune-row-group

BohuTANG · web-flow · commit 91be29689923 · 2022-12-14T15:34:30.000+08:00
feat(parquet): prune row groups before reading.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -60,6 +60,7 @@ members = [
     "src/query/storages/system",
     "src/query/storages/view",
     "src/query/storages/parquet",
+    "src/query/storages/pruner",
     "src/query/users",
     # databend-query
     "src/query/service",
diff --git a/src/query/catalog/src/plan/pushdown.rs b/src/query/catalog/src/plan/pushdown.rs
@@ -21,7 +21,7 @@ use crate::plan::Projection;
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
 pub struct PrewhereInfo {
-    /// columns to be ouput be prewhere scan
+    /// columns to be output by prewhere scan
     pub output_columns: Projection,
     /// columns used for prewhere
     pub prewhere_columns: Projection,
diff --git a/src/query/storages/fuse/fuse/Cargo.toml b/src/query/storages/fuse/fuse/Cargo.toml
@@ -31,6 +31,7 @@ common-sql = { path = "../../../sql" }
 common-storage = { path = "../../../../common/storage" }
 common-storages-cache = { path = "../../cache" }
 common-storages-index = { path = "../../index" }
+common-storages-pruner = { path = "../../pruner" }
 common-storages-table-meta = { path = "../../table-meta" }
 
 async-trait = { version = "0.1.57", package = "async-trait-fn" }
diff --git a/src/query/storages/fuse/fuse/src/pruning/mod.rs b/src/query/storages/fuse/fuse/src/pruning/mod.rs
@@ -12,10 +12,8 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
-mod limiter;
 mod pruner;
 mod pruning_executor;
-mod range_pruner;
 mod topn_pruner;
 
 pub use pruning_executor::BlockIndex;
diff --git a/src/query/storages/fuse/fuse/src/pruning/pruning_executor.rs b/src/query/storages/fuse/fuse/src/pruning/pruning_executor.rs
@@ -25,6 +25,10 @@ use common_catalog::table_context::TableContext;
 use common_datavalues::DataSchemaRef;
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_storages_pruner::LimiterPruner;
+use common_storages_pruner::LimiterPrunerCreator;
+use common_storages_pruner::RangePruner;
+use common_storages_pruner::RangePrunerCreator;
 use common_storages_table_meta::meta::BlockMeta;
 use common_storages_table_meta::meta::Location;
 use common_storages_table_meta::meta::SegmentInfo;
@@ -35,11 +39,7 @@ use tracing::Instrument;
 
 use super::pruner;
 use crate::io::MetaReaders;
-use crate::pruning::limiter;
-use crate::pruning::limiter::LimiterPruner;
 use crate::pruning::pruner::Pruner;
-use crate::pruning::range_pruner;
-use crate::pruning::range_pruner::RangePruner;
 use crate::pruning::topn_pruner;
 
 pub type BlockIndex = (usize, usize);
@@ -79,11 +79,11 @@ impl BlockPruner {
             .and_then(|p| p.limit);
 
         // prepare the limiter. in case that limit is none, an unlimited limiter will be returned
-        let limiter = limiter::new_limiter(limit);
+        let limiter = LimiterPrunerCreator::create(limit);
 
         // prepare the range filter.
         // if filter_expression is none, an dummy pruner will be returned, which prunes nothing
-        let range_pruner = range_pruner::new_range_pruner(ctx, filter_expressions, &schema)?;
+        let range_pruner = RangePrunerCreator::try_create(ctx, filter_expressions, &schema)?;
 
         // prepare the filter.
         // None will be returned, if filter is not applicable (e.g. unsuitable filter expression, index not available, etc.)
diff --git a/src/query/storages/fuse/fuse/src/pruning/topn_pruner.rs b/src/query/storages/fuse/fuse/src/pruning/topn_pruner.rs
@@ -24,18 +24,14 @@ use common_storages_table_meta::meta::ColumnStatistics;
 
 use crate::pruning::BlockIndex;
 
-pub(crate) struct TopNPrunner {
+pub struct TopNPrunner {
     schema: DataSchemaRef,
     sort: Vec<(Expression, bool, bool)>,
     limit: usize,
 }
 
 impl TopNPrunner {
-    pub(crate) fn new(
-        schema: DataSchemaRef,
-        sort: Vec<(Expression, bool, bool)>,
-        limit: usize,
-    ) -> Self {
+    pub fn new(schema: DataSchemaRef, sort: Vec<(Expression, bool, bool)>, limit: usize) -> Self {
         Self {
             schema,
             sort,
@@ -45,7 +41,7 @@ impl TopNPrunner {
 }
 
 impl TopNPrunner {
-    pub(crate) fn prune(
+    pub fn prune(
         &self,
         metas: Vec<(BlockIndex, Arc<BlockMeta>)>,
     ) -> Result<Vec<(BlockIndex, Arc<BlockMeta>)>> {
diff --git a/src/query/storages/parquet/Cargo.toml b/src/query/storages/parquet/Cargo.toml
@@ -24,6 +24,8 @@ common-meta-app = { path = "../../../meta/app" }
 common-pipeline-core = { path = "../../pipeline/core" }
 common-sql = { path = "../../sql" }
 common-storage = { path = "../../../common/storage" }
+common-storages-pruner = { path = "../pruner" }
+common-storages-table-meta = { path = "../table-meta" }
 
 async-trait = { version = "0.1.57", package = "async-trait-fn" }
 chrono = { workspace = true }
diff --git a/src/query/storages/parquet/src/parquet_reader/meta.rs b/src/query/storages/parquet/src/parquet_reader/meta.rs
@@ -12,13 +12,22 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
+use std::collections::HashSet;
 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_arrow::parquet::metadata::RowGroupMetaData;
+use common_datavalues::Column;
+use common_datavalues::ColumnRef;
+use common_datavalues::IntoColumn;
+use common_datavalues::UInt64Column;
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_storages_table_meta::meta::ColumnStatistics;
+use common_storages_table_meta::meta::StatisticsOfColumns;
 
 use crate::ParquetReader;
 
@@ -43,4 +52,79 @@ impl ParquetReader {
         });
         Ok(arrow_schema)
     }
+
+    /// Collect statistics of a batch of row groups of the specified columns.
+    ///
+    /// The retuened vector's length is the same as `rgs`.
+    pub fn collect_row_group_stats(
+        schema: &ArrowSchema,
+        rgs: &[RowGroupMetaData],
+        indices: &HashSet<usize>,
+    ) -> Result<Vec<StatisticsOfColumns>> {
+        let mut stats = Vec::with_capacity(rgs.len());
+        let mut stats_of_row_groups = HashMap::with_capacity(rgs.len());
+
+        for index in indices {
+            if rgs
+                .iter()
+                .any(|rg| rg.columns()[*index].metadata().statistics.is_none())
+            {
+                return Err(ErrorCode::InvalidArgument(
+                    "Some columns of the row groups have no statistics",
+                ));
+            }
+
+            let field = &schema.fields[*index];
+            let column_stats = pread::statistics::deserialize(field, rgs)?;
+            stats_of_row_groups.insert(*index, BatchStatistics::from(column_stats));
+        }
+
+        for (rg_idx, _) in rgs.iter().enumerate() {
+            let mut cols_stats = HashMap::with_capacity(stats.capacity());
+            for index in indices {
+                let col_stats = stats_of_row_groups[index].get(rg_idx);
+                cols_stats.insert(*index as u32, col_stats);
+            }
+            stats.push(cols_stats);
+        }
+
+        Ok(stats)
+    }
+}
+
+/// A temporary struct to present [`pread::statistics::Statistics`].
+///
+/// Convert the inner fields into Databend data structures.
+pub struct BatchStatistics {
+    pub null_count: UInt64Column,
+    pub distinct_count: UInt64Column,
+    pub min_values: ColumnRef,
+    pub max_values: ColumnRef,
+}
+
+impl BatchStatistics {
+    pub fn get(&self, index: usize) -> ColumnStatistics {
+        ColumnStatistics {
+            min: self.min_values.get(index),
+            max: self.max_values.get(index),
+            null_count: self.null_count.get_u64(index).unwrap(),
+            in_memory_size: 0, // this field is not used.
+            distinct_of_values: self.distinct_count.get_u64(index).ok(),
+        }
+    }
+}
+
+impl From<pread::statistics::Statistics> for BatchStatistics {
+    fn from(stats: pread::statistics::Statistics) -> Self {
+        let null_count = UInt64Column::from_arrow_array(&*stats.null_count);
+        let distinct_count = UInt64Column::from_arrow_array(&*stats.distinct_count);
+        let min_values = stats.min_value.clone().into_column();
+        let max_values = stats.max_value.clone().into_column();
+        Self {
+            null_count,
+            distinct_count,
+            min_values,
+            max_values,
+        }
+    }
 }
diff --git a/src/query/storages/parquet/src/parquet_reader/mod.rs b/src/query/storages/parquet/src/parquet_reader/mod.rs
@@ -58,7 +58,7 @@ pub struct ParquetReader {
     /// The actual schema used to read parquet. It will be converted to [`common_datavalues::DataSchema`] when output [`common_datablocks::DataBlock`].
     ///
     /// The reason of using [`ArrowSchema`] to read parquet is that
-    /// There are some types that Databend not support such as Timestmap of nanoseconds.
+    /// There are some types that Databend not support such as Timestamp of nanoseconds.
     /// Such types will be convert to supported types after deserialization.
     projected_arrow_schema: ArrowSchema,
     /// [`ColumnLeaves`] corresponding to the `projected_schema`.
diff --git a/src/query/storages/parquet/src/table_function/read.rs b/src/query/storages/parquet/src/table_function/read.rs
@@ -15,6 +15,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use common_arrow::arrow::io::parquet::read as pread;
 use common_catalog::plan::DataSourcePlan;
 use common_catalog::plan::Partitions;
 use common_catalog::plan::PartitionsShuffleKind;
@@ -26,6 +27,7 @@ use common_exception::Result;
 use common_pipeline_core::Pipeline;
 use common_sql::evaluator::EvalNode;
 use common_sql::evaluator::Evaluator;
+use common_storages_pruner::RangePrunerCreator;
 
 use super::ParquetTable;
 use super::TableContext;
@@ -134,11 +136,43 @@ impl ParquetTable {
         // `dummy_reader` is only used for prune columns in row groups.
         let (_, _, _, columns_to_read) =
             ParquetReader::do_projection(&plan.source_info.schema().to_arrow(), &columns_to_read)?;
+
+        // do parition at the begin of the whole pipeline.
+        let push_downs = plan.push_downs.clone();
+        let schema = plan.source_info.schema();
         pipeline.set_on_init(move || {
             let mut partitions = Vec::with_capacity(locations.len());
+
+            // build row group pruner.
+            let filter_expr = push_downs.as_ref().map(|extra| extra.filters.as_slice());
+            let row_group_pruner = RangePrunerCreator::try_create(&ctx_ref, filter_expr, &schema)?;
+
             for location in &locations {
                 let file_meta = ParquetReader::read_meta(location)?;
-                for rg in &file_meta.row_groups {
+                let arrow_schema = pread::infer_schema(&file_meta)?;
+                let mut row_group_pruned = vec![false; file_meta.row_groups.len()];
+
+                // If collecting stats fails or `should_keep` is true, we still read the row group.
+                // Otherwise, the row group will be pruned.
+                if let Ok(row_group_stats) = ParquetReader::collect_row_group_stats(
+                    &arrow_schema,
+                    &file_meta.row_groups,
+                    &columns_to_read,
+                ) {
+                    for (idx, (stats, rg)) in row_group_stats
+                        .iter()
+                        .zip(file_meta.row_groups.iter())
+                        .enumerate()
+                    {
+                        row_group_pruned[idx] =
+                            !row_group_pruner.should_keep(stats, rg.num_rows() as u64);
+                    }
+                }
+
+                for (idx, rg) in file_meta.row_groups.iter().enumerate() {
+                    if row_group_pruned[idx] {
+                        continue;
+                    }
                     let mut column_metas = HashMap::with_capacity(columns_to_read.len());
                     for index in &columns_to_read {
                         let c = &rg.columns()[*index];
diff --git a/src/query/storages/pruner/Cargo.toml b/src/query/storages/pruner/Cargo.toml
@@ -0,0 +1,21 @@
+[package]
+name = "common-storages-pruner"
+version = { workspace = true }
+authors = { workspace = true }
+license = { workspace = true }
+publish = { workspace = true }
+edition = { workspace = true }
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[lib]
+doctest = false
+test = false
+
+[dependencies]
+common-catalog = { path = "../../../common/../query/catalog" }
+common-datavalues = { path = "../../datavalues" }
+common-exception = { path = "../../../common/exception" }
+common-storages-index = { path = "../index" }
+common-storages-table-meta = { path = "../table-meta" }
+
+tracing = "0.1.36"
diff --git a/src/query/storages/pruner/src/lib.rs b/src/query/storages/pruner/src/lib.rs
@@ -0,0 +1,23 @@
+//  Copyright 2021 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.
+
+#![deny(unused_crate_dependencies)]
+
+mod limiter_pruner;
+mod range_pruner;
+
+pub use limiter_pruner::LimiterPruner;
+pub use limiter_pruner::LimiterPrunerCreator;
+pub use range_pruner::RangePruner;
+pub use range_pruner::RangePrunerCreator;
diff --git a/src/query/storages/pruner/src/limiter_pruner.rs b/src/query/storages/pruner/src/limiter_pruner.rs
@@ -59,9 +59,14 @@ impl Limiter for U64Limiter {
 }
 
 pub type LimiterPruner = Arc<dyn Limiter + Send + Sync>;
-pub fn new_limiter(limit: Option<usize>) -> LimiterPruner {
-    match limit {
-        Some(size) => Arc::new(U64Limiter::new(size as u64)),
-        _ => Arc::new(Unlimited),
+
+pub struct LimiterPrunerCreator;
+
+impl LimiterPrunerCreator {
+    pub fn create(limit: Option<usize>) -> LimiterPruner {
+        match limit {
+            Some(size) => Arc::new(U64Limiter::new(size as u64)),
+            _ => Arc::new(Unlimited),
+        }
     }
 }
diff --git a/src/query/storages/pruner/src/range_pruner.rs b/src/query/storages/pruner/src/range_pruner.rs
