Merge pull request #7302 from sundy-li/topn_pruner

feat(storage): add topn pruner

Author: BohuTANG

src/query/service/tests/it/storages/fuse/pruning.rs

@@ -26,6 +26,7 @@ use common_planners::col;
 use common_planners::lit;
 use common_planners::sub;
 use common_planners::CreateTablePlan;
+use common_planners::Expression;
 use common_planners::Extras;
 use databend_query::interpreters::CreateTableInterpreter;
 use databend_query::interpreters::Interpreter;
@@ -145,49 +146,57 @@ async fn test_block_pruner() -> Result<()> {
     let reader = MetaReaders::table_snapshot_reader(ctx.clone());
     let snapshot = reader.read(snapshot_loc.as_str(), None, 1).await?;
 
-    // nothing will be pruned
-    let push_downs = None;
-    let blocks = apply_block_pruning(
-        snapshot.clone(),
-        table.get_table_info().schema(),
-        &push_downs,
-        ctx.clone(),
-    )
-    .await?;
-    let rows = blocks.iter().map(|b| b.row_count as usize).sum::<usize>();
-    assert_eq!(rows, num_blocks * row_per_block);
-    assert_eq!(num_blocks, blocks.len());
-
-    // fully pruned
-    let mut extra = Extras::default();
-    // max value of col a is
-    let pred = col("a").gt(lit(30u64));
-    extra.filters = vec![pred];
-
-    let blocks = apply_block_pruning(
-        snapshot.clone(),
-        table.get_table_info().schema(),
-        &Some(extra),
-        ctx.clone(),
-    )
-    .await?;
-    assert_eq!(0, blocks.len());
+    // nothing is pruned
+    let mut e1 = Extras::default();
+    e1.filters = vec![col("a").gt(lit(30u64))];
 
     // some blocks pruned
-    let mut extra = Extras::default();
+    let mut e2 = Extras::default();
     let max_val_of_b = 6u64;
-    let pred = col("a").gt(lit(0u64)).and(col("b").gt(lit(max_val_of_b)));
-    extra.filters = vec![pred];
+    e2.filters = vec![col("a").gt(lit(0u64)).and(col("b").gt(lit(max_val_of_b)))];
+    let b2 = num_blocks - max_val_of_b as usize - 1;
+
+    // Sort asc Limit
+    let mut e3 = Extras::default();
+    e3.order_by = vec![Expression::Sort {
+        expr: Box::new(col("b")),
+        asc: true,
+        nulls_first: false,
+        origin_expr: Box::new(col("b")),
+    }];
+    e3.limit = Some(3);
+
+    // Sort desc Limit
+    let mut e4 = Extras::default();
+    e4.order_by = vec![Expression::Sort {
+        expr: Box::new(col("b")),
+        asc: false,
+        nulls_first: false,
+        origin_expr: Box::new(col("b")),
+    }];
+    e4.limit = Some(3);
+
+    let extras = vec![
+        (None, num_blocks, num_blocks * row_per_block),
+        (Some(e1), 0, 0),
+        (Some(e2), b2, b2 * row_per_block),
+        (Some(e3), 3, 3 * row_per_block),
+        (Some(e4), 4, 4 * row_per_block),
+    ];
 
-    let blocks = apply_block_pruning(
-        snapshot.clone(),
-        table.get_table_info().schema(),
-        &Some(extra),
-        ctx.clone(),
-    )
-    .await?;
+    for (extra, expected_blocks, expected_rows) in extras {
+        let blocks = apply_block_pruning(
+            snapshot.clone(),
+            table.get_table_info().schema(),
+            &extra,
+            ctx.clone(),
+        )
+        .await?;
 
-    assert_eq!((num_blocks - max_val_of_b as usize - 1), blocks.len());
+        let rows = blocks.iter().map(|b| b.row_count as usize).sum::<usize>();
+        assert_eq!(expected_rows, rows);
+        assert_eq!(expected_blocks, blocks.len());
+    }
 
     Ok(())
 }

src/query/storages/fuse/src/operations/read_partitions.rs

@@ -51,7 +51,6 @@ impl FuseTable {
                     .into_iter()
                     .map(|(_, v)| v)
                     .collect::<Vec<_>>();
-
                 let partitions_total = snapshot.summary.block_count as usize;
                 self.read_partitions_with_metas(ctx, push_downs, block_metas, partitions_total)
             }

src/query/storages/fuse/src/pruning/mod.rs

@@ -16,4 +16,6 @@ mod bloom_pruner;
 mod limiter;
 mod pruning_executor;
 mod range_pruner;
+mod topn_pruner;
+
 pub use pruning_executor::BlockPruner;

src/query/storages/fuse/src/pruning/pruning_executor.rs

@@ -33,6 +33,7 @@ use super::bloom_pruner;
 use crate::io::MetaReaders;
 use crate::pruning::limiter;
 use crate::pruning::range_pruner;
+use crate::pruning::topn_pruner;
 
 pub struct BlockPruner {
     table_snapshot: Arc<TableSnapshot>,
@@ -150,15 +151,33 @@ impl BlockPruner {
             .map_err(|e| ErrorCode::StorageOther(format!("block pruning failure, {}", e)))?;
 
         // 3. collect the result
-        tracing::debug_span!("collect_result").in_scope(|| {
-            // flatten the collected block metas
-            let metas = joint
-                .into_iter()
-                .collect::<Result<Vec<_>>>()?
-                .into_iter()
-                .flatten();
-            Ok(metas.collect())
-        })
+        let metas: Result<Vec<(usize, BlockMeta)>> = tracing::debug_span!("collect_result")
+            .in_scope(|| {
+                // flatten the collected block metas
+                let metas = joint
+                    .into_iter()
+                    .collect::<Result<Vec<_>>>()?
+                    .into_iter()
+                    .flatten();
+                Ok(metas.collect())
+            });
+        let metas = metas?;
+
+        // if there are ordering + limit clause, use topn pruner
+
+        if push_down
+            .as_ref()
+            .filter(|p| !p.order_by.is_empty() && p.limit.is_some())
+            .is_some()
+        {
+            let push_down = push_down.as_ref().unwrap();
+            let limit = push_down.limit.unwrap();
+            let sort = push_down.order_by.clone();
+            let tpruner = topn_pruner::TopNPrunner::new(schema, sort, limit);
+            return tpruner.prune(metas);
+        }
+
+        Ok(metas)
     }
 
     async fn all_the_blocks(

src/query/storages/fuse/src/pruning/topn_pruner.rs

@@ -0,0 +1,102 @@
+//  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_datavalues::DataSchemaRef;
+use common_exception::ErrorCode;
+use common_exception::Result;
+use common_fuse_meta::meta::BlockMeta;
+use common_fuse_meta::meta::ColumnStatistics;
+use common_planners::Expression;
+
+pub(crate) struct TopNPrunner {
+    schema: DataSchemaRef,
+    sort: Vec<Expression>,
+    limit: usize,
+}
+
+impl TopNPrunner {
+    pub(crate) fn new(schema: DataSchemaRef, sort: Vec<Expression>, limit: usize) -> Self {
+        Self {
+            schema,
+            sort,
+            limit,
+        }
+    }
+}
+
+impl TopNPrunner {
+    pub(crate) fn prune(&self, metas: Vec<(usize, BlockMeta)>) -> Result<Vec<(usize, BlockMeta)>> {
+        if self.sort.len() != 1 {
+            return Ok(metas);
+        }
+
+        if self.limit >= metas.len() {
+            return Ok(metas);
+        }
+
+        let (sort, asc, nulls_first) = match &self.sort[0] {
+            Expression::Sort {
+                expr,
+                asc,
+                nulls_first,
+                ..
+            } => (expr, asc, nulls_first),
+            _ => unreachable!(),
+        };
+
+        // Currently, we only support topn on single-column sort.
+        // TODO: support monadic + multi expression + order by cluster key sort.
+        let column = if let Expression::Column(c) = sort.as_ref() {
+            c
+        } else {
+            return Ok(metas);
+        };
+
+        let sort_idx = if let Ok(index) = self.schema.index_of(column.as_str()) {
+            index as u32
+        } else {
+            return Ok(metas);
+        };
+
+        let mut id_stats = metas
+            .iter()
+            .map(|(id, meta)| {
+                let stat = meta.col_stats.get(&sort_idx).ok_or_else(|| {
+                    ErrorCode::UnknownException(format!(
+                        "Unable to get the colStats by ColumnId: {}",
+                        sort_idx
+                    ))
+                })?;
+                Ok((*id, stat.clone(), meta.clone()))
+            })
+            .collect::<Result<Vec<(usize, ColumnStatistics, BlockMeta)>>>()?;
+
+        id_stats.sort_by(|a, b| {
+            if a.1.null_count + b.1.null_count != 0 && *nulls_first {
+                return a.1.null_count.cmp(&b.1.null_count).reverse();
+            }
+            // no nulls
+            if *asc {
+                a.1.min.cmp(&b.1.min)
+            } else {
+                a.1.max.cmp(&b.1.max).reverse()
+            }
+        });
+        Ok(id_stats
+            .iter()
+            .map(|s| (s.0, s.2.clone()))
+            .take(self.limit as usize)
+            .collect())
+    }
+}