feat(query): support cascading hierarchical grouping sets

src/query/service/src/schedulers/fragments/fragmenter.rs

@@ -111,12 +111,23 @@ impl Fragmenter {
             fragment_id: self.ctx.fragment_id().next_fragment_id(),
             exchange: None,
             query_id: self.query_id.clone(),
+            has_merge_input: false,
             source_fragments: self.fragments,
         });
 
         let edges = Self::collect_fragments_edge(fragments.values());
 
         for (source, target) in edges {
+            let has_merge_input = fragments
+                .get(&source)
+                .is_some_and(|fragment| matches!(fragment.exchange, Some(DataExchange::Merge(_))));
+
+            if has_merge_input {
+                if let Some(fragment) = fragments.get_mut(&target) {
+                    fragment.has_merge_input = true;
+                }
+            }
+
             let Some(fragment) = fragments.get_mut(&source) else {
                 continue;
             };
@@ -320,6 +331,7 @@ impl DeriveHandle for FragmentDeriveHandle {
                 source_fragments: vec![],
                 fragment_id: source_fragment_id,
                 query_id: self.query_id.clone(),
+                has_merge_input: false,
             };
 
             self.fragments.insert(source_fragment_id, source_fragment);
@@ -354,6 +366,7 @@ impl DeriveHandle for FragmentDeriveHandle {
                 fragment_id,
                 exchange: None,
                 query_id: self.query_id.clone(),
+                has_merge_input: false,
                 source_fragments: vec![],
             };
 

src/query/service/src/schedulers/fragments/plan_fragment.rs

@@ -23,6 +23,7 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::Column;
+use databend_common_expression::ColumnBuilder;
 use databend_common_expression::DataBlock;
 use databend_common_settings::ReplaceIntoShuffleStrategy;
 use databend_storages_common_table_meta::meta::BlockSlotDescription;
@@ -36,6 +37,7 @@ use crate::physical_plans::IPhysicalPlan;
 use crate::physical_plans::MutationSource;
 use crate::physical_plans::PhysicalPlan;
 use crate::physical_plans::PhysicalPlanCast;
+use crate::physical_plans::PhysicalPlanMeta;
 use crate::physical_plans::PhysicalPlanVisitor;
 use crate::physical_plans::Recluster;
 use crate::physical_plans::ReplaceDeduplicate;
@@ -77,6 +79,7 @@ pub struct PlanFragment {
     pub fragment_id: usize,
     pub exchange: Option<DataExchange>,
     pub query_id: String,
+    pub has_merge_input: bool,
 
     // The fragments to ask data from.
     pub source_fragments: Vec<PlanFragment>,
@@ -103,18 +106,46 @@ impl PlanFragment {
                 fragment_actions.add_action(action);
             }
             FragmentType::Intermediate => {
-                if self
-                    .source_fragments
-                    .iter()
-                    .any(|fragment| matches!(&fragment.exchange, Some(DataExchange::Merge(_))))
-                {
-                    // If this is a intermediate fragment with merge input,
-                    // we will only send it to coordinator node.
-                    let action = QueryFragmentAction::create(
-                        Fragmenter::get_local_executor(ctx),
-                        self.plan.clone(),
-                    );
+                if self.has_merge_input {
+                    // Only the coordinator can consume the merge input. Other shuffle
+                    // destinations still need this fragment to receive remote data.
+                    let local_executor = Fragmenter::get_local_executor(ctx);
+                    let action =
+                        QueryFragmentAction::create(local_executor.clone(), self.plan.clone());
                     fragment_actions.add_action(action);
+
+                    if let Some(exchange) = &self.exchange {
+                        let mut empty_plan = self.plan.clone();
+                        let Some(exchange_sink) =
+                            ExchangeSink::from_mut_physical_plan(&mut empty_plan)
+                        else {
+                            return Err(ErrorCode::Internal(
+                                "Intermediate fragment exchange plan has no ExchangeSink",
+                            ));
+                        };
+                        exchange_sink.input = PhysicalPlan::new(ConstantTableScan {
+                            meta: PhysicalPlanMeta::new("ConstantTableScan"),
+                            values: exchange_sink
+                                .schema
+                                .fields()
+                                .iter()
+                                .map(|field| {
+                                    ColumnBuilder::with_capacity(field.data_type(), 0).build()
+                                })
+                                .collect(),
+                            num_rows: 0,
+                            output_schema: exchange_sink.schema.clone(),
+                        });
+
+                        for executor in exchange.get_destinations() {
+                            if executor != local_executor {
+                                fragment_actions.add_action(QueryFragmentAction::create(
+                                    executor,
+                                    empty_plan.clone(),
+                                ));
+                            }
+                        }
+                    }
                 } else {
                     // Otherwise distribute the fragment to all the executors.
                     for executor in Fragmenter::get_executors(ctx) {

src/query/settings/src/settings_default.rs

@@ -666,6 +666,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("enable_cascading_grouping_sets", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(0),
+                    desc: "Builds hierarchical grouping sets from the closest available parent grouping.",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Both,
+                    range: Some(SettingRange::Numeric(0..=1)),
+                }),
                 ("storage_fetch_part_num", DefaultSettingValue {
                     value: UserSettingValue::UInt64(2),
                     desc: "Sets the number of partitions that are fetched in parallel from storage during query execution.",

src/query/settings/src/settings_getter_setter.rs

@@ -623,6 +623,10 @@ impl Settings {
         Ok(self.try_get_u64("grouping_sets_to_union")? == 1)
     }
 
+    pub fn get_enable_cascading_grouping_sets(&self) -> Result<bool> {
+        Ok(self.try_get_u64("enable_cascading_grouping_sets")? == 1)
+    }
+
     pub fn get_lazy_read_threshold(&self) -> Result<u64> {
         self.try_get_u64("lazy_read_threshold")
     }

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

@@ -36,6 +36,7 @@ use crate::optimizer::optimizers::CommonSubexpressionOptimizer;
 use crate::optimizer::optimizers::DPhpyOptimizer;
 use crate::optimizer::optimizers::EliminateSelfJoinOptimizer;
 use crate::optimizer::optimizers::distributed::BroadcastToShuffleOptimizer;
+use crate::optimizer::optimizers::distributed::MaterializedCTEDistributionOptimizer;
 use crate::optimizer::optimizers::operator::CleanupUnusedCTEOptimizer;
 use crate::optimizer::optimizers::operator::DeduplicateJoinConditionOptimizer;
 use crate::optimizer::optimizers::operator::FinalizeSpatialJoinOptimizer;
@@ -293,6 +294,8 @@ pub async fn optimize_query(opt_ctx: Arc<OptimizerContext>, s_expr: SExpr) -> Re
         )
         // Cascades optimizer may fail due to timeout, fallback to heuristic optimizer in this case.
         .add(CascadesOptimizer::new(opt_ctx.clone())?)
+        // Normalize distributed MaterializedCTE producers after physical properties are settled.
+        .add(MaterializedCTEDistributionOptimizer::new(opt_ctx.clone()))
         // Eliminate unnecessary scalar calculations to clean up the final plan
         .add_if(
             !opt_ctx.get_planning_agg_index(),

src/query/sql/src/planner/optimizer/optimizers/cascades/cascade.rs

@@ -25,7 +25,6 @@ use crate::optimizer::OptimizerContext;
 use crate::optimizer::cost::CostModel;
 use crate::optimizer::ir::Distribution;
 use crate::optimizer::ir::Memo;
-use crate::optimizer::ir::RelExpr;
 use crate::optimizer::ir::RequiredProperty;
 use crate::optimizer::ir::SExpr;
 use crate::optimizer::optimizers::cascades::cost::DefaultCostModel;
@@ -38,7 +37,6 @@ use crate::optimizer::optimizers::distributed::DistributedOptimizer;
 use crate::optimizer::optimizers::distributed::SortAndLimitPushDownOptimizer;
 use crate::optimizer::optimizers::rule::RuleSet;
 use crate::optimizer::optimizers::rule::TransformResult;
-use crate::plans::RelOperator;
 
 /// A cascades-style search engine to enumerate possible alternations of a relational expression and
 /// find the optimal one.
@@ -93,7 +91,7 @@ impl CascadesOptimizer {
         let result = self.optimize_internal(s_expr.clone());
 
         // Process different cases based on the result
-        let mut optimized_expr = match result {
+        let optimized_expr = match result {
             Ok(expr) => {
                 // After successful optimization, apply sort and limit push down if distributed optimization is enabled
                 if opt_ctx.get_enable_distributed_optimization() {
@@ -122,54 +120,9 @@ impl CascadesOptimizer {
             }
         };
 
-        optimized_expr = Self::remove_exchanges_for_serial_sequence(optimized_expr)?;
-
         Ok(optimized_expr)
     }
 
-    fn remove_exchanges_for_serial_sequence(s_expr: SExpr) -> Result<SExpr> {
-        if Self::has_sequence_with_serial_left_child(&s_expr)? {
-            Self::remove_all_exchanges(s_expr)
-        } else {
-            Ok(s_expr)
-        }
-    }
-
-    fn has_sequence_with_serial_left_child(s_expr: &SExpr) -> Result<bool> {
-        if let RelOperator::Sequence(_) = s_expr.plan.as_ref() {
-            let left_child = s_expr.left_child();
-            let rel_expr = RelExpr::with_s_expr(left_child);
-            let physical_prop = rel_expr.derive_physical_prop()?;
-
-            if physical_prop.distribution == Distribution::Serial {
-                return Ok(true);
-            }
-        }
-
-        for child in s_expr.children() {
-            if Self::has_sequence_with_serial_left_child(child)? {
-                return Ok(true);
-            }
-        }
-
-        Ok(false)
-    }
-
-    fn remove_all_exchanges(s_expr: SExpr) -> Result<SExpr> {
-        if let RelOperator::Exchange(_) = s_expr.plan.as_ref() {
-            return Self::remove_all_exchanges(s_expr.unary_child().clone());
-        }
-
-        let mut new_children = Vec::new();
-        for child in s_expr.children() {
-            let processed_child = Self::remove_all_exchanges(child.clone())?;
-            new_children.push(Arc::new(processed_child));
-        }
-
-        let result = s_expr.replace_children(new_children);
-        Ok(result)
-    }
-
     fn optimize_internal(&mut self, s_expr: SExpr) -> Result<SExpr> {
         // Update rule set based on current flags
         // This ensures we use the most up-to-date flag values, regardless of when the optimizer was created

src/query/sql/src/planner/optimizer/optimizers/distributed/materialized_cte.rs

@@ -0,0 +1,151 @@
+// 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.
+
+use std::sync::Arc;
+
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_expression::Scalar;
+use databend_common_expression::types::NumberScalar;
+
+use crate::optimizer::Optimizer;
+use crate::optimizer::OptimizerContext;
+use crate::optimizer::ir::Distribution;
+use crate::optimizer::ir::RelExpr;
+use crate::optimizer::ir::SExpr;
+use crate::optimizer::ir::SExprVisitor;
+use crate::optimizer::ir::VisitAction;
+use crate::plans::ConstantExpr;
+use crate::plans::Exchange;
+use crate::plans::RelOperator;
+use crate::plans::ScalarExpr;
+
+pub struct MaterializedCTEDistributionOptimizer {
+    ctx: Arc<OptimizerContext>,
+}
+
+impl MaterializedCTEDistributionOptimizer {
+    pub fn new(ctx: Arc<OptimizerContext>) -> Self {
+        Self { ctx }
+    }
+
+    pub fn optimize_sync(&self, s_expr: &SExpr) -> Result<SExpr> {
+        let mut result = if self.ctx.get_enable_distributed_optimization() {
+            s_expr
+                .accept(&mut SerialProducerRedistributor)?
+                .unwrap_or_else(|| s_expr.clone())
+        } else {
+            s_expr.clone()
+        };
+
+        let mut finder = SerialSequenceFinder::default();
+        result.accept(&mut finder)?;
+        if finder.found {
+            result = result.accept(&mut ExchangeRemover)?.unwrap_or(result);
+        }
+
+        Ok(result)
+    }
+}
+
+/// A Sequence producer must participate in the distributed fragment graph. When
+/// its CTE definition ends in a scalar operator, redistribute the result after
+/// that operator instead of forcing the entire query to run without Exchanges.
+/// Hashing by a constant preserves the producer's rows without broadcasting or
+/// changing the scalar operator's empty-input behavior.
+struct SerialProducerRedistributor;
+
+impl SExprVisitor for SerialProducerRedistributor {
+    fn visit(&mut self, _expr: &SExpr) -> Result<VisitAction> {
+        Ok(VisitAction::Continue)
+    }
+
+    fn post_visit(&mut self, expr: &SExpr) -> Result<VisitAction> {
+        if !matches!(expr.plan(), RelOperator::Sequence(_)) {
+            return Ok(VisitAction::Continue);
+        }
+
+        let left = expr.left_child();
+        let physical_prop = RelExpr::with_s_expr(left).derive_physical_prop()?;
+        if physical_prop.distribution != Distribution::Serial {
+            return Ok(VisitAction::Continue);
+        }
+        if !matches!(left.plan(), RelOperator::MaterializedCTE(_)) {
+            return Err(ErrorCode::Internal(
+                "Sequence left child is expected to be MaterializedCTE".to_string(),
+            ));
+        }
+
+        let hash_key = ScalarExpr::ConstantExpr(ConstantExpr {
+            value: Scalar::Number(NumberScalar::UInt32(0)),
+            span: None,
+        });
+        let exchange = left
+            .unary_child_arc()
+            .ref_build_unary(Exchange::GlobalHash(vec![hash_key]));
+        let left = left.replace_children([Arc::new(exchange)]);
+        Ok(VisitAction::Replace(expr.replace_left_child(left)))
+    }
+}
+
+#[derive(Default)]
+struct SerialSequenceFinder {
+    found: bool,
+}
+
+impl SExprVisitor for SerialSequenceFinder {
+    fn visit(&mut self, expr: &SExpr) -> Result<VisitAction> {
+        if self.found {
+            return Ok(VisitAction::SkipChildren);
+        }
+
+        if matches!(expr.plan(), RelOperator::Sequence(_)) {
+            let left = expr.left_child();
+            let physical_prop = RelExpr::with_s_expr(left).derive_physical_prop()?;
+            if physical_prop.distribution == Distribution::Serial {
+                self.found = true;
+                return Ok(VisitAction::SkipChildren);
+            }
+        }
+
+        Ok(VisitAction::Continue)
+    }
+}
+
+struct ExchangeRemover;
+
+impl SExprVisitor for ExchangeRemover {
+    fn visit(&mut self, _expr: &SExpr) -> Result<VisitAction> {
+        Ok(VisitAction::Continue)
+    }
+
+    fn post_visit(&mut self, expr: &SExpr) -> Result<VisitAction> {
+        if matches!(expr.plan(), RelOperator::Exchange(_)) {
+            Ok(VisitAction::Replace(expr.unary_child().clone()))
+        } else {
+            Ok(VisitAction::Continue)
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl Optimizer for MaterializedCTEDistributionOptimizer {
+    fn name(&self) -> String {
+        "MaterializedCTEDistributionOptimizer".to_string()
+    }
+
+    async fn optimize(&mut self, s_expr: &SExpr) -> Result<SExpr> {
+        self.optimize_sync(s_expr)
+    }
+}