refactor: introduce ReadPolicy to parquet reader. (#13020)

* Support tablestats of parquet rs.

* Remove `is_inner_project`.

* Introduce read policy and split parquet readers.

* Support predicate and topk policy.

* Add test cases and fix bugs.

* Fix typo.

* Satisfy the code review.

* Simplify the codes and fix bug.

---------

Co-authored-by: sundyli <[email protected]>

Author: RinChanNOWWW

src/query/catalog/src/plan/pushdown.rs

@@ -132,11 +132,16 @@ impl PushDownInfo {
                 }
 
                 let leaf_fields = schema.leaf_fields();
-                let column_id = leaf_fields
+                let (leaf_id, f) = leaf_fields
                     .iter()
-                    .find(|&p| p == field)
-                    .unwrap()
-                    .column_id();
+                    .enumerate()
+                    .find(|&(_, p)| p == field)
+                    .unwrap();
+                // Databend column id is not equal to parquet leaf id when there is nested type.
+                if f.column_id as usize != leaf_id {
+                    return None;
+                }
+                let column_id = f.column_id;
 
                 let top_k = TopK {
                     limit: self.limit.unwrap(),

src/query/catalog/src/table.rs

@@ -405,7 +405,7 @@ pub enum NavigationPoint {
     TimePoint(DateTime<Utc>),
 }
 
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, Default)]
 pub struct TableStatistics {
     pub num_rows: Option<u64>,
     pub data_size: Option<u64>,

src/query/sql/src/planner/binder/table.rs

@@ -972,13 +972,8 @@ impl Binder {
             }
         }
 
-        let mut stat = table.table().table_statistics()?;
-        if let Some(rows) = statistics_provider.num_rows() {
-            // For external storage (parquet)
-            if let Some(stat) = &mut stat {
-                stat.num_rows = Some(rows);
-            }
-        };
+        let stat = table.table().table_statistics()?;
+
         Ok((
             SExpr::create_leaf(Arc::new(
                 Scan {

src/query/storages/iceberg/src/table.rs

@@ -43,7 +43,7 @@ use common_storage::DataOperator;
 use common_storages_parquet::ParquetFilesPart;
 use common_storages_parquet::ParquetPart;
 use common_storages_parquet::ParquetRSPruner;
-use common_storages_parquet::ParquetRSReader;
+use common_storages_parquet::ParquetRSReaderBuilder;
 use icelake::catalog::Catalog;
 use opendal::Operator;
 use storages_common_pruner::RangePrunerCreator;
@@ -189,12 +189,16 @@ impl IcebergTable {
             options,
         )?;
 
-        let builder =
-            ParquetRSReader::builder(ctx.clone(), self.op.operator(), table_schema, &arrow_schema)?
-                .with_push_downs(plan.push_downs.as_ref())
-                .with_pruner(Some(pruner));
+        let mut builder = ParquetRSReaderBuilder::create(
+            ctx.clone(),
+            self.op.operator(),
+            table_schema,
+            &arrow_schema,
+        )?
+        .with_push_downs(plan.push_downs.as_ref())
+        .with_pruner(Some(pruner));
 
-        let praquet_reader = Arc::new(builder.build()?);
+        let praquet_reader = Arc::new(builder.build_full_reader()?);
 
         // TODO: we need to support top_k.
         let output_schema = Arc::new(DataSchema::from(plan.schema()));

src/query/storages/iceberg/src/table_source.rs

@@ -28,7 +28,7 @@ use common_pipeline_core::processors::processor::Event;
 use common_pipeline_core::processors::processor::ProcessorPtr;
 use common_pipeline_core::processors::Processor;
 use common_storages_parquet::ParquetPart;
-use common_storages_parquet::ParquetRSReader;
+use common_storages_parquet::ParquetRSFullReader;
 use opendal::Reader;
 use parquet::arrow::async_reader::ParquetRecordBatchStream;
 
@@ -45,7 +45,7 @@ pub struct IcebergTableSource {
 
     // Used to read parquet.
     output_schema: DataSchemaRef,
-    parquet_reader: Arc<ParquetRSReader>,
+    parquet_reader: Arc<ParquetRSFullReader>,
     stream: Option<ParquetRecordBatchStream<Reader>>,
 }
 
@@ -54,7 +54,7 @@ impl IcebergTableSource {
         ctx: Arc<dyn TableContext>,
         output: Arc<OutputPort>,
         output_schema: DataSchemaRef,
-        parquet_reader: Arc<ParquetRSReader>,
+        parquet_reader: Arc<ParquetRSFullReader>,
     ) -> Result<ProcessorPtr> {
         let scan_progress = ctx.get_scan_progress();
         Ok(ProcessorPtr::create(Box::new(IcebergTableSource {

src/query/storages/parquet/src/lib.rs

@@ -16,6 +16,8 @@
 #![feature(try_blocks)]
 #![feature(impl_trait_in_assoc_type)]
 #![feature(let_chains)]
+#![feature(core_intrinsics)]
+#![feature(int_roundings)]
 
 mod parquet2;
 mod parquet_part;
@@ -25,7 +27,9 @@ mod utils;
 pub use parquet2::Parquet2Table;
 pub use parquet_part::ParquetFilesPart;
 pub use parquet_part::ParquetPart;
+pub use parquet_rs::ParquetRSFullReader;
 pub use parquet_rs::ParquetRSPruner;
-pub use parquet_rs::ParquetRSReader;
+pub use parquet_rs::ParquetRSReaderBuilder;
 pub use parquet_rs::ParquetRSRowGroupPart;
+pub use parquet_rs::ParquetRSRowGroupReader;
 pub use parquet_rs::ParquetRSTable;

src/query/storages/parquet/src/parquet2/parquet_table/create.rs

@@ -63,10 +63,14 @@ impl Parquet2Table {
             &first_file,
         )
         .await?;
+
+        let num_rows = file_metas.iter().map(|m| m.num_rows as u64).sum();
         let column_statistics_provider =
             create_parquet2_statistics_provider(file_metas, &arrow_schema)?;
 
-        let table_info = create_parquet_table_info(arrow_schema.clone(), &stage_info);
+        let mut table_info = create_parquet_table_info(arrow_schema.clone(), &stage_info);
+        table_info.meta.statistics.number_of_rows = num_rows;
+
         Ok(Arc::new(Parquet2Table {
             table_info,
             arrow_schema,

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

@@ -19,7 +19,9 @@ mod pruning;
 mod source;
 mod statistics;
 
-pub use parquet_reader::ParquetRSReader;
+pub use parquet_reader::ParquetRSFullReader;
+pub use parquet_reader::ParquetRSReaderBuilder;
+pub use parquet_reader::ParquetRSRowGroupReader;
 pub use parquet_table::ParquetRSTable;
 pub use partition::ParquetRSRowGroupPart;
 pub use pruning::ParquetRSPruner;

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

@@ -13,9 +13,13 @@
 // limitations under the License.
 
 mod predicate;
+mod read_policy;
 mod reader;
 mod row_group;
 mod topk;
+mod utils;
 
-pub use reader::ParquetRSReader;
+pub use read_policy::*;
+pub use reader::ParquetRSFullReader;
 pub use reader::ParquetRSReaderBuilder;
+pub use reader::ParquetRSRowGroupReader;

src/query/storages/parquet/src/parquet_rs/parquet_reader/predicate.rs

@@ -12,56 +12,46 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
+
 use arrow_array::BooleanArray;
 use arrow_array::RecordBatch;
-use arrow_schema::FieldRef;
+use common_arrow::arrow::bitmap::Bitmap;
+use common_catalog::plan::PrewhereInfo;
+use common_catalog::plan::Projection;
 use common_exception::Result;
 use common_expression::types::DataType;
 use common_expression::DataBlock;
 use common_expression::Evaluator;
 use common_expression::Expr;
-use common_expression::FieldIndex;
 use common_expression::FunctionContext;
+use common_expression::TableSchema;
 use common_functions::BUILTIN_FUNCTIONS;
+use parquet::arrow::parquet_to_arrow_field_levels;
 use parquet::arrow::FieldLevels;
 use parquet::arrow::ProjectionMask;
+use parquet::schema::types::SchemaDescriptor;
 
-use super::reader::transform_record_batch;
+use super::utils::bitmap_to_boolean_array;
+use super::utils::transform_record_batch;
+use super::utils::FieldPaths;
+use crate::parquet_rs::parquet_reader::utils::compute_output_field_paths;
 
 pub struct ParquetPredicate {
     func_ctx: FunctionContext,
 
     /// Columns used for eval predicate.
     projection: ProjectionMask,
-    /// Projected field levels.
     field_levels: FieldLevels,
 
     /// Predicate filter expression.
     filter: Expr,
-    field_paths: Vec<(FieldRef, Vec<FieldIndex>)>,
+    field_paths: Option<FieldPaths>,
 
-    inner_projected: bool,
+    schema: TableSchema,
 }
 
 impl ParquetPredicate {
-    pub fn new(
-        func_ctx: FunctionContext,
-        projection: ProjectionMask,
-        field_levels: FieldLevels,
-        filter: Expr,
-        field_paths: Vec<(FieldRef, Vec<FieldIndex>)>,
-        inner_projected: bool,
-    ) -> Self {
-        Self {
-            func_ctx,
-            projection,
-            field_levels,
-            filter,
-            field_paths,
-            inner_projected,
-        }
-    }
-
     pub fn projection(&self) -> &ProjectionMask {
         &self.projection
     }
@@ -70,18 +60,59 @@ impl ParquetPredicate {
         &self.field_levels
     }
 
-    pub fn evaluate(&self, batch: &RecordBatch) -> Result<BooleanArray> {
-        let block = if self.inner_projected {
-            transform_record_batch(batch, &self.field_paths)?
-        } else {
-            let (block, _) = DataBlock::from_record_batch(batch)?;
-            block
-        };
-        let evaluator = Evaluator::new(&block, &self.func_ctx, &BUILTIN_FUNCTIONS);
+    pub fn field_paths(&self) -> &Option<FieldPaths> {
+        &self.field_paths
+    }
+
+    pub fn evaluate_block(&self, block: &DataBlock) -> Result<Bitmap> {
+        let evaluator = Evaluator::new(block, &self.func_ctx, &BUILTIN_FUNCTIONS);
         let res = evaluator
             .run(&self.filter)?
-            .convert_to_full_column(&DataType::Boolean, batch.num_rows())
-            .into_arrow_rs()?;
-        Ok(BooleanArray::from(res.to_data()))
+            .convert_to_full_column(&DataType::Boolean, block.num_rows())
+            .as_boolean()
+            .cloned()
+            .unwrap();
+        Ok(res)
     }
+
+    pub fn evaluate(&self, batch: &RecordBatch) -> Result<BooleanArray> {
+        let block = transform_record_batch(batch, &self.field_paths)?;
+        let res = self.evaluate_block(&block)?;
+        Ok(bitmap_to_boolean_array(res))
+    }
+
+    pub fn schema(&self) -> &TableSchema {
+        &self.schema
+    }
+}
+
+/// Build [`PrewhereInfo`] into [`ParquetPredicate`] and get the leave columnd ids.
+pub fn build_predicate(
+    func_ctx: FunctionContext,
+    prewhere: &PrewhereInfo,
+    table_schema: &TableSchema,
+    schema_desc: &SchemaDescriptor,
+) -> Result<(Arc<ParquetPredicate>, Vec<usize>)> {
+    let inner_projection = matches!(prewhere.output_columns, Projection::InnerColumns(_));
+    let schema = prewhere.prewhere_columns.project_schema(table_schema);
+    let filter = prewhere
+        .filter
+        .as_expr(&BUILTIN_FUNCTIONS)
+        .project_column_ref(|name| schema.index_of(name).unwrap());
+    let (projection, leaves) = prewhere.prewhere_columns.to_arrow_projection(schema_desc);
+    let field_paths =
+        compute_output_field_paths(schema_desc, &projection, &schema, inner_projection)?;
+    let field_levels = parquet_to_arrow_field_levels(schema_desc, projection.clone(), None)?;
+
+    Ok((
+        Arc::new(ParquetPredicate {
+            func_ctx,
+            projection,
+            filter,
+            field_levels,
+            field_paths,
+            schema,
+        }),
+        leaves,
+    ))
 }

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

@@ -0,0 +1,26 @@
+// 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.
+
+mod no_prefetch;
+mod predicate_and_topk;
+mod topk_only;
+mod utils;
+
+pub mod policy;
+pub use no_prefetch::NoPrefetchPolicy;
+pub use no_prefetch::NoPretchPolicyBuilder;
+pub use predicate_and_topk::PredicateAndTopkPolicy;
+pub use predicate_and_topk::PredicateAndTopkPolicyBuilder;
+pub use topk_only::TopkOnlyPolicy;
+pub use topk_only::TopkOnlyPolicyBuilder;