fix(query): aviod building TransformSortMergeLimit with a big number …

…limit (#17339)

* fix

* update

* update

* fix

* fix

src/query/pipeline/transforms/src/processors/transforms/transform_sort_merge.rs

@@ -43,7 +43,7 @@ use crate::processors::sort::Merger;
 pub struct TransformSortMerge<R: Rows> {
     schema: DataSchemaRef,
     enable_loser_tree: bool,
-
+    limit: Option<usize>,
     block_size: usize,
     buffer: Vec<Option<(DataBlock, Column)>>,
 
@@ -62,10 +62,12 @@ impl<R: Rows> TransformSortMerge<R> {
         _sort_desc: Arc<Vec<SortColumnDescription>>,
         block_size: usize,
         enable_loser_tree: bool,
+        limit: Option<usize>,
     ) -> Self {
         TransformSortMerge {
             schema,
             enable_loser_tree,
+            limit,
             block_size,
             buffer: vec![],
             aborting: Arc::new(AtomicBool::new(false)),
@@ -171,7 +173,8 @@ impl<R: Rows> TransformSortMerge<R> {
         let streams = self.buffer.drain(..).collect::<Vec<BlockStream>>();
         let mut result = Vec::with_capacity(size_hint);
 
-        let mut merger = Merger::<A, _>::create(self.schema.clone(), streams, batch_size, None);
+        let mut merger =
+            Merger::<A, _>::create(self.schema.clone(), streams, batch_size, self.limit);
 
         while let Some(block) = merger.next_block()? {
             if unlikely(self.aborting.load(Ordering::Relaxed)) {
@@ -218,7 +221,7 @@ pub fn sort_merge(
         0,
         0,
         sort_spilling_batch_bytes,
-        MergeSortCommonImpl::create(schema, sort_desc, block_size, enable_loser_tree),
+        MergeSortCommonImpl::create(schema, sort_desc, block_size, enable_loser_tree, None),
     )?;
     for block in data_blocks {
         processor.transform(block)?;

src/query/pipeline/transforms/src/processors/transforms/transform_sort_merge_base.rs

@@ -337,7 +337,7 @@ impl TransformSortMergeBuilder {
             !self.schema.has_field(ORDER_COL_NAME)
         });
 
-        if self.limit.is_some() {
+        if self.limit.map(|limit| limit < 10000).unwrap_or_default() {
             self.build_sort_limit()
         } else {
             self.build_sort()
@@ -400,6 +400,7 @@ impl TransformSortMergeBuilder {
                     self.sort_desc,
                     self.block_size,
                     self.enable_loser_tree,
+                    self.limit,
                 ),
             )?,
         ))

src/query/pipeline/transforms/src/processors/transforms/transform_sort_merge_limit.rs

@@ -127,7 +127,6 @@ impl<R: Rows> CursorOrder<R> for LocalCursorOrder {
 
 impl<R: Rows> TransformSortMergeLimit<R> {
     pub fn create(block_size: usize, limit: usize) -> Self {
-        debug_assert!(limit <= 10000, "Too large sort merge limit: {}", limit);
         TransformSortMergeLimit {
             heap: FixedHeap::new(limit),
             buffer: HashMap::with_capacity(limit),

src/query/service/src/pipelines/builders/builder_window.rs

@@ -170,7 +170,9 @@ impl PipelineBuilder {
             })
             .collect::<Result<Vec<_>>>()?;
 
-        if let Some(top_n) = &window_partition.top_n {
+        if let Some(top_n) = &window_partition.top_n
+            && top_n.top < 10000
+        {
             self.main_pipeline.exchange(
                 num_processors,
                 WindowPartitionTopNExchange::create(

src/query/service/tests/it/pipelines/transforms/sort/spill.rs

@@ -71,6 +71,7 @@ fn create_sort_spill_pipeline(
                 sort_desc.clone(),
                 block_size,
                 enable_loser_tree,
+                None,
             ),
         )
     })?;

src/query/settings/src/settings_default.rs

@@ -324,6 +324,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("max_push_down_limit", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(10000),
+                    desc: "Sets the maximum number of rows limit that can be pushed down to the leaf operator.",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Both,
+                    range: Some(SettingRange::Numeric(0..=u64::MAX)),
+                }),
                 ("join_spilling_memory_ratio", DefaultSettingValue {
                     value: UserSettingValue::UInt64(60),
                     desc: "Sets the maximum memory ratio in bytes that hash join can use before spilling data to storage during query execution, 0 is unlimited",

src/query/settings/src/settings_getter_setter.rs

@@ -319,6 +319,10 @@ impl Settings {
         Ok(self.unchecked_try_get_u64("disable_join_reorder")? != 0)
     }
 
+    pub fn get_max_push_down_limit(&self) -> Result<usize> {
+        Ok(self.try_get_u64("max_push_down_limit")? as usize)
+    }
+
     pub fn get_join_spilling_memory_ratio(&self) -> Result<usize> {
         Ok(self.try_get_u64("join_spilling_memory_ratio")? as usize)
     }

src/query/sql/src/planner/optimizer/optimizer.rs

@@ -71,6 +71,7 @@ pub struct OptimizerContext {
     pub(crate) enable_distributed_optimization: bool,
     enable_join_reorder: bool,
     enable_dphyp: bool,
+    pub(crate) max_push_down_limit: usize,
     planning_agg_index: bool,
     #[educe(Debug(ignore))]
     pub(crate) sample_executor: Option<Arc<dyn QueryExecutor>>,
@@ -85,6 +86,7 @@ impl OptimizerContext {
             enable_distributed_optimization: false,
             enable_join_reorder: true,
             enable_dphyp: true,
+            max_push_down_limit: 10000,
             sample_executor: None,
             planning_agg_index: false,
         }
@@ -114,6 +116,11 @@ impl OptimizerContext {
         self.planning_agg_index = true;
         self
     }
+
+    pub fn with_max_push_down_limit(mut self, max_push_down_limit: usize) -> Self {
+        self.max_push_down_limit = max_push_down_limit;
+        self
+    }
 }
 
 /// A recursive optimizer that will apply the given rules recursively.

src/query/sql/src/planner/optimizer/rule/factory.rs

@@ -56,8 +56,6 @@ use crate::optimizer::OptimizerContext;
 
 pub struct RuleFactory;
 
-pub const MAX_PUSH_DOWN_LIMIT: usize = 10000;
-
 impl RuleFactory {
     pub fn create_rule(id: RuleID, ctx: OptimizerContext) -> Result<RulePtr> {
         match id {
@@ -81,15 +79,15 @@ impl RuleFactory {
             }
             RuleID::PushDownLimitOuterJoin => Ok(Box::new(RulePushDownLimitOuterJoin::new())),
             RuleID::PushDownLimitEvalScalar => Ok(Box::new(RulePushDownLimitEvalScalar::new())),
-            RuleID::PushDownLimitSort => {
-                Ok(Box::new(RulePushDownLimitSort::new(MAX_PUSH_DOWN_LIMIT)))
-            }
-            RuleID::PushDownLimitWindow => {
-                Ok(Box::new(RulePushDownLimitWindow::new(MAX_PUSH_DOWN_LIMIT)))
-            }
-            RuleID::RulePushDownRankLimitAggregate => {
-                Ok(Box::new(RulePushDownRankLimitAggregate::new()))
-            }
+            RuleID::PushDownLimitSort => Ok(Box::new(RulePushDownLimitSort::new(
+                ctx.max_push_down_limit,
+            ))),
+            RuleID::PushDownLimitWindow => Ok(Box::new(RulePushDownLimitWindow::new(
+                ctx.max_push_down_limit,
+            ))),
+            RuleID::RulePushDownRankLimitAggregate => Ok(Box::new(
+                RulePushDownRankLimitAggregate::new(ctx.max_push_down_limit),
+            )),
             RuleID::PushDownFilterAggregate => Ok(Box::new(RulePushDownFilterAggregate::new())),
             RuleID::PushDownFilterWindow => Ok(Box::new(RulePushDownFilterWindow::new())),
             RuleID::PushDownFilterWindowTopN => {

src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_limit_aggregate.rs

@@ -41,10 +41,11 @@ use crate::plans::SortItem;
 pub struct RulePushDownRankLimitAggregate {
     id: RuleID,
     matchers: Vec<Matcher>,
+    max_limit: usize,
 }
 
 impl RulePushDownRankLimitAggregate {
-    pub fn new() -> Self {
+    pub fn new(max_limit: usize) -> Self {
         Self {
             id: RuleID::RulePushDownRankLimitAggregate,
             matchers: vec![
@@ -73,6 +74,7 @@ impl RulePushDownRankLimitAggregate {
                     }],
                 },
             ],
+            max_limit,
         }
     }
 
@@ -84,40 +86,44 @@ impl RulePushDownRankLimitAggregate {
         state: &mut TransformResult,
     ) -> databend_common_exception::Result<()> {
         let limit: Limit = s_expr.plan().clone().try_into()?;
-        if let Some(mut count) = limit.limit {
-            count += limit.offset;
-            let agg = s_expr.child(0)?;
-            let mut agg_limit: Aggregate = agg.plan().clone().try_into()?;
-
-            let sort_items = agg_limit
-                .group_items
-                .iter()
-                .map(|g| SortItem {
-                    index: g.index,
-                    asc: true,
-                    nulls_first: false,
-                })
-                .collect::<Vec<_>>();
-            agg_limit.rank_limit = Some((sort_items.clone(), count));
-
-            let sort = Sort {
-                items: sort_items.clone(),
-                limit: Some(count),
-                after_exchange: None,
-                pre_projection: None,
-                window_partition: None,
-            };
-
-            let agg = SExpr::create_unary(
-                Arc::new(RelOperator::Aggregate(agg_limit)),
-                Arc::new(agg.child(0)?.clone()),
-            );
-            let sort = SExpr::create_unary(Arc::new(RelOperator::Sort(sort)), agg.into());
-            let mut result = s_expr.replace_children(vec![Arc::new(sort)]);
-
-            result.set_applied_rule(&self.id);
-            state.add_result(result);
+        let Some(mut count) = limit.limit else {
+            return Ok(());
+        };
+        count += limit.offset;
+        if count > self.max_limit {
+            return Ok(());
         }
+        let agg = s_expr.child(0)?;
+        let mut agg_limit: Aggregate = agg.plan().clone().try_into()?;
+
+        let sort_items = agg_limit
+            .group_items
+            .iter()
+            .map(|g| SortItem {
+                index: g.index,
+                asc: true,
+                nulls_first: false,
+            })
+            .collect::<Vec<_>>();
+        agg_limit.rank_limit = Some((sort_items.clone(), count));
+
+        let sort = Sort {
+            items: sort_items.clone(),
+            limit: Some(count),
+            after_exchange: None,
+            pre_projection: None,
+            window_partition: None,
+        };
+
+        let agg = SExpr::create_unary(
+            Arc::new(RelOperator::Aggregate(agg_limit)),
+            Arc::new(agg.child(0)?.clone()),
+        );
+        let sort = SExpr::create_unary(Arc::new(RelOperator::Sort(sort)), agg.into());
+        let mut result = s_expr.replace_children(vec![Arc::new(sort)]);
+
+        result.set_applied_rule(&self.id);
+        state.add_result(result);
         Ok(())
     }
 
@@ -137,53 +143,55 @@ impl RulePushDownRankLimitAggregate {
             _ => return Ok(()),
         };
 
+        let Some(limit) = sort.limit else {
+            return Ok(());
+        };
+
         let mut agg_limit: Aggregate = agg_limit_expr.plan().clone().try_into()?;
 
-        if let Some(limit) = sort.limit {
-            let is_order_subset = sort
-                .items
-                .iter()
-                .all(|k| agg_limit.group_items.iter().any(|g| g.index == k.index));
+        let is_order_subset = sort
+            .items
+            .iter()
+            .all(|k| agg_limit.group_items.iter().any(|g| g.index == k.index));
+        if !is_order_subset {
+            return Ok(());
+        }
 
-            if !is_order_subset {
-                return Ok(());
-            }
-            let mut sort_items = Vec::with_capacity(agg_limit.group_items.len());
-            let mut not_found_sort_items = vec![];
-            for i in 0..agg_limit.group_items.len() {
-                let group_item = &agg_limit.group_items[i];
-                if let Some(sort_item) = sort.items.iter().find(|k| k.index == group_item.index) {
-                    sort_items.push(SortItem {
-                        index: group_item.index,
-                        asc: sort_item.asc,
-                        nulls_first: sort_item.nulls_first,
-                    });
-                } else {
-                    not_found_sort_items.push(SortItem {
-                        index: group_item.index,
-                        asc: true,
-                        nulls_first: false,
-                    });
-                }
+        let mut sort_items = Vec::with_capacity(agg_limit.group_items.len());
+        let mut not_found_sort_items = vec![];
+        for i in 0..agg_limit.group_items.len() {
+            let group_item = &agg_limit.group_items[i];
+            if let Some(sort_item) = sort.items.iter().find(|k| k.index == group_item.index) {
+                sort_items.push(SortItem {
+                    index: group_item.index,
+                    asc: sort_item.asc,
+                    nulls_first: sort_item.nulls_first,
+                });
+            } else {
+                not_found_sort_items.push(SortItem {
+                    index: group_item.index,
+                    asc: true,
+                    nulls_first: false,
+                });
             }
-            sort_items.extend(not_found_sort_items);
+        }
+        sort_items.extend(not_found_sort_items);
 
-            agg_limit.rank_limit = Some((sort_items, limit));
+        agg_limit.rank_limit = Some((sort_items, limit));
 
-            let agg = SExpr::create_unary(
-                Arc::new(RelOperator::Aggregate(agg_limit)),
-                Arc::new(agg_limit_expr.child(0)?.clone()),
-            );
+        let agg = SExpr::create_unary(
+            Arc::new(RelOperator::Aggregate(agg_limit)),
+            Arc::new(agg_limit_expr.child(0)?.clone()),
+        );
 
-            let mut result = if has_eval_scalar {
-                let eval_scalar = s_expr.child(0)?.replace_children(vec![Arc::new(agg)]);
-                s_expr.replace_children(vec![Arc::new(eval_scalar)])
-            } else {
-                s_expr.replace_children(vec![Arc::new(agg)])
-            };
-            result.set_applied_rule(&self.id);
-            state.add_result(result);
-        }
+        let mut result = if has_eval_scalar {
+            let eval_scalar = s_expr.child(0)?.replace_children(vec![Arc::new(agg)]);
+            s_expr.replace_children(vec![Arc::new(eval_scalar)])
+        } else {
+            s_expr.replace_children(vec![Arc::new(agg)])
+        };
+        result.set_applied_rule(&self.id);
+        state.add_result(result);
         Ok(())
     }
 }

src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_limit_scan.rs

@@ -62,17 +62,19 @@ impl Rule for RulePushDownLimitScan {
 
     fn apply(&self, s_expr: &SExpr, state: &mut TransformResult) -> Result<()> {
         let limit: Limit = s_expr.plan().clone().try_into()?;
-        if let Some(mut count) = limit.limit {
-            let child = s_expr.child(0)?;
-            let mut get: Scan = child.plan().clone().try_into()?;
-            count += limit.offset;
-            get.limit = Some(get.limit.map_or(count, |c| cmp::max(c, count)));
-            let get = SExpr::create_leaf(Arc::new(RelOperator::Scan(get)));
+        let Some(mut count) = limit.limit else {
+            return Ok(());
+        };
+        count += limit.offset;
 
-            let mut result = s_expr.replace_children(vec![Arc::new(get)]);
-            result.set_applied_rule(&self.id);
-            state.add_result(result);
-        }
+        let child = s_expr.child(0)?;
+        let mut get: Scan = child.plan().clone().try_into()?;
+        get.limit = Some(get.limit.map_or(count, |c| cmp::max(c, count)));
+        let get = SExpr::create_leaf(Arc::new(RelOperator::Scan(get)));
+
+        let mut result = s_expr.replace_children(vec![Arc::new(get)]);
+        result.set_applied_rule(&self.id);
+        state.add_result(result);
         Ok(())
     }