diff --git a/src/query/service/src/interpreters/interpreter_index_refresh.rs b/src/query/service/src/interpreters/interpreter_index_refresh.rs index 7e9446e17a8..77a96fbf4b0 100644 --- a/src/query/service/src/interpreters/interpreter_index_refresh.rs +++ b/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, ) -> std::result::Result> { + 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); }; diff --git a/src/query/service/src/physical_plans/format/format_fuse_prune.rs b/src/query/service/src/physical_plans/format/format_fuse_prune.rs new file mode 100644 index 00000000000..cceba76d4fa --- /dev/null +++ b/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 { + Box::new(FusePruneFormatter { inner }) + } +} + +impl PhysicalFormat for FusePruneFormatter<'_> { + fn get_meta(&self) -> &PhysicalPlanMeta { + self.inner.get_meta() + } + + fn format(&self, ctx: &mut FormatContext<'_>) -> Result> { + 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> { + Ok(FormatTreeNode::new(self.inner.get_name())) + } + + fn partial_format(&self, _ctx: &mut FormatContext<'_>) -> Result> { + Ok(FormatTreeNode::new(self.inner.get_name())) + } +} diff --git a/src/query/service/src/physical_plans/format/mod.rs b/src/query/service/src/physical_plans/format/mod.rs index 8ebfa745bbe..edbe1711652 100644 --- a/src/query/service/src/physical_plans/format/mod.rs +++ b/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::*; diff --git a/src/query/service/src/physical_plans/mod.rs b/src/query/service/src/physical_plans/mod.rs index 04c2e26c018..982e27fe017 100644 --- a/src/query/service/src/physical_plans/mod.rs +++ b/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::*; diff --git a/src/query/service/src/physical_plans/physical_distributed_pruning.rs b/src/query/service/src/physical_plans/physical_distributed_pruning.rs new file mode 100644 index 00000000000..f5d2d5fc52d --- /dev/null +++ b/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, +) -> 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, + 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)); + } +} diff --git a/src/query/service/src/physical_plans/physical_fuse_block_read.rs b/src/query/service/src/physical_plans/physical_fuse_block_read.rs new file mode 100644 index 00000000000..b4a1a0898b5 --- /dev/null +++ b/src/query/service/src/physical_plans/physical_fuse_block_read.rs @@ -0,0 +1,186 @@ +// 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::any::Any; +use std::collections::BTreeMap; +use std::collections::HashMap; + +use databend_common_catalog::plan::DataSourcePlan; +use databend_common_catalog::plan::InternalColumn; +use databend_common_catalog::plan::Partitions; +use databend_common_exception::Result; +use databend_common_expression::DataSchema; +use databend_common_expression::DataSchemaRef; +use databend_common_expression::FieldIndex; +use databend_common_sql::IndexType; +use databend_common_sql::executor::physical_plans::FragmentKind; +use databend_common_storages_fuse::FuseTable; + +use super::Exchange; +use super::FusePrune; +use super::TableScan; +use super::explain::PlanStatsInfo; +use super::physical_plan::IPhysicalPlan; +use super::physical_plan::PhysicalPlan; +use super::physical_plan::PhysicalPlanMeta; +use super::physical_table_scan::build_scan_output_pipeline; +use crate::pipelines::PipelineBuilder; +use crate::servers::flight::v1::exchange::FusePartExchangeInjector; +use crate::sessions::TableContextTableFactory; + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct FuseBlockRead { + pub meta: PhysicalPlanMeta, + pub input: PhysicalPlan, + pub scan_id: usize, + pub name_mapping: BTreeMap, + pub source: Box, + pub internal_column: Option>, + pub table_index: Option, + pub stat_info: Option, +} + +#[typetag::serde] +impl IPhysicalPlan for FuseBlockRead { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_meta(&self) -> &PhysicalPlanMeta { + &self.meta + } + + fn get_meta_mut(&mut self) -> &mut PhysicalPlanMeta { + &mut self.meta + } + + fn output_schema(&self) -> Result { + TableScan::output_fields(self.source.schema(), &self.name_mapping).map(DataSchema::new_ref) + } + + fn children(&self) -> Box + '_> { + Box::new(std::iter::once(&self.input)) + } + + fn children_mut(&mut self) -> Box + '_> { + Box::new(std::iter::once(&mut self.input)) + } + + fn try_find_single_data_source(&self) -> Option<&DataSourcePlan> { + Some(&self.source) + } + + fn get_desc(&self) -> Result { + Ok(format!( + "{}.{}", + self.source.source_info.catalog_name(), + self.source.source_info.desc() + )) + } + + fn get_labels(&self) -> Result>> { + Ok(HashMap::from([ + (String::from("Full table name"), vec![format!( + "{}.{}", + self.source.source_info.catalog_name(), + self.source.source_info.desc() + )]), + ( + format!( + "Columns ({} / {})", + self.output_schema()?.num_fields(), + std::cmp::max( + self.output_schema()?.num_fields(), + self.source.source_info.schema().num_fields(), + ) + ), + self.name_mapping.keys().cloned().collect(), + ), + ])) + } + + fn derive(&self, mut children: Vec) -> PhysicalPlan { + assert_eq!(children.len(), 1); + PhysicalPlan::new(FuseBlockRead { + meta: self.meta.clone(), + input: children.pop().unwrap(), + scan_id: self.scan_id, + name_mapping: self.name_mapping.clone(), + source: self.source.clone(), + internal_column: self.internal_column.clone(), + table_index: self.table_index, + stat_info: self.stat_info.clone(), + }) + } + + fn build_pipeline2(&self, builder: &mut PipelineBuilder) -> Result<()> { + let old_injector = builder.exchange_injector.clone(); + builder.exchange_injector = FusePartExchangeInjector::create(); + self.input.build_pipeline(builder)?; + builder.exchange_injector = old_injector; + + let table = builder.ctx.build_table_from_source_plan(&self.source)?; + let fuse_table = FuseTable::try_from_table(table.as_ref())?; + fuse_table.do_read_data_from_partitions( + builder.ctx.clone(), + &self.source, + &mut builder.main_pipeline, + true, + )?; + build_scan_output_pipeline( + builder, + &self.source, + &self.name_mapping, + &self.internal_column, + ) + } +} + +impl FuseBlockRead { + pub fn create(scan: TableScan) -> PhysicalPlan { + let TableScan { + scan_id, + name_mapping, + mut source, + internal_column, + table_index, + stat_info, + .. + } = scan; + + let input = PhysicalPlan::new(Exchange { + meta: PhysicalPlanMeta::new("Exchange"), + input: FusePrune::create(source.clone()), + kind: FragmentKind::Normal, + keys: vec![], + ignore_exchange: false, + allow_adjust_parallelism: false, + }); + + // The destination reader only needs schema, push-downs and snapshot information. Keeping + // the global lazy segment list here would duplicate it into every destination fragment. + source.parts = Partitions::default(); + + PhysicalPlan::new(FuseBlockRead { + meta: PhysicalPlanMeta::new("FuseBlockRead"), + input, + scan_id, + name_mapping, + source, + internal_column, + table_index, + stat_info, + }) + } +} diff --git a/src/query/service/src/physical_plans/physical_fuse_prune.rs b/src/query/service/src/physical_plans/physical_fuse_prune.rs new file mode 100644 index 00000000000..b193229752d --- /dev/null +++ b/src/query/service/src/physical_plans/physical_fuse_prune.rs @@ -0,0 +1,142 @@ +// 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::any::Any; +use std::collections::HashMap; + +use databend_common_catalog::plan::DataSourcePlan; +use databend_common_catalog::plan::PartStatistics; +use databend_common_catalog::plan::PartitionsShuffleKind; +use databend_common_exception::Result; +use databend_common_expression::DataSchemaRef; +use databend_common_storages_fuse::FuseTable; + +use super::format::FusePruneFormatter; +use super::format::PhysicalFormat; +use super::physical_plan::IPhysicalPlan; +use super::physical_plan::PhysicalPlan; +use super::physical_plan::PhysicalPlanMeta; +use crate::pipelines::PipelineBuilder; +use crate::servers::flight::v1::exchange::FusePartExchangeInjector; +use crate::sessions::TableContextPartitionStats; +use crate::sessions::TableContextTableFactory; + +/// A Fuse source operator that emits only the block partitions surviving pruning. +/// +/// Its output consists of empty data blocks carrying `BlockPartitionMeta`; actual block data is +/// read by `FuseBlockRead` after the metadata exchange. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct FusePrune { + pub meta: PhysicalPlanMeta, + pub source: Box, +} + +#[typetag::serde] +impl IPhysicalPlan for FusePrune { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_meta(&self) -> &PhysicalPlanMeta { + &self.meta + } + + fn get_meta_mut(&mut self) -> &mut PhysicalPlanMeta { + &mut self.meta + } + + fn output_schema(&self) -> Result { + Ok(DataSchemaRef::default()) + } + + fn formatter(&self) -> Result> { + Ok(FusePruneFormatter::create(self)) + } + + fn try_find_single_data_source(&self) -> Option<&DataSourcePlan> { + Some(&self.source) + } + + fn get_all_data_source(&self, sources: &mut Vec<(u32, Box)>) { + sources.push((self.get_id(), self.source.clone())); + } + + fn set_pruning_stats(&mut self, stats: &mut HashMap) { + if let Some(stat) = stats.remove(&self.get_id()) { + self.source.statistics = stat; + } + } + + fn is_warehouse_distributed_plan(&self) -> bool { + self.source.parts.kind == PartitionsShuffleKind::BroadcastWarehouse + } + + fn get_desc(&self) -> Result { + Ok(format!( + "{}.{}", + self.source.source_info.catalog_name(), + self.source.source_info.desc() + )) + } + + fn get_labels(&self) -> Result>> { + Ok(HashMap::from([ + (String::from("Full table name"), vec![format!( + "{}.{}", + self.source.source_info.catalog_name(), + self.source.source_info.desc() + )]), + (String::from("Total partitions"), vec![ + self.source.statistics.partitions_total.to_string(), + ]), + ])) + } + + fn derive(&self, children: Vec) -> PhysicalPlan { + assert!(children.is_empty()); + PhysicalPlan::new(self.clone()) + } + + fn build_pipeline2(&self, builder: &mut PipelineBuilder) -> Result<()> { + let table = builder.ctx.build_table_from_source_plan(&self.source)?; + builder.ctx.set_partitions(self.source.parts.clone())?; + + if let Some(prune_pipeline) = table.build_prune_pipeline( + builder.ctx.clone(), + &self.source, + &mut builder.main_pipeline, + self.get_id(), + )? { + builder.pipelines.push(prune_pipeline); + } + + let fuse_table = FuseTable::try_from_table(table.as_ref())?; + fuse_table.do_read_pruned_partitions( + builder.ctx.clone(), + &self.source, + &mut builder.main_pipeline, + )?; + builder.exchange_injector = FusePartExchangeInjector::create(); + Ok(()) + } +} + +impl FusePrune { + pub fn create(source: Box) -> PhysicalPlan { + PhysicalPlan::new(FusePrune { + meta: PhysicalPlanMeta::new("FusePrune"), + source, + }) + } +} diff --git a/src/query/service/src/physical_plans/physical_limit.rs b/src/query/service/src/physical_plans/physical_limit.rs index 43c1a1c3d2a..bc95c87c530 100644 --- a/src/query/service/src/physical_plans/physical_limit.rs +++ b/src/query/service/src/physical_plans/physical_limit.rs @@ -229,7 +229,13 @@ impl PhysicalPlanBuilder { plan: &PhysicalPlan, sources: &mut HashMap, ) { - if let Some(scan) = crate::physical_plans::TableScan::from_physical_plan(plan) { + if let Some(scan) = crate::physical_plans::FuseBlockRead::from_physical_plan(plan) { + if let Some(table_index) = scan.table_index { + sources + .entry(table_index) + .or_insert_with(|| (*scan.source).clone()); + } + } else if let Some(scan) = crate::physical_plans::TableScan::from_physical_plan(plan) { if let Some(table_index) = scan.table_index { sources .entry(table_index) diff --git a/src/query/service/src/physical_plans/physical_plan.rs b/src/query/service/src/physical_plans/physical_plan.rs index 0a69d098afb..4fea8762a14 100644 --- a/src/query/service/src/physical_plans/physical_plan.rs +++ b/src/query/service/src/physical_plans/physical_plan.rs @@ -37,6 +37,7 @@ use serde::Serializer; use crate::physical_plans::EvalScalar; use crate::physical_plans::ExchangeSink; use crate::physical_plans::Filter; +use crate::physical_plans::FuseBlockRead; use crate::physical_plans::MutationSource; use crate::physical_plans::TableScan; use crate::physical_plans::format::FormatContext; @@ -74,11 +75,16 @@ pub trait DeriveHandle: Send + Sync + 'static { } /// The scan a runtime scan filter (TopN boundary / limit early-stop) may -/// attach to: only `Filter`/`EvalScalar` may sit in between; anything else -/// (Sort, Join, Exchange, ...) stops the traversal. Guarded by +/// attach to: a distributed Fuse metadata exchange is part of the scan itself, while only +/// `Filter`/`EvalScalar` may sit above it; anything else (Sort, Join, Exchange, ...) stops the +/// traversal. Guarded by /// `runtime_scan_data_source_only_crosses_row_preserving_wrappers`. #[recursive::recursive] pub fn runtime_scan_data_source(plan: &PhysicalPlan) -> Option<&DataSourcePlan> { + if let Some(scan) = plan.as_any().downcast_ref::() { + return Some(&scan.source); + } + if let Some(scan) = plan.as_any().downcast_ref::() { return Some(&scan.source); } diff --git a/src/query/service/src/physical_plans/physical_plan_builder.rs b/src/query/service/src/physical_plans/physical_plan_builder.rs index dd2177daf6c..8881b79578b 100644 --- a/src/query/service/src/physical_plans/physical_plan_builder.rs +++ b/src/query/service/src/physical_plans/physical_plan_builder.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use databend_common_catalog::plan::PartStatistics; @@ -31,6 +32,7 @@ use databend_storages_common_table_meta::meta::TableMetaTimestamps; use databend_storages_common_table_meta::meta::TableSnapshot; use crate::physical_plans::explain::PlanStatsInfo; +use crate::physical_plans::optimize_distributed_fuse_pruning; use crate::physical_plans::physical_plan::PhysicalPlan; use crate::sessions::TableContext; @@ -44,6 +46,9 @@ pub struct PhysicalPlanBuilder { pub cte_required_columns: HashMap, pub is_cte_required_columns_collected: bool, pub build_depth: usize, + /// Scan IDs that can use the post-pruning block metadata exchange if their finalized + /// physical-plan fragment runs on all executors. + pub distributed_fuse_pruning_scans: HashSet, } impl PhysicalPlanBuilder { @@ -58,6 +63,7 @@ impl PhysicalPlanBuilder { cte_required_columns: HashMap::new(), is_cte_required_columns_collected: false, build_depth: 0, + distributed_fuse_pruning_scans: HashSet::new(), } } @@ -74,6 +80,7 @@ impl PhysicalPlanBuilder { let is_root_build = self.build_depth == 0; if is_root_build { self.ctx.clear_pruned_partitions_stats(); + self.distributed_fuse_pruning_scans.clear(); } if !self.is_cte_required_columns_collected { @@ -89,6 +96,7 @@ impl PhysicalPlanBuilder { let mut plan = build_result?; if is_root_build { + plan = optimize_distributed_fuse_pruning(&plan, &self.distributed_fuse_pruning_scans); plan.adjust_plan_id(&mut 0); self.publish_synchronous_pruning_stats(&plan); } diff --git a/src/query/service/src/physical_plans/physical_table_scan.rs b/src/query/service/src/physical_plans/physical_table_scan.rs index 9bacefbfccc..89de4fb03a5 100644 --- a/src/query/service/src/physical_plans/physical_table_scan.rs +++ b/src/query/service/src/physical_plans/physical_table_scan.rs @@ -25,7 +25,9 @@ use databend_common_catalog::plan::DataSourceInfo; use databend_common_catalog::plan::DataSourcePlan; use databend_common_catalog::plan::Filters; use databend_common_catalog::plan::InternalColumn; +use databend_common_catalog::plan::PartInfoType; use databend_common_catalog::plan::PartStatistics; +use databend_common_catalog::plan::Partitions; use databend_common_catalog::plan::PartitionsShuffleKind; use databend_common_catalog::plan::PrewhereInfo; use databend_common_catalog::plan::Projection; @@ -88,6 +90,20 @@ use crate::sessions::TableContextPartitionStats; use crate::sessions::TableContextSettings; use crate::sessions::TableContextTableFactory; +fn should_use_distributed_block_meta_shuffle( + enable_distributed_pruning: bool, + enable_prune_pipeline: bool, + is_multi_node: bool, + partitions: &Partitions, + is_fuse: bool, +) -> bool { + enable_distributed_pruning + && enable_prune_pipeline + && is_multi_node + && is_fuse + && partitions.partitions_type() == PartInfoType::LazyLevel +} + #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub struct TableScan { pub meta: PhysicalPlanMeta, @@ -207,32 +223,45 @@ impl IPhysicalPlan for TableScan { true, )?; - let schema = self.source.schema(); - // Fill internal columns if needed. - if let Some(internal_columns) = &self.internal_column { - builder.main_pipeline.add_transformer(|| { - TransformAddInternalColumns::new(internal_columns.clone(), schema.clone()) - }); - } + build_scan_output_pipeline( + builder, + &self.source, + &self.name_mapping, + &self.internal_column, + ) + } +} - let mut projection = self - .name_mapping - .keys() - .map(|name| schema.index_of(name.as_str())) - .collect::>>()?; - projection.sort(); - - // if projection is sequential, no need to add projection - if projection != (0..schema.fields().len()).collect::>() { - let ops = vec![BlockOperator::Project { projection }]; - let num_input_columns = schema.num_fields(); - builder.main_pipeline.add_transformer(|| { - CompoundBlockOperator::new(ops.clone(), builder.func_ctx.clone(), num_input_columns) - }); - } +pub(crate) fn build_scan_output_pipeline( + builder: &mut PipelineBuilder, + source: &DataSourcePlan, + name_mapping: &BTreeMap, + internal_column: &Option>, +) -> Result<()> { + let schema = source.schema(); + // Fill internal columns if needed. + if let Some(internal_columns) = internal_column { + builder.main_pipeline.add_transformer(|| { + TransformAddInternalColumns::new(internal_columns.clone(), schema.clone()) + }); + } - Ok(()) + let mut projection = name_mapping + .keys() + .map(|name| schema.index_of(name.as_str())) + .collect::>>()?; + projection.sort(); + + // if projection is sequential, no need to add projection + if projection != (0..schema.fields().len()).collect::>() { + let ops = vec![BlockOperator::Project { projection }]; + let num_input_columns = schema.num_fields(); + builder.main_pipeline.add_transformer(|| { + CompoundBlockOperator::new(ops.clone(), builder.func_ctx.clone(), num_input_columns) + }); } + + Ok(()) } impl TableScan { @@ -530,6 +559,18 @@ impl PhysicalPlanBuilder { metadata.set_table_source(scan.table_index, source.clone()); } + let use_distributed_block_meta_shuffle = should_use_distributed_block_meta_shuffle( + self.ctx.get_settings().get_enable_distributed_pruning()?, + self.ctx.get_settings().get_enable_prune_pipeline()?, + !self.ctx.get_cluster().is_empty(), + &source.parts, + FuseTable::try_from_table(table.as_ref()).is_ok(), + ); + + if use_distributed_block_meta_shuffle { + self.distributed_fuse_pruning_scans.insert(scan.scan_id); + } + let mut plan = TableScan::create( scan.scan_id, name_mapping, diff --git a/src/query/service/src/pipelines/builders/merge_into_join_optimizations.rs b/src/query/service/src/pipelines/builders/merge_into_join_optimizations.rs index 1ae65ab7bb5..0d6248ad6a1 100644 --- a/src/query/service/src/pipelines/builders/merge_into_join_optimizations.rs +++ b/src/query/service/src/pipelines/builders/merge_into_join_optimizations.rs @@ -14,6 +14,7 @@ use databend_common_storages_fuse::operations::need_reserve_block_info; +use crate::physical_plans::FuseBlockRead; use crate::physical_plans::HashJoin; use crate::physical_plans::PhysicalPlanCast; use crate::physical_plans::TableScan; @@ -22,10 +23,15 @@ use crate::pipelines::PipelineBuilder; impl PipelineBuilder { pub(crate) fn merge_into_get_optimization_flag(&self, join: &HashJoin) -> (bool, bool) { // for merge into target table as build side. - if let Some(scan) = TableScan::from_physical_plan(&join.build) { - return match scan.table_index { - None | Some(databend_common_sql::DUMMY_TABLE_INDEX) => (false, false), - Some(table_index) => match need_reserve_block_info(self.ctx.clone(), table_index) { + let table_index = TableScan::from_physical_plan(&join.build) + .and_then(|scan| scan.table_index) + .or_else(|| { + FuseBlockRead::from_physical_plan(&join.build).and_then(|read| read.table_index) + }); + if let Some(table_index) = table_index { + return match table_index { + databend_common_sql::DUMMY_TABLE_INDEX => (false, false), + table_index => match need_reserve_block_info(self.ctx.clone(), table_index) { // due to issue https://github.com/datafuselabs/databend/issues/15643, // target build optimization of merge-into is disabled diff --git a/src/query/service/src/schedulers/fragments/fragmenter.rs b/src/query/service/src/schedulers/fragments/fragmenter.rs index 872d5ffc42e..1d176020403 100644 --- a/src/query/service/src/schedulers/fragments/fragmenter.rs +++ b/src/query/service/src/schedulers/fragments/fragmenter.rs @@ -31,6 +31,7 @@ use crate::physical_plans::DeriveHandle; use crate::physical_plans::Exchange; use crate::physical_plans::ExchangeSink; use crate::physical_plans::ExchangeSource; +use crate::physical_plans::FusePrune; use crate::physical_plans::IPhysicalPlan; use crate::physical_plans::MaterializedCTE; use crate::physical_plans::MutationSource; @@ -396,6 +397,10 @@ impl PhysicalPlanVisitor for FragmentTypeVisitor { self.fragment_type = FragmentType::Source; } + if FusePrune::check_physical_plan(v) { + self.fragment_type = FragmentType::Source; + } + if ConstantTableScan::check_physical_plan(v) { self.fragment_type = FragmentType::Source; } diff --git a/src/query/service/src/schedulers/fragments/plan_fragment.rs b/src/query/service/src/schedulers/fragments/plan_fragment.rs index 8d21462fc38..1502d5b1589 100644 --- a/src/query/service/src/schedulers/fragments/plan_fragment.rs +++ b/src/query/service/src/schedulers/fragments/plan_fragment.rs @@ -33,6 +33,7 @@ use crate::physical_plans::CompactSource; use crate::physical_plans::ConstantTableScan; use crate::physical_plans::DeriveHandle; use crate::physical_plans::ExchangeSink; +use crate::physical_plans::FusePrune; use crate::physical_plans::IPhysicalPlan; use crate::physical_plans::MutationSource; use crate::physical_plans::PhysicalPlan; @@ -62,8 +63,8 @@ pub enum FragmentType { /// doesn't contain any `TableScan` operator. Intermediate, - /// Leaf fragment of a query plan, which contains - /// a `TableScan` operator. + /// Leaf fragment of a query plan, which contains a row-producing scan or a metadata-producing + /// `FusePrune` source operator. Source, /// Intermediate fragment of a replace into plan, which contains a `ReplaceInto` operator. ReplaceInto, @@ -551,6 +552,9 @@ impl PlanFragment { if let Some(scan) = TableScan::from_physical_plan(plan) { self.data_sources .insert(plan.get_id(), DataSource::Table(*scan.source.clone())); + } else if let Some(prune) = FusePrune::from_physical_plan(plan) { + self.data_sources + .insert(plan.get_id(), DataSource::Table(*prune.source.clone())); } else if let Some(scan) = ConstantTableScan::from_physical_plan(plan) { self.data_sources.insert( plan.get_id(), @@ -651,6 +655,22 @@ impl DeriveHandle for ReadSourceDeriveHandle { source: Box::new(source), ..table_scan.clone() })); + } else if let Some(fuse_prune) = FusePrune::from_physical_plan(v) { + let Some(source) = self.sources.remove(&fuse_prune.get_id()) else { + unreachable!( + "Cannot find data source for Fuse prune plan {}", + fuse_prune.get_id() + ) + }; + + let Ok(source) = DataSourcePlan::try_from(source) else { + unreachable!("Cannot create data source plan"); + }; + + return Ok(PhysicalPlan::new(FusePrune { + source: Box::new(source), + ..fuse_prune.clone() + })); } else if let Some(table_scan) = ConstantTableScan::from_physical_plan(v) { let Some(source) = self.sources.remove(&table_scan.get_id()) else { unreachable!( diff --git a/src/query/service/src/servers/flight/v1/exchange/fuse_part_exchange.rs b/src/query/service/src/servers/flight/v1/exchange/fuse_part_exchange.rs new file mode 100644 index 00000000000..928c8334350 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/exchange/fuse_part_exchange.rs @@ -0,0 +1,392 @@ +// 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::HashMap; +use std::sync::Arc; + +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::BlockMetaInfoDowncast; +use databend_common_expression::DataBlock; +use databend_common_pipeline::core::Pipeline; +use databend_common_settings::FlightCompression; +use databend_common_storages_fuse::FuseBlockPartInfo; +use databend_common_storages_fuse::operations::BlockPartitionMeta; + +use super::DataExchange; +use super::DefaultExchangeInjector; +use super::ExchangeInjector; +use super::ExchangeSorting; +use super::MergeExchangeParams; +use super::ShuffleExchangeParams; +use crate::clusters::ClusterHelper; +use crate::servers::flight::v1::scatter::FlightScatter; +use crate::sessions::QueryContext; +use crate::sessions::TableContextCluster; + +pub struct FusePartExchangeInjector { + default_injector: Arc, +} + +impl FusePartExchangeInjector { + pub fn create() -> Arc { + Arc::new(Self { + default_injector: DefaultExchangeInjector::create(), + }) + } +} + +fn build_bucket_to_output( + destination_ids: &[String], + cache_ids: &HashMap, +) -> Result> { + let mut outputs_by_cache_id = destination_ids + .iter() + .enumerate() + .map(|(output, id)| { + cache_ids + .get(id) + .map(|cache_id| (cache_id.clone(), id.clone(), output)) + .ok_or_else(|| { + ErrorCode::Internal(format!( + "Cannot find cache id for exchange destination {id}" + )) + }) + }) + .collect::>>()?; + outputs_by_cache_id + .sort_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))); + + Ok(outputs_by_cache_id + .into_iter() + .map(|(_, _, output)| output) + .collect()) +} + +impl ExchangeInjector for FusePartExchangeInjector { + fn flight_scatter( + &self, + ctx: &Arc, + exchange: &DataExchange, + ) -> Result>> { + let DataExchange::NodeToNodeExchange(exchange) = exchange else { + return Err(ErrorCode::Internal( + "Fuse block partition exchange requires a node-to-node exchange", + )); + }; + + let cache_ids = ctx + .get_cluster() + .get_nodes() + .iter() + .map(|node| (node.id.clone(), node.cache_id.clone())) + .collect::>(); + let bucket_to_output = build_bucket_to_output(&exchange.destination_ids, &cache_ids)?; + Ok(Arc::new(Box::new(FusePartFlightScatter { + bucket_to_output, + }))) + } + + fn exchange_sorting(&self) -> Option> { + self.default_injector.exchange_sorting() + } + + fn apply_merge_serializer( + &self, + params: &MergeExchangeParams, + compression: Option, + pipeline: &mut Pipeline, + ) -> Result<()> { + self.default_injector + .apply_merge_serializer(params, compression, pipeline) + } + + fn apply_shuffle_serializer( + &self, + params: &ShuffleExchangeParams, + compression: Option, + pipeline: &mut Pipeline, + ) -> Result<()> { + self.default_injector + .apply_shuffle_serializer(params, compression, pipeline) + } + + fn apply_merge_deserializer( + &self, + params: &MergeExchangeParams, + pipeline: &mut Pipeline, + ) -> Result<()> { + self.default_injector + .apply_merge_deserializer(params, pipeline) + } + + fn apply_shuffle_deserializer( + &self, + params: &ShuffleExchangeParams, + pipeline: &mut Pipeline, + ) -> Result<()> { + self.default_injector + .apply_shuffle_deserializer(params, pipeline) + } +} + +struct FusePartFlightScatter { + /// Hash buckets are ordered by persistent cache id, while exchange outputs use destination id + /// order. This map translates a stable hash bucket into the corresponding output position. + bucket_to_output: Vec, +} + +impl FlightScatter for FusePartFlightScatter { + fn name(&self) -> &'static str { + "FusePartFlightScatter" + } + + fn execute(&self, mut data_block: DataBlock) -> Result> { + if !data_block.is_empty() { + return Err(ErrorCode::Internal( + "Fuse block partition exchange received a non-empty data block", + )); + } + if self.bucket_to_output.is_empty() { + return Err(ErrorCode::Internal( + "Fuse block partition exchange has no destination", + )); + } + + let meta = data_block.take_meta().ok_or_else(|| { + ErrorCode::Internal("Fuse block partition exchange received data without metadata") + })?; + let meta = BlockPartitionMeta::downcast_from(meta).ok_or_else(|| { + ErrorCode::Internal("Fuse block partition exchange received unexpected metadata") + })?; + + let mut partitions = vec![Vec::new(); self.bucket_to_output.len()]; + for part in meta.part_ptr { + FuseBlockPartInfo::from_part(&part)?; + let bucket = (part.hash() % self.bucket_to_output.len() as u64) as usize; + partitions[self.bucket_to_output[bucket]].push(part); + } + + Ok(partitions + .into_iter() + .map(|parts| match parts.is_empty() { + true => DataBlock::empty(), + false => DataBlock::empty_with_meta(BlockPartitionMeta::create(parts)), + }) + .collect()) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use databend_common_catalog::plan::PartInfoPtr; + use databend_common_expression::BlockMetaInfoDowncast; + use databend_common_expression::BlockMetaInfoPtr; + use databend_common_expression::Scalar; + use databend_common_io::prelude::bincode_deserialize_from_slice; + use databend_common_io::prelude::bincode_serialize_into_buf; + use databend_common_storages_fuse::FuseLazyPartInfo; + use databend_storages_common_pruner::BlockMetaIndex; + use databend_storages_common_pruner::VirtualBlockMetaIndex; + use databend_storages_common_table_meta::meta::ColumnMeta; + use databend_storages_common_table_meta::meta::ColumnStatistics; + use databend_storages_common_table_meta::meta::Compression; + use databend_storages_common_table_meta::meta::SingleColumnMeta; + + use super::*; + + fn part(location: &str, block_idx: usize) -> PartInfoPtr { + FuseBlockPartInfo::create( + location.to_string(), + Some((format!("{location}.bloom"), 1)), + 32, + 10, + HashMap::from([(1, ColumnMeta::Parquet(SingleColumnMeta::new(11, 22, 10)))]), + Some(HashMap::from([( + 1, + ColumnStatistics::new(Scalar::from(1_i64), Scalar::from(9_i64), 0, 80, Some(9)), + )])), + Compression::Lz4Raw, + Some((Scalar::from(1_i64), Scalar::from(9_i64))), + Some(BlockMetaIndex { + block_idx, + block_id: block_idx + 100, + block_location: location.to_string(), + segment_location: "segment-1".to_string(), + snapshot_location: Some("snapshot-1".to_string()), + range: Some(2..8), + virtual_block_meta: Some(VirtualBlockMetaIndex { + virtual_block_location: format!("{location}.virtual"), + ..Default::default() + }), + ..Default::default() + }), + None, + ) + } + + fn take_parts(mut block: DataBlock) -> Vec { + let Some(meta) = block.take_meta() else { + return vec![]; + }; + BlockPartitionMeta::downcast_from(meta).unwrap().part_ptr + } + + #[test] + fn test_fuse_part_scatter_routes_every_part_once() -> Result<()> { + let bucket_to_output = vec![2, 0, 3, 1]; + let scatter = FusePartFlightScatter { + bucket_to_output: bucket_to_output.clone(), + }; + let parts = (0..100) + .map(|index| part(&format!("block-{index}"), index)) + .collect::>(); + let expected = parts + .iter() + .map(|part| { + let bucket = (part.hash() % bucket_to_output.len() as u64) as usize; + ( + FuseBlockPartInfo::from_part(part).unwrap().location.clone(), + bucket_to_output[bucket], + ) + }) + .collect::>(); + + let outputs = scatter.execute(DataBlock::empty_with_meta(BlockPartitionMeta::create( + parts, + )))?; + assert_eq!(outputs.len(), bucket_to_output.len()); + + let mut seen = HashSet::new(); + for (output, block) in outputs.into_iter().enumerate() { + for part in take_parts(block) { + let part = FuseBlockPartInfo::from_part(&part)?; + assert_eq!(expected[&part.location], output); + assert!(seen.insert(part.location.clone())); + assert_eq!(part.block_meta_index.as_ref().unwrap().range, Some(2..8)); + } + } + assert_eq!(seen.len(), expected.len()); + Ok(()) + } + + #[test] + fn test_cache_id_order_is_independent_of_exchange_order() -> Result<()> { + let cache_ids = HashMap::from([ + ("node-a".to_string(), "cache-2".to_string()), + ("node-b".to_string(), "cache-1".to_string()), + ("node-c".to_string(), "cache-3".to_string()), + ]); + let first = vec![ + "node-a".to_string(), + "node-b".to_string(), + "node-c".to_string(), + ]; + let second = vec![ + "node-c".to_string(), + "node-a".to_string(), + "node-b".to_string(), + ]; + + let first_nodes = build_bucket_to_output(&first, &cache_ids)? + .into_iter() + .map(|output| first[output].clone()) + .collect::>(); + let second_nodes = build_bucket_to_output(&second, &cache_ids)? + .into_iter() + .map(|output| second[output].clone()) + .collect::>(); + + assert_eq!(first_nodes, vec!["node-b", "node-a", "node-c"]); + assert_eq!(first_nodes, second_nodes); + Ok(()) + } + + #[test] + fn test_fuse_part_scatter_does_not_emit_meta_for_empty_destinations() -> Result<()> { + let scatter = FusePartFlightScatter { + bucket_to_output: vec![0, 1, 2, 3], + }; + let outputs = scatter.execute(DataBlock::empty_with_meta(BlockPartitionMeta::create( + vec![part("only-block", 0)], + )))?; + + assert_eq!(outputs.len(), 4); + assert_eq!( + outputs + .iter() + .filter(|block| block.get_meta().is_some()) + .count(), + 1 + ); + Ok(()) + } + + #[test] + fn test_fuse_part_scatter_rejects_missing_meta_and_non_block_parts() { + let scatter = FusePartFlightScatter { + bucket_to_output: vec![0, 1], + }; + assert!(scatter.execute(DataBlock::empty()).is_err()); + + let lazy_part = FuseLazyPartInfo::create(0, ("segment-0".to_string(), 1)); + let block = DataBlock::empty_with_meta(BlockPartitionMeta::create(vec![lazy_part])); + assert!(scatter.execute(block).is_err()); + } + + #[test] + fn test_block_partition_meta_bincode_round_trip() -> Result<()> { + let meta: Option = Some(BlockPartitionMeta::create(vec![part( + "block-round-trip", + 7, + )])); + let mut encoded = Vec::new(); + bincode_serialize_into_buf(&mut encoded, &meta)?; + let decoded: Option = bincode_deserialize_from_slice(&encoded)?; + let decoded = BlockPartitionMeta::downcast_from(decoded.unwrap()).unwrap(); + let part = FuseBlockPartInfo::from_part(&decoded.part_ptr[0])?; + + assert_eq!(part.location, "block-round-trip"); + assert_eq!( + part.bloom_filter_index_location.as_ref().unwrap().0, + "block-round-trip.bloom" + ); + assert_eq!(part.bloom_filter_index_size, 32); + assert_eq!(part.columns_meta[&1].offset_length(), (11, 22)); + assert_eq!( + part.columns_stat.as_ref().unwrap()[&1].min(), + &Scalar::from(1_i64) + ); + assert_eq!( + part.sort_min_max.clone(), + Some((Scalar::from(1_i64), Scalar::from(9_i64))) + ); + assert_eq!(part.block_meta_index.as_ref().unwrap(), &BlockMetaIndex { + block_idx: 7, + block_id: 107, + block_location: "block-round-trip".to_string(), + segment_location: "segment-1".to_string(), + snapshot_location: Some("snapshot-1".to_string()), + range: Some(2..8), + virtual_block_meta: Some(VirtualBlockMetaIndex { + virtual_block_location: "block-round-trip.virtual".to_string(), + ..Default::default() + }), + ..Default::default() + }); + Ok(()) + } +} diff --git a/src/query/service/src/servers/flight/v1/exchange/mod.rs b/src/query/service/src/servers/flight/v1/exchange/mod.rs index e79e78ffda2..4a194502a52 100644 --- a/src/query/service/src/servers/flight/v1/exchange/mod.rs +++ b/src/query/service/src/servers/flight/v1/exchange/mod.rs @@ -26,6 +26,7 @@ mod exchange_source_reader; mod exchange_transform; mod exchange_transform_scatter; mod exchange_transform_shuffle; +mod fuse_part_exchange; mod hash_send_sink; mod hash_send_source; mod hash_send_transform; @@ -52,6 +53,7 @@ pub use exchange_sorting::ExchangeSorting; pub use exchange_transform_scatter::ScatterTransform; pub use exchange_transform_shuffle::ExchangeShuffleMeta; pub use exchange_transform_shuffle::ExchangeShuffleTransform; +pub use fuse_part_exchange::FusePartExchangeInjector; pub use hash_send_sink::HashSendSink; pub use hash_send_source::HashSendSource; pub use hash_send_transform::HashSendTransform; diff --git a/src/query/storages/fuse/src/operations/mod.rs b/src/query/storages/fuse/src/operations/mod.rs index b4ea947fd7c..66b53f4e805 100644 --- a/src/query/storages/fuse/src/operations/mod.rs +++ b/src/query/storages/fuse/src/operations/mod.rs @@ -52,6 +52,7 @@ pub use compact::CompactOptions; pub use merge_into::*; pub use mutation::*; pub use mutation_source::*; +pub use read::BlockPartitionMeta; pub use read::DeserializeDataTransform; pub use read::ReadState; pub use read::need_reserve_block_info; diff --git a/src/query/storages/fuse/src/operations/read/block_partition_meta.rs b/src/query/storages/fuse/src/operations/read/block_partition_meta.rs index 91269afb54e..43575702ea0 100644 --- a/src/query/storages/fuse/src/operations/read/block_partition_meta.rs +++ b/src/query/storages/fuse/src/operations/read/block_partition_meta.rs @@ -18,8 +18,8 @@ use std::fmt::Formatter; use databend_common_catalog::plan::PartInfoPtr; use databend_common_expression::BlockMetaInfo; use databend_common_expression::BlockMetaInfoPtr; -use databend_common_expression::local_block_meta_serde; +#[derive(serde::Serialize, serde::Deserialize)] pub struct BlockPartitionMeta { pub part_ptr: Vec, } @@ -38,7 +38,5 @@ impl Debug for BlockPartitionMeta { } } -local_block_meta_serde!(BlockPartitionMeta); - #[typetag::serde(name = "block_partition_meta")] impl BlockMetaInfo for BlockPartitionMeta {} diff --git a/src/query/storages/fuse/src/operations/read/fuse_source.rs b/src/query/storages/fuse/src/operations/read/fuse_source.rs index fdaf6662734..2b83facfec8 100644 --- a/src/query/storages/fuse/src/operations/read/fuse_source.rs +++ b/src/query/storages/fuse/src/operations/read/fuse_source.rs @@ -59,10 +59,43 @@ pub fn build_fuse_source_pipeline( ) -> Result<()> { (max_threads, max_io_requests) = adjust_threads_and_request(max_threads, max_io_requests, plan); - let waker = pipeline.get_waker(); let batch_size = ctx.get_settings().get_storage_fetch_part_num()? as usize; + let source_batch_size = if receiver.is_some() { 1 } else { batch_size }; + build_fuse_partitions_source_pipeline( + ctx.clone(), + pipeline, + plan, + max_io_requests, + receiver, + source_batch_size, + )?; + + build_fuse_read_transform_pipeline( + ctx, + storage_format, + table_schema, + pipeline, + block_reader, + max_threads, + max_io_requests, + plan, + index_reader, + virtual_reader, + false, + ) +} + +pub(crate) fn build_fuse_partitions_source_pipeline( + ctx: Arc, + pipeline: &mut Pipeline, + plan: &DataSourcePlan, + max_io_requests: usize, + receiver: Option>>, + batch_size: usize, +) -> Result<()> { + let waker = pipeline.get_waker(); let stream: Arc = match receiver { - Some(rx) => Arc::new(ReceiverPartitionStream::new(rx)), + Some(rx) => Arc::new(ReceiverPartitionStream::with_batch_size(rx, batch_size)), None => { let partitions = dispatch_partitions(ctx.clone(), plan, max_io_requests); let partitions = StealablePartitions::new(partitions, ctx.clone()); @@ -88,6 +121,25 @@ pub fn build_fuse_source_pipeline( } pipeline.add_pipe(source_builder.finalize()); + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn build_fuse_read_transform_pipeline( + ctx: Arc, + storage_format: FuseStorageFormat, + table_schema: Arc, + pipeline: &mut Pipeline, + block_reader: Arc, + max_threads: usize, + max_io_requests: usize, + plan: &DataSourcePlan, + index_reader: Arc>, + virtual_reader: Arc>, + record_partitions: bool, +) -> Result<()> { + pipeline.try_resize(max_io_requests)?; + let block_format = match storage_format { FuseStorageFormat::Parquet => FuseParquetBlockFormat::create(), FuseStorageFormat::Unsupported => { @@ -111,6 +163,7 @@ pub fn build_fuse_source_pipeline( table_schema.clone(), block_reader.clone(), read_block_context.clone(), + record_partitions, input, output, ) diff --git a/src/query/storages/fuse/src/operations/read/mod.rs b/src/query/storages/fuse/src/operations/read/mod.rs index 83234291a0b..a16b79a7098 100644 --- a/src/query/storages/fuse/src/operations/read/mod.rs +++ b/src/query/storages/fuse/src/operations/read/mod.rs @@ -28,6 +28,7 @@ mod data_source_with_meta; mod partition_stream; mod util; +pub use block_partition_meta::BlockPartitionMeta; pub use fuse_rows_fetcher::row_fetch_processor; pub use fuse_source::build_fuse_source_pipeline; pub use parquet_data_source_deserializer::DeserializeDataTransform; diff --git a/src/query/storages/fuse/src/operations/read/partition_stream.rs b/src/query/storages/fuse/src/operations/read/partition_stream.rs index d56fe3b514a..eb3b6664a34 100644 --- a/src/query/storages/fuse/src/operations/read/partition_stream.rs +++ b/src/query/storages/fuse/src/operations/read/partition_stream.rs @@ -65,22 +65,38 @@ impl PartitionStream for StealPartitionStream { pub struct ReceiverPartitionStream { receiver: Receiver>, + max_batch_size: usize, } impl ReceiverPartitionStream { - pub fn new(receiver: Receiver>) -> Self { - Self { receiver } + pub fn with_batch_size(receiver: Receiver>, max_batch_size: usize) -> Self { + Self { + receiver, + max_batch_size: max_batch_size.max(1), + } } } #[async_trait::async_trait] impl PartitionStream for ReceiverPartitionStream { async fn fetch(&self, _id: usize) -> Result>> { - match self.receiver.recv().await { - Ok(Ok(part)) => Ok(Some(vec![part])), - Ok(Err(e)) => Err(e), - Err(_) => Ok(None), + let first = match self.receiver.recv().await { + Ok(Ok(part)) => part, + Ok(Err(e)) => return Err(e), + Err(_) => return Ok(None), + }; + + let mut parts = Vec::with_capacity(self.max_batch_size); + parts.push(first); + while parts.len() < self.max_batch_size { + match self.receiver.try_recv() { + Ok(Ok(part)) => parts.push(part), + Ok(Err(e)) => return Err(e), + Err(_) => break, + } } + + Ok(Some(parts)) } } @@ -230,3 +246,53 @@ impl Processor for PartitionStreamSource { self.id = id; } } + +#[cfg(test)] +mod tests { + use databend_common_exception::ErrorCode; + + use super::*; + use crate::FuseLazyPartInfo; + + #[tokio::test] + async fn test_receiver_partition_stream_batches_parts() -> Result<()> { + let (tx, rx) = async_channel::unbounded(); + for index in 0..3 { + tx.send(Ok(FuseLazyPartInfo::create( + index, + (format!("segment-{index}"), 1), + ))) + .await + .unwrap(); + } + drop(tx); + + let stream = ReceiverPartitionStream::with_batch_size(rx, 2); + assert_eq!(stream.fetch(0).await?.unwrap().len(), 2); + assert_eq!(stream.fetch(0).await?.unwrap().len(), 1); + assert!(stream.fetch(0).await?.is_none()); + Ok(()) + } + + #[tokio::test] + async fn test_receiver_partition_stream_forwards_errors() { + let (tx, rx) = async_channel::unbounded(); + tx.send(Err(ErrorCode::Internal("prune failed"))) + .await + .unwrap(); + drop(tx); + + let stream = ReceiverPartitionStream::with_batch_size(rx, 2); + assert!(stream.fetch(0).await.is_err()); + } + + #[tokio::test] + async fn test_receiver_partition_stream_empty_channel_finishes() -> Result<()> { + let (tx, rx) = async_channel::unbounded(); + drop(tx); + + let stream = ReceiverPartitionStream::with_batch_size(rx, 8); + assert!(stream.fetch(0).await?.is_none()); + Ok(()) + } +} diff --git a/src/query/storages/fuse/src/operations/read/read_data_transform.rs b/src/query/storages/fuse/src/operations/read/read_data_transform.rs index a812506815d..9fe84576da0 100644 --- a/src/query/storages/fuse/src/operations/read/read_data_transform.rs +++ b/src/query/storages/fuse/src/operations/read/read_data_transform.rs @@ -14,6 +14,8 @@ use std::sync::Arc; +use databend_common_base::runtime::profile::Profile; +use databend_common_base::runtime::profile::ProfileStatisticsName; use databend_common_catalog::plan::PartInfoPtr; use databend_common_catalog::runtime_filter_info::RuntimeScanFilters; use databend_common_catalog::table_context::TableContext; @@ -46,6 +48,7 @@ pub struct ReadDataTransform { scan_id: IndexType, context: Arc, runtime_scan_filters: RuntimeScanFilters, + record_partitions: bool, } impl ReadDataTransform { @@ -56,6 +59,7 @@ impl ReadDataTransform { table_schema: Arc, block_reader: Arc, read_block_context: Arc, + record_partitions: bool, input: Arc, output: Arc, ) -> Result { @@ -72,6 +76,7 @@ impl ReadDataTransform { scan_id, context: ctx, runtime_scan_filters, + record_partitions, }, ))) } @@ -177,6 +182,10 @@ impl AsyncTransform for ReadDataTransform { .and_then(|meta| (!meta.part_ptr.is_empty()).then(|| meta.part_ptr.clone())) .ok_or_else(|| ErrorCode::Internal("AsyncReadDataTransform got wrong meta data"))?; + if self.record_partitions { + Profile::record_usize_profile(ProfileStatisticsName::ScanPartitions, parts.len()); + } + self.read_parts(parts).await } } diff --git a/src/query/storages/fuse/src/operations/read_data.rs b/src/query/storages/fuse/src/operations/read_data.rs index 875a985245f..d4f886a3286 100644 --- a/src/query/storages/fuse/src/operations/read_data.rs +++ b/src/query/storages/fuse/src/operations/read_data.rs @@ -22,8 +22,10 @@ use databend_common_catalog::plan::PushDownInfo; use databend_common_catalog::plan::ReadPartitionsPruningMode; use databend_common_catalog::table::Table; use databend_common_catalog::table_context::TableContext; +use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_pipeline::core::Pipeline; +use databend_common_pipeline::sources::EmptySource; use crate::FuseLazyPartInfo; use crate::FuseTable; @@ -32,6 +34,14 @@ use crate::io::AggIndexReader; use crate::io::BlockReader; use crate::io::VirtualColumnReader; use crate::operations::read::build_fuse_source_pipeline; +use crate::operations::read::fuse_source::build_fuse_partitions_source_pipeline; +use crate::operations::read::fuse_source::build_fuse_read_transform_pipeline; + +type FuseDataReaders = ( + Arc, + Arc>, + Arc>, +); impl FuseTable { pub fn create_block_reader( @@ -74,30 +84,13 @@ impl FuseTable { Ok(std::cmp::max(max_threads, max_io_requests)) } - #[inline] - pub fn do_read_data( + fn build_data_readers( &self, ctx: Arc, plan: &DataSourcePlan, - pipeline: &mut Pipeline, put_cache: bool, - ) -> Result<()> { - let snapshot_loc = plan.statistics.snapshot.clone(); - let mut lazy_init_segments = Vec::with_capacity(plan.parts.len()); - - for part in &plan.parts.partitions { - if let Some(lazy_part_info) = part.as_any().downcast_ref::() { - lazy_init_segments.push(SegmentLocation { - segment_idx: lazy_part_info.segment_index, - location: lazy_part_info.segment_location.clone(), - snapshot_loc: snapshot_loc.clone(), - }); - } - } - + ) -> Result { let block_reader = self.build_block_reader(ctx.clone(), plan, put_cache)?; - let max_io_requests = self.adjust_io_request(&ctx)?; - let index_reader = Arc::new( plan.push_downs .as_ref() @@ -113,13 +106,12 @@ impl FuseTable { }) .transpose()?, ); - let virtual_reader = Arc::new( PushDownInfo::virtual_columns_of_push_downs(&plan.push_downs) .as_ref() .map(|virtual_column| { VirtualColumnReader::try_create( - ctx.clone(), + ctx, self.operator.clone(), block_reader.schema(), plan, @@ -130,6 +122,86 @@ impl FuseTable { .transpose()?, ); + Ok((block_reader, index_reader, virtual_reader)) + } + + /// Build a source that only emits the block partitions produced by the pruning pipeline. + pub fn do_read_pruned_partitions( + &self, + ctx: Arc, + plan: &DataSourcePlan, + pipeline: &mut Pipeline, + ) -> Result<()> { + self.check_format_supported()?; + let Some(receiver) = self.pruned_result_receiver.lock().take() else { + return match plan.parts.is_empty() { + true => pipeline.add_source(EmptySource::create, 1), + false => Err(ErrorCode::Internal( + "Distributed Fuse prune operator did not produce a partition receiver", + )), + }; + }; + + let batch_size = ctx.get_settings().get_storage_fetch_part_num()? as usize; + // Metadata delivery is cheap and a single receiver preserves useful batching. Block reads + // are expanded back to the configured IO parallelism after the exchange. + build_fuse_partitions_source_pipeline(ctx, pipeline, plan, 1, Some(receiver), batch_size) + } + + /// Attach block reading to a pipeline that already emits `BlockPartitionMeta`. + pub fn do_read_data_from_partitions( + &self, + ctx: Arc, + plan: &DataSourcePlan, + pipeline: &mut Pipeline, + put_cache: bool, + ) -> Result<()> { + self.check_format_supported()?; + let (block_reader, index_reader, virtual_reader) = + self.build_data_readers(ctx.clone(), plan, put_cache)?; + let max_threads = ctx.get_settings().get_max_threads()? as usize; + let max_io_requests = self.adjust_io_request(&ctx)?; + + build_fuse_read_transform_pipeline( + ctx, + self.storage_format, + self.schema_with_stream(), + pipeline, + block_reader, + max_threads, + max_io_requests, + plan, + index_reader, + virtual_reader, + true, + ) + } + + #[inline] + pub fn do_read_data( + &self, + ctx: Arc, + plan: &DataSourcePlan, + pipeline: &mut Pipeline, + put_cache: bool, + ) -> Result<()> { + let snapshot_loc = plan.statistics.snapshot.clone(); + let mut lazy_init_segments = Vec::with_capacity(plan.parts.len()); + + for part in &plan.parts.partitions { + if let Some(lazy_part_info) = part.as_any().downcast_ref::() { + lazy_init_segments.push(SegmentLocation { + segment_idx: lazy_part_info.segment_index, + location: lazy_part_info.segment_location.clone(), + snapshot_loc: snapshot_loc.clone(), + }); + } + } + + let (block_reader, index_reader, virtual_reader) = + self.build_data_readers(ctx.clone(), plan, put_cache)?; + let max_io_requests = self.adjust_io_request(&ctx)?; + let enable_prune_pipeline = ctx.get_settings().get_enable_prune_pipeline()?; let rx = if !enable_prune_pipeline && !lazy_init_segments.is_empty() { // If the prune pipeline is disabled and is lazy init segments, we need to fallback diff --git a/tests/sqllogictests/suites/mode/cluster/explain_analyze.test b/tests/sqllogictests/suites/mode/cluster/explain_analyze.test index 6b67d3aec0b..ba78651c57c 100644 --- a/tests/sqllogictests/suites/mode/cluster/explain_analyze.test +++ b/tests/sqllogictests/suites/mode/cluster/explain_analyze.test @@ -64,6 +64,11 @@ drop table if exists article; statement ok drop table if exists author; +# Keep this legacy EXPLAIN ANALYZE profile focused on scan metrics. The distributed metadata +# exchange has its own plan and cluster regression coverage. +statement ok +set enable_distributed_pruning = 0; + statement ok create table if not exists article (article_id int, author_id int, viewer_id int, view_date date); @@ -245,3 +250,6 @@ drop table if exists article; statement ok drop table if exists author; + +statement ok +set enable_distributed_pruning = 1; diff --git a/tests/sqllogictests/suites/mode/cluster/fuse_pruned_block_meta_shuffle.test b/tests/sqllogictests/suites/mode/cluster/fuse_pruned_block_meta_shuffle.test new file mode 100644 index 00000000000..6e0081f3775 --- /dev/null +++ b/tests/sqllogictests/suites/mode/cluster/fuse_pruned_block_meta_shuffle.test @@ -0,0 +1,129 @@ +statement ok +create or replace database fuse_pruned_block_meta_shuffle; + +statement ok +use fuse_pruned_block_meta_shuffle; + +statement ok +set enable_prune_pipeline = 1; + +statement ok +set enable_distributed_pruning = 1; + +# One block per segment makes the lazy segment count comfortably larger than the cluster size. +statement ok +create or replace table target (id int, value int) row_per_block = 10 block_per_segment = 1; + +statement ok +insert into target select number, number * 2 from numbers(200); + +query II +select count(*), sum(value) from target where id between 37 and 83; +---- +47 5640 + +query II rowsort +select id, value from target where id % 31 = 0 order by id limit 4; +---- +0 0 +31 62 +62 124 +93 186 + +query I +select count(*) from target where id < 0; +---- +0 + +query II rowsort +select id % 4, count(*) from target where id < 40 group by id % 4; +---- +0 10 +1 10 +2 10 +3 10 + +# Disabling distributed pruning keeps the original eager/local scan path. +statement ok +set enable_distributed_pruning = 0; + +query II +select count(*), sum(value) from target where id between 37 and 83; +---- +47 5640 + +statement ok +set enable_distributed_pruning = 1; + +statement ok +set enable_prune_pipeline = 0; + +query II +select count(*), sum(value) from target where id between 37 and 83; +---- +47 5640 + +statement ok +set enable_prune_pipeline = 1; + +# Force one large, many-block segment plus several tiny segments. Static segment/block pruning +# leaves the large segment's owner with most metadata before the post-prune shuffle. +statement ok +set max_threads = 1; + +statement ok +create or replace table merge_target (id int, value int) +row_per_block = 10 block_per_segment = 1000; + +statement ok +insert into merge_target select number, number * 2 from numbers(1000); + +statement ok +insert into merge_target select number + 10000, number from numbers(10); + +statement ok +insert into merge_target select number + 11000, number from numbers(10); + +statement ok +insert into merge_target select number + 12000, number from numbers(10); + +statement ok +insert into merge_target select number + 13000, number from numbers(10); + +statement ok +set max_threads = 4; + +statement ok +create or replace table merge_source (id int, value int); + +statement ok +insert into merge_source select number, number + 20000 from numbers(500); + +statement ok +insert into merge_source values (5000, 25000); + +query II +merge into merge_target +using merge_source +on merge_target.id = merge_source.id +when matched then update set merge_target.value = merge_source.value +when not matched then insert (id, value) values (merge_source.id, merge_source.value); +---- +1 500 + +query I +select count(*) from merge_target; +---- +1041 + +query II +select id, value from merge_target where id in (0, 499, 500, 5000, 10000) order by id; +---- +0 20000 +499 20499 +500 1000 +5000 25000 +10000 0 + +statement ok +drop database fuse_pruned_block_meta_shuffle; diff --git a/tests/suites/1_stateful/02_query/02_0011_explain_analyze_part_info.py b/tests/suites/1_stateful/02_query/02_0011_explain_analyze_part_info.py index b14b807decf..7465db8d305 100755 --- a/tests/suites/1_stateful/02_query/02_0011_explain_analyze_part_info.py +++ b/tests/suites/1_stateful/02_query/02_0011_explain_analyze_part_info.py @@ -35,12 +35,14 @@ mycursor.execute(f"insert into test_explain_analyze_0011_1 values ({i})") def explain_output(res): - cnt = 0 pruning_fields = ["partitions total", "partitions scanned", "pruning stats"] + # Distributed pruning can report a field from both its prune and read operators. + found_fields = set() for row in res: - if any(field in row[0] for field in pruning_fields): - cnt += 1 - print(cnt) + for field in pruning_fields: + if field in row[0]: + found_fields.add(field) + print(len(found_fields)) mycursor.execute( """EXPLAIN ANALYZE SELECT * FROM test_explain_analyze_0011_1 a WHERE a.id > 5"""