feat: reshuffle block metadata after pruning

src/query/service/src/interpreters/interpreter_index_refresh.rs

@@ -46,9 +46,11 @@ use databend_storages_common_table_meta::meta::Location;
 
 use crate::interpreters::Interpreter;
 use crate::physical_plans::DeriveHandle;
+use crate::physical_plans::FuseBlockRead;
 use crate::physical_plans::PhysicalPlan;
 use crate::physical_plans::PhysicalPlanBuilder;
 use crate::physical_plans::PhysicalPlanCast;
+use crate::physical_plans::PhysicalPlanMeta;
 use crate::physical_plans::TableScan;
 use crate::pipelines::PipelineBuildResult;
 use crate::schedulers::build_query_pipeline_without_render_result_set;
@@ -400,6 +402,20 @@ impl DeriveHandle for ReadSourceDeriveHandle {
         v: &PhysicalPlan,
         children: Vec<PhysicalPlan>,
     ) -> std::result::Result<PhysicalPlan, Vec<PhysicalPlan>> {
+        if let Some(read) = FuseBlockRead::from_physical_plan(v) {
+            // Refresh replaces lazy segments with eagerly pruned block partitions, so the
+            // metadata-pruning exchange is no longer applicable.
+            return Ok(PhysicalPlan::new(TableScan {
+                meta: PhysicalPlanMeta::with_plan_id("TableScan", read.meta.plan_id),
+                scan_id: read.scan_id,
+                name_mapping: read.name_mapping.clone(),
+                source: Box::new(self.source.clone()),
+                internal_column: read.internal_column.clone(),
+                table_index: read.table_index,
+                stat_info: read.stat_info.clone(),
+            }));
+        }
+
         let Some(table_scan) = TableScan::from_physical_plan(v) else {
             return Err(children);
         };

src/query/service/src/physical_plans/format/format_fuse_prune.rs

@@ -0,0 +1,63 @@
+// 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 databend_common_ast::ast::FormatTreeNode;
+use databend_common_exception::Result;
+
+use crate::physical_plans::FusePrune;
+use crate::physical_plans::IPhysicalPlan;
+use crate::physical_plans::PhysicalPlanMeta;
+use crate::physical_plans::format::FormatContext;
+use crate::physical_plans::format::PhysicalFormat;
+use crate::physical_plans::format::part_stats_info_to_format_tree;
+
+pub struct FusePruneFormatter<'a> {
+    inner: &'a FusePrune,
+}
+
+impl<'a> FusePruneFormatter<'a> {
+    pub fn create(inner: &'a FusePrune) -> Box<dyn PhysicalFormat + 'a> {
+        Box::new(FusePruneFormatter { inner })
+    }
+}
+
+impl PhysicalFormat for FusePruneFormatter<'_> {
+    fn get_meta(&self) -> &PhysicalPlanMeta {
+        self.inner.get_meta()
+    }
+
+    fn format(&self, ctx: &mut FormatContext<'_>) -> Result<FormatTreeNode<String>> {
+        let table_name = ctx
+            .metadata
+            .table(self.inner.source.table_index)
+            .qualified_name();
+        let mut children = vec![FormatTreeNode::new(format!("table: {table_name}"))];
+        children.extend(part_stats_info_to_format_tree(
+            &self.inner.source.statistics,
+        ));
+
+        Ok(FormatTreeNode::with_children(
+            "FusePrune".to_string(),
+            children,
+        ))
+    }
+
+    fn format_join(&self, _ctx: &mut FormatContext<'_>) -> Result<FormatTreeNode<String>> {
+        Ok(FormatTreeNode::new(self.inner.get_name()))
+    }
+
+    fn partial_format(&self, _ctx: &mut FormatContext<'_>) -> Result<FormatTreeNode<String>> {
+        Ok(FormatTreeNode::new(self.inner.get_name()))
+    }
+}

src/query/service/src/physical_plans/format/mod.rs

@@ -36,6 +36,7 @@ mod format_exchange_sink;
 mod format_exchange_source;
 mod format_expression_scan;
 mod format_filter;
+mod format_fuse_prune;
 mod format_hash_join;
 mod format_limit;
 mod format_materialized_cte;
@@ -86,6 +87,7 @@ pub use format_exchange_sink::*;
 pub use format_exchange_source::*;
 pub use format_expression_scan::*;
 pub use format_filter::*;
+pub use format_fuse_prune::*;
 pub use format_hash_join::*;
 pub use format_limit::*;
 pub use format_materialized_cte::*;

src/query/service/src/physical_plans/mod.rs

@@ -28,12 +28,15 @@ mod physical_constant_table_scan;
 mod physical_copy_into_location;
 mod physical_copy_into_table;
 mod physical_distributed_insert_select;
+mod physical_distributed_pruning;
 mod physical_eval_scalar;
 mod physical_exchange;
 mod physical_exchange_sink;
 mod physical_exchange_source;
 mod physical_expression_scan;
 mod physical_filter;
+mod physical_fuse_block_read;
+mod physical_fuse_prune;
 mod physical_hash_join;
 mod physical_join;
 mod physical_limit;
@@ -82,11 +85,14 @@ pub use physical_copy_into_location::CopyIntoLocation;
 pub use physical_copy_into_table::*;
 pub use physical_cte_consumer::MaterializeCTERef;
 pub use physical_distributed_insert_select::DistributedInsertSelect;
+pub use physical_distributed_pruning::optimize_distributed_fuse_pruning;
 pub use physical_eval_scalar::EvalScalar;
 pub use physical_exchange::Exchange;
 pub use physical_exchange_sink::ExchangeSink;
 pub use physical_exchange_source::ExchangeSource;
 pub use physical_filter::Filter;
+pub use physical_fuse_block_read::FuseBlockRead;
+pub use physical_fuse_prune::FusePrune;
 pub use physical_hash_join::*;
 pub use physical_limit::Limit;
 pub use physical_materialized_cte::*;

src/query/service/src/physical_plans/physical_distributed_pruning.rs

@@ -0,0 +1,143 @@
+// 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::collections::HashSet;
+
+use super::BroadcastSink;
+use super::Exchange;
+use super::FuseBlockRead;
+use super::MaterializedCTE;
+use super::PhysicalPlan;
+use super::PhysicalPlanCast;
+use super::TableScan;
+
+/// Inject the Fuse metadata exchange only after the optimizer has finalized the row-data
+/// exchanges. A scan inside an existing exchange input can be split into a pruning source and a
+/// reader fragment that runs on all executors. A scan in the root fragment may be coordinator-only
+/// and must keep the original `TableScan` execution path.
+pub fn optimize_distributed_fuse_pruning(
+    plan: &PhysicalPlan,
+    eligible_scan_ids: &HashSet<usize>,
+) -> PhysicalPlan {
+    if eligible_scan_ids.is_empty() {
+        return plan.clone();
+    }
+
+    rewrite(plan, eligible_scan_ids, false)
+}
+
+#[recursive::recursive]
+fn rewrite(
+    plan: &PhysicalPlan,
+    eligible_scan_ids: &HashSet<usize>,
+    can_host_distributed_reader: bool,
+) -> PhysicalPlan {
+    if can_host_distributed_reader
+        && let Some(scan) = TableScan::from_physical_plan(plan)
+        && eligible_scan_ids.contains(&scan.scan_id)
+    {
+        return FuseBlockRead::create(scan.clone());
+    }
+
+    // Replacing a scan inside an Exchange input splits that non-root fragment into a pruning
+    // source and an intermediate reader fragment. Fragmenter schedules the reader on all
+    // executors, so every metadata exchange destination has a receiver. Materialized CTE and
+    // broadcast sink inputs are also non-root fragments without requiring a parent Exchange.
+    let children_can_host_distributed_reader = can_host_distributed_reader
+        || Exchange::check_physical_plan(plan)
+        || MaterializedCTE::check_physical_plan(plan)
+        || BroadcastSink::check_physical_plan(plan);
+    let children = plan
+        .children()
+        .map(|child| {
+            rewrite(
+                child,
+                eligible_scan_ids,
+                children_can_host_distributed_reader,
+            )
+        })
+        .collect();
+
+    plan.derive(children)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use databend_common_catalog::plan::DataSourceInfo;
+    use databend_common_catalog::plan::DataSourcePlan;
+    use databend_common_catalog::plan::PartStatistics;
+    use databend_common_catalog::plan::Partitions;
+    use databend_common_expression::TableSchema;
+    use databend_common_meta_app::schema::TableInfo;
+    use databend_common_sql::executor::physical_plans::FragmentKind;
+
+    use super::*;
+    use crate::physical_plans::FusePrune;
+    use crate::physical_plans::PhysicalPlanMeta;
+
+    fn table_scan(scan_id: usize) -> PhysicalPlan {
+        let schema = Arc::new(TableSchema::empty());
+        let mut table_info = TableInfo::simple("default", "t", schema.clone());
+        table_info.meta.engine = "FUSE".to_string();
+        let source = DataSourcePlan {
+            source_info: DataSourceInfo::TableSource(table_info),
+            output_schema: schema,
+            parts: Partitions::default(),
+            statistics: PartStatistics::default(),
+            description: String::new(),
+            tbl_args: None,
+            push_downs: None,
+            internal_columns: None,
+            base_block_ids: None,
+            block_meta_options: Default::default(),
+            table_index: 0,
+            scan_id,
+        };
+
+        TableScan::create(
+            scan_id,
+            Default::default(),
+            Box::new(source),
+            None,
+            None,
+            None,
+        )
+    }
+
+    #[test]
+    fn injects_metadata_exchange_only_into_distributed_fragments() {
+        let eligible = HashSet::from([1]);
+
+        let local = optimize_distributed_fuse_pruning(&table_scan(1), &eligible);
+        assert!(TableScan::check_physical_plan(&local));
+
+        let distributed = PhysicalPlan::new(Exchange {
+            meta: PhysicalPlanMeta::new("Exchange"),
+            input: table_scan(1),
+            kind: FragmentKind::Merge,
+            keys: vec![],
+            ignore_exchange: false,
+            allow_adjust_parallelism: true,
+        });
+        let distributed = optimize_distributed_fuse_pruning(&distributed, &eligible);
+
+        let outer_exchange = Exchange::from_physical_plan(&distributed).unwrap();
+        let block_read = FuseBlockRead::from_physical_plan(&outer_exchange.input).unwrap();
+        let metadata_exchange = Exchange::from_physical_plan(&block_read.input).unwrap();
+        assert_eq!(metadata_exchange.kind, FragmentKind::Normal);
+        assert!(FusePrune::check_physical_plan(&metadata_exchange.input));
+    }
+}