diff --git a/src/common/base/src/runtime/profile/profiles.rs b/src/common/base/src/runtime/profile/profiles.rs index ed2a1e4e8828f..d1a80985f751e 100644 --- a/src/common/base/src/runtime/profile/profiles.rs +++ b/src/common/base/src/runtime/profile/profiles.rs @@ -41,6 +41,13 @@ pub enum ProfileStatisticsName { ScanBytesFromDataCache, ScanBytesFromMemory, + RowFetchRows, + RowFetchBlocks, + RowFetchEstimatedBytes, + RowFetchBatches, + RowFetchLocalBatches, + RowFetchDistributedBatches, + RemoteSpillWriteCount, RemoteSpillWriteBytes, RemoteSpillWriteTime, @@ -65,6 +72,9 @@ pub enum ProfileStatisticsName { MemoryUsage, ExternalServerRetryCount, ExternalServerRequestCount, + + RowFetchInputBatches, + RowFetchAffinityReassignedBlocks, } #[derive(Clone, Hash, Eq, PartialEq, serde::Serialize, serde::Deserialize, Debug)] @@ -228,6 +238,48 @@ pub fn get_statistics_desc() -> Arc unit: StatisticsUnit::Bytes, plain_statistics: true, }), + (ProfileStatisticsName::RowFetchRows, ProfileDesc { + display_name: "row fetch rows", + desc: "The number of rows requested by RowFetch", + index: ProfileStatisticsName::RowFetchRows as usize, + unit: StatisticsUnit::Rows, + plain_statistics: true, + }), + (ProfileStatisticsName::RowFetchBlocks, ProfileDesc { + display_name: "row fetch blocks", + desc: "The number of distinct block fetches requested by RowFetch batches", + index: ProfileStatisticsName::RowFetchBlocks as usize, + unit: StatisticsUnit::Count, + plain_statistics: true, + }), + (ProfileStatisticsName::RowFetchEstimatedBytes, ProfileDesc { + display_name: "row fetch estimated bytes", + desc: "The estimated uncompressed bytes of columns requested by RowFetch", + index: ProfileStatisticsName::RowFetchEstimatedBytes as usize, + unit: StatisticsUnit::Bytes, + plain_statistics: true, + }), + (ProfileStatisticsName::RowFetchBatches, ProfileDesc { + display_name: "row fetch batches", + desc: "The number of RowFetch batches", + index: ProfileStatisticsName::RowFetchBatches as usize, + unit: StatisticsUnit::Count, + plain_statistics: true, + }), + (ProfileStatisticsName::RowFetchLocalBatches, ProfileDesc { + display_name: "row fetch local batches", + desc: "The number of RowFetch batches kept on the coordinator by adaptive routing", + index: ProfileStatisticsName::RowFetchLocalBatches as usize, + unit: StatisticsUnit::Count, + plain_statistics: true, + }), + (ProfileStatisticsName::RowFetchDistributedBatches, ProfileDesc { + display_name: "row fetch distributed batches", + desc: "The number of RowFetch batches hash-distributed by adaptive routing", + index: ProfileStatisticsName::RowFetchDistributedBatches as usize, + unit: StatisticsUnit::Count, + plain_statistics: true, + }), (ProfileStatisticsName::RemoteSpillWriteCount, ProfileDesc { display_name: "numbers remote spilled by write", desc: "The number of remote spilled by write", @@ -368,6 +420,20 @@ pub fn get_statistics_desc() -> Arc unit: StatisticsUnit::Count, plain_statistics: true, }), + (ProfileStatisticsName::RowFetchInputBatches, ProfileDesc { + display_name: "row fetch input batches", + desc: "The number of input batches coalesced before adaptive RowFetch routing", + index: ProfileStatisticsName::RowFetchInputBatches as usize, + unit: StatisticsUnit::Count, + plain_statistics: true, + }), + (ProfileStatisticsName::RowFetchAffinityReassignedBlocks, ProfileDesc { + display_name: "row fetch affinity reassigned blocks", + desc: "The number of RowFetch blocks moved from their primary hash destination to bound load skew", + index: ProfileStatisticsName::RowFetchAffinityReassignedBlocks as usize, + unit: StatisticsUnit::Count, + plain_statistics: true, + }), ])) }).clone() } diff --git a/src/common/metrics/src/metrics/storage.rs b/src/common/metrics/src/metrics/storage.rs index b310bd7b48890..3dfe904043a78 100644 --- a/src/common/metrics/src/metrics/storage.rs +++ b/src/common/metrics/src/metrics/storage.rs @@ -147,6 +147,22 @@ static REMOTE_IO_READ_MILLISECONDS: LazyLock = LazyLock::new(|| register_histogram_in_milliseconds("fuse_remote_io_read_milliseconds")); static REMOTE_IO_DESERIALIZE_MILLISECONDS: LazyLock = LazyLock::new(|| register_histogram_in_milliseconds("fuse_remote_io_deserialize_milliseconds")); +static ROW_FETCH_ROWS: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_rows")); +static ROW_FETCH_BLOCKS: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_blocks")); +static ROW_FETCH_ESTIMATED_BYTES: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_estimated_bytes")); +static ROW_FETCH_BATCHES: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_batches")); +static ROW_FETCH_LOCAL_BATCHES: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_local_batches")); +static ROW_FETCH_DISTRIBUTED_BATCHES: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_distributed_batches")); +static ROW_FETCH_INPUT_BATCHES: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_input_batches")); +static ROW_FETCH_AFFINITY_REASSIGNED_BLOCKS: LazyLock = + LazyLock::new(|| register_counter("fuse_row_fetch_affinity_reassigned_blocks")); static BLOCK_WRITE_NUMS: LazyLock = LazyLock::new(|| register_counter("fuse_block_write_nums")); static BLOCK_WRITE_BYTES: LazyLock = @@ -589,6 +605,38 @@ pub fn metrics_inc_remote_io_deserialize_milliseconds(c: u64) { REMOTE_IO_DESERIALIZE_MILLISECONDS.observe(c as f64); } +pub fn metrics_inc_row_fetch_rows(c: u64) { + ROW_FETCH_ROWS.inc_by(c); +} + +pub fn metrics_inc_row_fetch_blocks(c: u64) { + ROW_FETCH_BLOCKS.inc_by(c); +} + +pub fn metrics_inc_row_fetch_estimated_bytes(c: u64) { + ROW_FETCH_ESTIMATED_BYTES.inc_by(c); +} + +pub fn metrics_inc_row_fetch_batches(c: u64) { + ROW_FETCH_BATCHES.inc_by(c); +} + +pub fn metrics_inc_row_fetch_local_batches(c: u64) { + ROW_FETCH_LOCAL_BATCHES.inc_by(c); +} + +pub fn metrics_inc_row_fetch_distributed_batches(c: u64) { + ROW_FETCH_DISTRIBUTED_BATCHES.inc_by(c); +} + +pub fn metrics_inc_row_fetch_input_batches(c: u64) { + ROW_FETCH_INPUT_BATCHES.inc_by(c); +} + +pub fn metrics_inc_row_fetch_affinity_reassigned_blocks(c: u64) { + ROW_FETCH_AFFINITY_REASSIGNED_BLOCKS.inc_by(c); +} + /// Block metrics. pub fn metrics_inc_block_write_nums(c: u64) { BLOCK_WRITE_NUMS.inc_by(c); diff --git a/src/query/service/src/physical_plans/format/format_exchange.rs b/src/query/service/src/physical_plans/format/format_exchange.rs index 01f47b0b5fedf..00591e17920c1 100644 --- a/src/query/service/src/physical_plans/format/format_exchange.rs +++ b/src/query/service/src/physical_plans/format/format_exchange.rs @@ -57,6 +57,19 @@ impl<'a> PhysicalFormat for ExchangeFormatter<'a> { .collect::>() .join(", ") ), + FragmentKind::RowFetch { + local_block_threshold, + .. + } => format!( + "Adaptive RowFetch(Hash({}), local blocks <= {})", + self.inner + .keys + .iter() + .map(|key| { key.as_expr(&BUILTIN_FUNCTIONS).sql_display() }) + .collect::>() + .join(", "), + local_block_threshold + ), FragmentKind::Expansive => "Broadcast".to_string(), FragmentKind::Merge => "Merge".to_string(), FragmentKind::GlobalShuffle => format!( diff --git a/src/query/service/src/physical_plans/physical_limit.rs b/src/query/service/src/physical_plans/physical_limit.rs index 237e87cd3e1ea..55a152cd95417 100644 --- a/src/query/service/src/physical_plans/physical_limit.rs +++ b/src/query/service/src/physical_plans/physical_limit.rs @@ -16,19 +16,31 @@ use std::any::Any; use std::collections::HashMap; use databend_common_catalog::plan::DataSourcePlan; +use databend_common_catalog::plan::NUM_ROW_ID_PREFIX_BITS; use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_expression::ColumnRef; +use databend_common_expression::Constant; use databend_common_expression::DataField; +use databend_common_expression::Expr; use databend_common_expression::ROW_ID_COL_NAME; +use databend_common_expression::Scalar; +use databend_common_expression::type_check::check_function; +use databend_common_expression::types::DataType; +use databend_common_expression::types::NumberDataType; +use databend_common_functions::BUILTIN_FUNCTIONS; use databend_common_pipeline::core::ProcessorPtr; use databend_common_pipeline_transforms::filters::TransformLimit; use databend_common_sql::ColumnEntry; use databend_common_sql::ColumnSet; use databend_common_sql::IndexType; use databend_common_sql::Symbol; +use databend_common_sql::executor::physical_plans::FragmentKind; use databend_common_sql::optimizer::ir::SExpr; +use crate::physical_plans::Exchange; use crate::physical_plans::PhysicalPlanBuilder; +use crate::physical_plans::Sort; use crate::physical_plans::explain::PlanStatsInfo; use crate::physical_plans::format::LimitFormatter; use crate::physical_plans::format::PhysicalFormat; @@ -37,6 +49,7 @@ use crate::physical_plans::physical_plan::PhysicalPlan; use crate::physical_plans::physical_plan::PhysicalPlanCast; use crate::physical_plans::physical_plan::PhysicalPlanMeta; use crate::physical_plans::physical_row_fetch::RowFetch; +use crate::physical_plans::physical_sort::SortStep; use crate::pipelines::PipelineBuilder; #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] @@ -130,6 +143,44 @@ impl IPhysicalPlan for Limit { } impl PhysicalPlanBuilder { + fn contains_merge_exchange(plan: &PhysicalPlan) -> bool { + if let Some(exchange) = Exchange::from_physical_plan(plan) + && exchange.kind == FragmentKind::Merge + { + return true; + } + + plan.children().any(Self::contains_merge_exchange) + } + + fn build_row_fetch_shuffle_key( + input_schema: &databend_common_expression::DataSchema, + row_id_col_offset: usize, + ) -> Result { + let row_id_field = input_schema.field(row_id_col_offset); + let row_id_expr = Expr::ColumnRef(ColumnRef { + span: None, + id: row_id_col_offset, + data_type: row_id_field.data_type().clone(), + display_name: row_id_field.name().to_string(), + }); + let block_prefix = check_function( + None, + "bit_shift_right", + &[], + &[ + row_id_expr, + Expr::Constant(Constant { + span: None, + scalar: Scalar::Number(((64 - NUM_ROW_ID_PREFIX_BITS) as u64).into()), + data_type: DataType::Number(NumberDataType::UInt64), + }), + ], + &BUILTIN_FUNCTIONS, + )?; + Ok(block_prefix.as_remote_expr()) + } + pub async fn build_limit( &mut self, s_expr: &SExpr, @@ -164,6 +215,7 @@ impl PhysicalPlanBuilder { } // If `lazy_columns` is not empty, build a `RowFetch` plan on top of the `Limit` plan. + let order_by = Sort::from_physical_plan(&input_plan).map(|sort| sort.order_by.clone()); let mut plan = PhysicalPlan::new(Limit { meta: PhysicalPlanMeta::new("Limit"), input: input_plan, @@ -220,6 +272,16 @@ impl PhysicalPlanBuilder { let mut table_indexes: Vec = lazy_columns_by_table.keys().copied().collect(); table_indexes.sort_unstable(); + let max_threads = self.ctx.get_settings().get_max_threads()? as usize; + // A local RowFetch already reads up to `max_threads` blocks concurrently. Avoid a + // cluster round trip for small limits, where all possible blocks fit in only a few + // local I/O waves. The runtime metrics below let us tune this conservative guard. + let distribute_row_fetch = table_indexes.len() == 1 + && limit + .limit + .is_some_and(|rows| rows > max_threads.saturating_mul(4)) + && !self.ctx.get_cluster().is_empty() + && Self::contains_merge_exchange(&plan); let mut row_id_col_offsets = HashMap::new(); for table_index in table_indexes.iter() { @@ -288,6 +350,22 @@ impl PhysicalPlanBuilder { )) })?; + if distribute_row_fetch { + let shuffle_key = + Self::build_row_fetch_shuffle_key(input_schema.as_ref(), row_id_col_offset)?; + plan = PhysicalPlan::new(Exchange { + input: plan, + kind: FragmentKind::RowFetch { + row_id_col_offset, + local_block_threshold: max_threads.saturating_mul(8).max(128), + }, + keys: vec![shuffle_key], + ignore_exchange: false, + allow_adjust_parallelism: true, + meta: PhysicalPlanMeta::new("Exchange"), + }); + } + plan = PhysicalPlan::new(RowFetch { meta: PhysicalPlanMeta::new("RowFetch"), input: plan, @@ -299,6 +377,33 @@ impl PhysicalPlanBuilder { enable_block_id_repartition: false, stat_info: Some(stat_info.clone()), }); + + if distribute_row_fetch { + plan = PhysicalPlan::new(Exchange { + input: plan, + kind: FragmentKind::Merge, + keys: vec![], + ignore_exchange: false, + allow_adjust_parallelism: true, + meta: PhysicalPlanMeta::new("Exchange"), + }); + + if let Some(order_by) = order_by.clone() + && !order_by.is_empty() + { + plan = PhysicalPlan::new(Sort { + input: plan, + order_by, + limit: None, + step: SortStep::Single, + pre_projection: None, + broadcast_id: None, + enable_fixed_rows: false, + stat_info: Some(stat_info.clone()), + meta: PhysicalPlanMeta::new("Sort"), + }); + } + } } Ok(plan) diff --git a/src/query/service/src/schedulers/fragments/fragmenter.rs b/src/query/service/src/schedulers/fragments/fragmenter.rs index 44a463eac6dff..1d3679f0957d4 100644 --- a/src/query/service/src/schedulers/fragments/fragmenter.rs +++ b/src/query/service/src/schedulers/fragments/fragmenter.rs @@ -49,6 +49,7 @@ use crate::servers::flight::v1::exchange::BroadcastExchange; use crate::servers::flight::v1::exchange::DataExchange; use crate::servers::flight::v1::exchange::MergeExchange; use crate::servers::flight::v1::exchange::NodeToNodeExchange; +use crate::servers::flight::v1::exchange::RowFetchExchange; use crate::sessions::QueryContext; use crate::sessions::TableContext; use crate::sessions::TableContextCluster; @@ -59,18 +60,13 @@ use crate::sessions::TableContextSettings; pub struct Fragmenter { ctx: Arc, query_id: String, - fragments: Vec, } impl Fragmenter { pub fn try_create(ctx: Arc) -> Result { let query_id = ctx.get_id(); - Ok(Self { - ctx, - fragments: vec![], - query_id, - }) + Ok(Self { ctx, query_id }) } /// Get ids of executor nodes. @@ -111,18 +107,38 @@ impl Fragmenter { fragment_id: self.ctx.fragment_id().next_fragment_id(), exchange: None, query_id: self.query_id.clone(), - source_fragments: self.fragments, + has_merge_input: false, }); let edges = Self::collect_fragments_edge(fragments.values()); - for (source, target) in edges { - let Some(fragment) = fragments.get_mut(&source) else { + for (source, target) in &edges { + let Some(fragment) = fragments.get_mut(source) else { continue; }; if let Some(exchange_sink) = ExchangeSink::from_mut_physical_plan(&mut fragment.plan) { - exchange_sink.destination_fragment_id = target; + exchange_sink.destination_fragment_id = *target; + } + } + + // An intermediate fragment that consumes a Merge exchange can only run + // on that exchange's coordination node. + let merge_input_targets = edges + .iter() + .filter_map(|(source, target)| { + fragments + .get(source) + .is_some_and(|fragment| { + matches!(fragment.exchange.as_ref(), Some(DataExchange::Merge(_))) + }) + .then_some(*target) + }) + .collect::>(); + + for target in merge_input_targets { + if let Some(fragment) = fragments.get_mut(&target) { + fragment.has_merge_input = true; } } @@ -222,7 +238,11 @@ impl FragmentDeriveHandle { Ok(match exchange_sink.kind { FragmentKind::Init => None, - FragmentKind::Normal => { + FragmentKind::Normal + | FragmentKind::RowFetch { + row_id_col_offset: _, + local_block_threshold: _, + } => { let destination_ids = get_executors(cluster); let mut destination_channels = Vec::with_capacity(destination_ids.len()); @@ -237,6 +257,16 @@ impl FragmentDeriveHandle { destination_channels, shuffle_keys: exchange_sink.keys.clone(), allow_adjust_parallelism: exchange_sink.allow_adjust_parallelism, + row_fetch: match exchange_sink.kind { + FragmentKind::RowFetch { + row_id_col_offset, + local_block_threshold, + } => Some(RowFetchExchange { + row_id_col_offset, + local_block_threshold, + }), + _ => None, + }, })) } FragmentKind::GlobalShuffle => { @@ -255,6 +285,7 @@ impl FragmentDeriveHandle { destination_channels, shuffle_keys: exchange_sink.keys.clone(), allow_adjust_parallelism: exchange_sink.allow_adjust_parallelism, + row_fetch: None, })) } FragmentKind::Merge => Some(MergeExchange::create( @@ -317,7 +348,7 @@ impl DeriveHandle for FragmentDeriveHandle { plan, exchange, fragment_type, - source_fragments: vec![], + has_merge_input: false, fragment_id: source_fragment_id, query_id: self.query_id.clone(), }; @@ -354,7 +385,7 @@ impl DeriveHandle for FragmentDeriveHandle { fragment_id, exchange: None, query_id: self.query_id.clone(), - source_fragments: vec![], + has_merge_input: false, }; self.fragments.insert(fragment_id, fragment); diff --git a/src/query/service/src/schedulers/fragments/plan_fragment.rs b/src/query/service/src/schedulers/fragments/plan_fragment.rs index b08ced7c01435..e3829bf484365 100644 --- a/src/query/service/src/schedulers/fragments/plan_fragment.rs +++ b/src/query/service/src/schedulers/fragments/plan_fragment.rs @@ -36,6 +36,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; @@ -78,8 +79,8 @@ pub struct PlanFragment { pub exchange: Option, pub query_id: String, - // The fragments to ask data from. - pub source_fragments: Vec, + /// Whether this intermediate fragment consumes a coordinator-only Merge exchange. + pub has_merge_input: bool, } impl PlanFragment { @@ -88,10 +89,6 @@ impl PlanFragment { ctx: Arc, actions: &mut QueryFragmentsActions, ) -> Result<()> { - // for input in self.source_fragments.iter() { - // input.get_actions(ctx.clone(), actions)?; - // } - let mut fragment_actions = QueryFragmentActions::create(self.fragment_id); match &self.fragment_type { @@ -103,18 +100,33 @@ 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. A node-shuffle + // destination still needs this fragment ID locally to construct its + // exchange receivers, so other executors get an empty source stub. + let local_executor = Fragmenter::get_local_executor(ctx.clone()); + let action = + QueryFragmentAction::create(local_executor.clone(), self.plan.clone()); fragment_actions.add_action(action); + + if matches!( + self.exchange.as_ref(), + Some( + DataExchange::Broadcast(_) + | DataExchange::NodeToNodeExchange(_) + | DataExchange::GlobalShuffleExchange(_) + ) + ) { + let empty_plan = self.empty_exchange_source_plan()?; + for executor in Fragmenter::get_executors(ctx) { + 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) { @@ -148,6 +160,27 @@ impl PlanFragment { actions.add_fragment_actions(fragment_actions) } + fn empty_exchange_source_plan(&self) -> Result { + let exchange_sink = ExchangeSink::from_physical_plan(&self.plan).ok_or_else(|| { + ErrorCode::Internal("Merge-dependent intermediate fragment has no ExchangeSink") + })?; + let schema = exchange_sink.schema.clone(); + let values = DataBlock::empty_with_schema(&schema) + .columns() + .iter() + .map(BlockEntry::to_column) + .collect(); + let empty_input = PhysicalPlan::new(ConstantTableScan { + meta: PhysicalPlanMeta::new("EmptyExchangeSource"), + values, + num_rows: 0, + output_schema: schema, + }); + let mut empty_sink = exchange_sink.clone(); + empty_sink.input = empty_input; + Ok(PhysicalPlan::new(empty_sink)) + } + /// Redistribute partitions of current source fragment to executors. fn redistribute_source_fragment( &self, diff --git a/src/query/service/src/servers/flight/v1/exchange/data_exchange.rs b/src/query/service/src/servers/flight/v1/exchange/data_exchange.rs index a3c375de4e0e4..fe80d80c0c300 100644 --- a/src/query/service/src/servers/flight/v1/exchange/data_exchange.rs +++ b/src/query/service/src/servers/flight/v1/exchange/data_exchange.rs @@ -87,6 +87,14 @@ pub struct NodeToNodeExchange { pub shuffle_keys: Vec, pub destination_channels: Vec<(String, Vec)>, pub allow_adjust_parallelism: bool, + #[serde(default)] + pub row_fetch: Option, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct RowFetchExchange { + pub row_id_col_offset: usize, + pub local_block_threshold: usize, } #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -137,3 +145,23 @@ impl BroadcastExchange { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_node_to_node_exchange_deserializes_without_row_fetch() { + let legacy_payload = r#"{ + "id": "exchange-id", + "destination_ids": [], + "shuffle_keys": [], + "destination_channels": [], + "allow_adjust_parallelism": false + }"#; + + let exchange: NodeToNodeExchange = serde_json::from_str(legacy_payload).unwrap(); + + assert!(exchange.row_fetch.is_none()); + } +} diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_injector.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_injector.rs index b1c554a4e6b9c..3dbb753da7fac 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_injector.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_injector.rs @@ -25,11 +25,13 @@ use crate::servers::flight::v1::exchange::ShuffleExchangeParams; use crate::servers::flight::v1::exchange::serde::TransformExchangeDeserializer; use crate::servers::flight::v1::exchange::serde::TransformExchangeSerializer; use crate::servers::flight::v1::exchange::serde::TransformScatterExchangeSerializer; +use crate::servers::flight::v1::scatter::AdaptiveRowFetchFlightScatter; use crate::servers::flight::v1::scatter::BroadcastFlightScatter; use crate::servers::flight::v1::scatter::FlightScatter; use crate::servers::flight::v1::scatter::HashFlightScatter; use crate::sessions::QueryContext; use crate::sessions::TableContextCluster; +use crate::sessions::TableContextQueryIdentity; use crate::sessions::TableContextSettings; pub trait ExchangeInjector: Send + Sync + 'static { @@ -95,12 +97,23 @@ impl ExchangeInjector for DefaultExchangeInjector { .iter() .position(|x| x == local_id) .unwrap(); - HashFlightScatter::try_create( + let hash_scatter = HashFlightScatter::try_create( ctx.get_function_context()?, exchange.shuffle_keys.clone(), exchange.destination_ids.len(), local_pos, - )? + )?; + match &exchange.row_fetch { + Some(row_fetch) => AdaptiveRowFetchFlightScatter::create( + hash_scatter, + ctx.get_id(), + row_fetch.row_id_col_offset, + row_fetch.local_block_threshold, + local_pos, + exchange.destination_ids.len(), + ), + None => hash_scatter, + } } })) } diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs index 51688ee138e87..30f47c0ed6f53 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs @@ -1382,6 +1382,7 @@ impl FragmentCoordinator { shuffle_scatter: exchange_injector .flight_scatter(&info.query_ctx, data_exchange)?, allow_adjust_parallelism: exchange.allow_adjust_parallelism, + row_fetch: exchange.row_fetch.clone(), }), )), DataExchange::GlobalShuffleExchange(exchange) => Ok(Some( diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs index f9bb7e15ac774..b06a2c32c303f 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs @@ -23,6 +23,7 @@ use databend_common_expression::RemoteExpr; use crate::servers::flight::FlightReceiver; use crate::servers::flight::FlightSender; use crate::servers::flight::v1::exchange::ExchangeInjector; +use crate::servers::flight::v1::exchange::RowFetchExchange; use crate::servers::flight::v1::scatter::FlightScatter; #[derive(Clone)] @@ -36,6 +37,7 @@ pub struct ShuffleExchangeParams { pub shuffle_scatter: Arc>, pub exchange_injector: Arc, pub allow_adjust_parallelism: bool, + pub row_fetch: Option, } #[derive(Clone)] diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_transform_shuffle.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_transform_shuffle.rs index 982c31ec905de..36dc97427f5e4 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_transform_shuffle.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_transform_shuffle.rs @@ -34,11 +34,13 @@ use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Pipeline; use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline_transforms::TransformPipelineHelper; use super::exchange_params::ShuffleExchangeParams; use super::exchange_sorting::ExchangeSorting; use super::exchange_sorting::TransformExchangeSorting; use super::exchange_transform_scatter::ScatterTransform; +use super::row_fetch_exchange_coalescer::RowFetchExchangeCoalescer; use super::serde::ExchangeSerializeMeta; use crate::sessions::QueryContext; use crate::sessions::TableContextSettings; @@ -382,6 +384,13 @@ pub fn exchange_shuffle( params: &ShuffleExchangeParams, pipeline: &mut Pipeline, ) -> Result<()> { + if params.row_fetch.is_some() { + pipeline.try_resize(1)?; + pipeline.add_accumulating_transformer(|| { + RowFetchExchangeCoalescer::create(params.query_id.clone()) + }); + } + // append scatter transform pipeline.add_transform(|input, output| { Ok(ScatterTransform::create( 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 e79e78ffda251..dad26eb0c6e25 100644 --- a/src/query/service/src/servers/flight/v1/exchange/mod.rs +++ b/src/query/service/src/servers/flight/v1/exchange/mod.rs @@ -30,6 +30,7 @@ mod hash_send_sink; mod hash_send_source; mod hash_send_transform; mod outbound_send_channels; +mod row_fetch_exchange_coalescer; mod statistics_receiver; mod statistics_sender; @@ -43,6 +44,7 @@ pub use data_exchange::BroadcastExchange; pub use data_exchange::DataExchange; pub use data_exchange::MergeExchange; pub use data_exchange::NodeToNodeExchange; +pub use data_exchange::RowFetchExchange; pub use exchange_injector::DefaultExchangeInjector; pub use exchange_injector::ExchangeInjector; pub use exchange_manager::DataExchangeManager; @@ -55,3 +57,4 @@ pub use exchange_transform_shuffle::ExchangeShuffleTransform; pub use hash_send_sink::HashSendSink; pub use hash_send_source::HashSendSource; pub use hash_send_transform::HashSendTransform; +pub use row_fetch_exchange_coalescer::RowFetchExchangeCoalescer; diff --git a/src/query/service/src/servers/flight/v1/exchange/row_fetch_exchange_coalescer.rs b/src/query/service/src/servers/flight/v1/exchange/row_fetch_exchange_coalescer.rs new file mode 100644 index 0000000000000..f9558c4dffffb --- /dev/null +++ b/src/query/service/src/servers/flight/v1/exchange/row_fetch_exchange_coalescer.rs @@ -0,0 +1,107 @@ +// 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_base::runtime::profile::Profile; +use databend_common_base::runtime::profile::ProfileStatisticsName; +use databend_common_exception::Result; +use databend_common_expression::DataBlock; +use databend_common_metrics::storage::metrics_inc_row_fetch_input_batches; +use databend_common_pipeline_transforms::processors::AccumulatingTransform; +use log::info; + +pub struct RowFetchExchangeCoalescer { + query_id: String, + blocks: Vec, + rows: usize, +} + +impl RowFetchExchangeCoalescer { + pub fn create(query_id: String) -> Self { + Self { + query_id, + blocks: Vec::new(), + rows: 0, + } + } +} + +impl AccumulatingTransform for RowFetchExchangeCoalescer { + const NAME: &'static str = "RowFetchExchangeCoalescer"; + + fn transform(&mut self, data: DataBlock) -> Result> { + if !data.is_empty() { + self.rows += data.num_rows(); + self.blocks.push(data); + } + Ok(vec![]) + } + + fn on_finish(&mut self, output: bool) -> Result> { + if !output || self.blocks.is_empty() { + self.blocks.clear(); + self.rows = 0; + return Ok(vec![]); + } + + let input_batches = self.blocks.len(); + let block = match input_batches { + 1 => self.blocks.pop().unwrap(), + _ => DataBlock::concat(&self.blocks)?, + }; + self.blocks.clear(); + + Profile::record_usize_profile(ProfileStatisticsName::RowFetchInputBatches, input_batches); + metrics_inc_row_fetch_input_batches(input_batches as u64); + info!( + "RowFetch exchange coalesced query_id={} input_batches={} rows={}", + self.query_id, input_batches, self.rows + ); + self.rows = 0; + + Ok(vec![block]) + } +} + +#[cfg(test)] +mod tests { + use databend_common_expression::FromData; + use databend_common_expression::types::UInt64Type; + + use super::*; + + fn block(values: &[u64]) -> DataBlock { + DataBlock::new_from_columns(vec![UInt64Type::from_data(values.to_vec())]) + } + + #[test] + fn coalesces_non_empty_blocks_once() -> Result<()> { + let mut coalescer = RowFetchExchangeCoalescer::create("test-query".to_string()); + assert!(coalescer.transform(DataBlock::empty())?.is_empty()); + assert!(coalescer.transform(block(&[1, 2]))?.is_empty()); + assert!(coalescer.transform(block(&[3]))?.is_empty()); + + let output = coalescer.on_finish(true)?; + assert_eq!(output.len(), 1); + assert_eq!(output[0].num_rows(), 3); + Ok(()) + } + + #[test] + fn does_not_emit_for_empty_input() -> Result<()> { + let mut coalescer = RowFetchExchangeCoalescer::create("test-query".to_string()); + assert!(coalescer.transform(DataBlock::empty())?.is_empty()); + assert!(coalescer.on_finish(true)?.is_empty()); + Ok(()) + } +} diff --git a/src/query/service/src/servers/flight/v1/scatter/flight_scatter_row_fetch.rs b/src/query/service/src/servers/flight/v1/scatter/flight_scatter_row_fetch.rs new file mode 100644 index 0000000000000..c19a3a663ba70 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/scatter/flight_scatter_row_fetch.rs @@ -0,0 +1,342 @@ +// 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::BTreeSet; +use std::collections::HashMap; + +use databend_common_base::runtime::profile::Profile; +use databend_common_base::runtime::profile::ProfileStatisticsName; +use databend_common_catalog::plan::split_row_id; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::DataBlock; +use databend_common_expression::types::DataType; +use databend_common_expression::types::NumberDataType; +use databend_common_metrics::storage::metrics_inc_row_fetch_affinity_reassigned_blocks; +use databend_common_metrics::storage::metrics_inc_row_fetch_distributed_batches; +use databend_common_metrics::storage::metrics_inc_row_fetch_local_batches; +use log::info; + +use super::FlightScatter; + +struct RoutingDecision { + local: bool, + indices: Vec, + distinct_blocks: usize, + destination_rows: Vec, + destination_blocks: Vec, + affinity_reassigned_blocks: usize, +} + +pub struct AdaptiveRowFetchFlightScatter { + hash_scatter: Box, + query_id: String, + row_id_col_offset: usize, + local_block_threshold: usize, + local_pos: usize, + scatter_size: usize, +} + +impl AdaptiveRowFetchFlightScatter { + pub fn create( + hash_scatter: Box, + query_id: String, + row_id_col_offset: usize, + local_block_threshold: usize, + local_pos: usize, + scatter_size: usize, + ) -> Box { + Box::new(Self { + hash_scatter, + query_id, + row_id_col_offset, + local_block_threshold, + local_pos, + scatter_size, + }) + } + + fn row_id_prefixes(&self, data_block: &DataBlock) -> Result> { + let entry = data_block + .columns() + .get(self.row_id_col_offset) + .ok_or_else(|| { + ErrorCode::Internal(format!( + "Adaptive RowFetch row ID column offset {} is out of bounds for {} columns", + self.row_id_col_offset, + data_block.num_columns() + )) + })?; + let column = entry.to_column(); + let row_ids = match entry.data_type() { + DataType::Number(NumberDataType::UInt64) => { + column.into_number().unwrap().into_u_int64().unwrap() + } + DataType::Nullable(inner) + if matches!(inner.as_ref(), DataType::Number(NumberDataType::UInt64)) => + { + column + .into_nullable() + .unwrap() + .column + .into_number() + .unwrap() + .into_u_int64() + .unwrap() + } + data_type => { + return Err(ErrorCode::Internal(format!( + "Adaptive RowFetch row ID column must be UInt64, but got {data_type}" + ))); + } + }; + + Ok(row_ids + .iter() + .map(|row_id| split_row_id(*row_id).0) + .collect()) + } + + fn local_decision(&self, prefixes: Vec, distinct_blocks: usize) -> RoutingDecision { + let mut destination_rows = vec![0; self.scatter_size]; + let mut destination_blocks = vec![0; self.scatter_size]; + destination_rows[self.local_pos] = prefixes.len(); + destination_blocks[self.local_pos] = distinct_blocks; + + RoutingDecision { + local: true, + indices: vec![self.local_pos as u64; prefixes.len()], + distinct_blocks, + destination_rows, + destination_blocks, + affinity_reassigned_blocks: 0, + } + } + + fn distributed_decision( + &self, + data_block: &DataBlock, + prefixes: Vec, + ) -> Result { + let primary_indices = self + .hash_scatter + .scatter_indices(data_block)? + .ok_or_else(|| ErrorCode::Internal("RowFetch hash scatter does not expose indices"))?; + if primary_indices.len() != prefixes.len() { + return Err(ErrorCode::Internal(format!( + "RowFetch hash scatter produced {} indices for {} rows", + primary_indices.len(), + prefixes.len() + ))); + } + + let mut blocks = HashMap::::new(); + let mut destination_rows = vec![0usize; self.scatter_size]; + let mut destination_blocks = vec![0usize; self.scatter_size]; + for (prefix, primary) in prefixes + .iter() + .copied() + .zip(primary_indices.iter().copied()) + { + let primary = primary as usize; + if primary >= self.scatter_size { + return Err(ErrorCode::Internal(format!( + "RowFetch hash scatter destination {} is out of bounds for {} destinations", + primary, self.scatter_size + ))); + } + + destination_rows[primary] += 1; + match blocks.get(&prefix) { + Some(block_primary) => { + if *block_primary != primary { + return Err(ErrorCode::Internal(format!( + "Rows from RowFetch block {} mapped to different primary destinations", + prefix + ))); + } + } + None => { + blocks.insert(prefix, primary); + destination_blocks[primary] += 1; + } + } + } + + let distinct_blocks = blocks.len(); + + Ok(RoutingDecision { + local: false, + indices: primary_indices, + distinct_blocks, + destination_rows, + destination_blocks, + affinity_reassigned_blocks: 0, + }) + } + + fn route(&self, data_block: &DataBlock) -> Result { + if self.scatter_size == 0 || self.local_pos >= self.scatter_size { + return Err(ErrorCode::Internal(format!( + "Invalid adaptive RowFetch destinations: local_pos={}, scatter_size={}", + self.local_pos, self.scatter_size + ))); + } + + let prefixes = self.row_id_prefixes(data_block)?; + let distinct_blocks = prefixes.iter().copied().collect::>().len(); + if distinct_blocks <= self.local_block_threshold { + return Ok(self.local_decision(prefixes, distinct_blocks)); + } + self.distributed_decision(data_block, prefixes) + } + + fn record_decision(&self, decision: &RoutingDecision, rows: usize) { + let mode = if decision.local { + Profile::record_usize_profile(ProfileStatisticsName::RowFetchLocalBatches, 1); + metrics_inc_row_fetch_local_batches(1); + "local" + } else { + Profile::record_usize_profile(ProfileStatisticsName::RowFetchDistributedBatches, 1); + Profile::record_usize_profile( + ProfileStatisticsName::RowFetchAffinityReassignedBlocks, + decision.affinity_reassigned_blocks, + ); + metrics_inc_row_fetch_distributed_batches(1); + metrics_inc_row_fetch_affinity_reassigned_blocks( + decision.affinity_reassigned_blocks as u64, + ); + "distributed" + }; + + info!( + "Adaptive RowFetch routing query_id={} mode={} rows={} distinct_blocks={} local_block_threshold={} destinations={} destination_rows={:?} destination_blocks={:?} affinity_reassigned_blocks={}", + self.query_id, + mode, + rows, + decision.distinct_blocks, + self.local_block_threshold, + self.scatter_size, + decision.destination_rows, + decision.destination_blocks, + decision.affinity_reassigned_blocks + ); + } +} + +impl FlightScatter for AdaptiveRowFetchFlightScatter { + fn name(&self) -> &'static str { + "AdaptiveRowFetch" + } + + fn execute(&self, data_block: DataBlock) -> Result> { + let decision = self.route(&data_block)?; + self.record_decision(&decision, data_block.num_rows()); + + let block_meta = data_block.get_meta().cloned(); + let blocks = DataBlock::scatter(&data_block, &decision.indices, self.scatter_size)?; + blocks + .into_iter() + .map(|block| block.add_meta(block_meta.clone())) + .collect() + } + + fn scatter_indices(&self, data_block: &DataBlock) -> Result>> { + Ok(Some(self.route(data_block)?.indices)) + } +} + +#[cfg(test)] +mod tests { + use databend_common_catalog::plan::compute_row_id; + use databend_common_catalog::plan::compute_row_id_prefix; + use databend_common_expression::FromData; + use databend_common_expression::types::UInt64Type; + + use super::*; + + struct PrimaryZeroScatter; + + impl FlightScatter for PrimaryZeroScatter { + fn name(&self) -> &'static str { + "PrimaryZero" + } + + fn execute(&self, data_block: DataBlock) -> Result> { + let indices = self.scatter_indices(&data_block)?.unwrap(); + DataBlock::scatter(&data_block, &indices, 3) + } + + fn scatter_indices(&self, data_block: &DataBlock) -> Result>> { + Ok(Some(vec![0; data_block.num_rows()])) + } + } + + fn row_id_block(block_ids: &[u64]) -> DataBlock { + let row_ids = block_ids + .iter() + .map(|block_id| compute_row_id(compute_row_id_prefix(0, *block_id), 0)) + .collect(); + DataBlock::new_from_columns(vec![UInt64Type::from_data(row_ids)]) + } + + fn scatter(local_block_threshold: usize) -> AdaptiveRowFetchFlightScatter { + AdaptiveRowFetchFlightScatter { + hash_scatter: Box::new(PrimaryZeroScatter), + query_id: "test-query".to_string(), + row_id_col_offset: 0, + local_block_threshold, + local_pos: 1, + scatter_size: 3, + } + } + + #[test] + fn keeps_compact_row_fetch_local() -> Result<()> { + let block = row_id_block(&[1, 1, 2, 2]); + assert_eq!(scatter(2).scatter_indices(&block)?, Some(vec![1, 1, 1, 1])); + Ok(()) + } + + #[test] + fn preserves_primary_hash_affinity_for_dispersed_blocks() -> Result<()> { + let block = row_id_block(&[1, 1, 2, 3, 4]); + let decision = scatter(2).route(&block)?; + + assert!(!decision.local); + assert_eq!(decision.distinct_blocks, 4); + assert_eq!(decision.indices, vec![0; 5]); + assert_eq!(decision.destination_rows, vec![5, 0, 0]); + assert_eq!(decision.destination_blocks, vec![4, 0, 0]); + assert_eq!(decision.affinity_reassigned_blocks, 0); + Ok(()) + } + + #[test] + fn does_not_reassign_skewed_primary_blocks() -> Result<()> { + let mut block_ids = vec![1; 100]; + block_ids.extend(vec![2; 50]); + block_ids.extend([3, 4]); + let block = row_id_block(&block_ids); + let decision = scatter(2).route(&block)?; + + assert_eq!(decision.destination_rows.iter().sum::(), 152); + assert_eq!(decision.destination_blocks.iter().sum::(), 4); + assert_eq!(decision.destination_rows, vec![152, 0, 0]); + assert_eq!(decision.destination_blocks, vec![4, 0, 0]); + assert_eq!(decision.affinity_reassigned_blocks, 0); + assert_eq!(decision.indices, vec![0; 152]); + Ok(()) + } +} diff --git a/src/query/service/src/servers/flight/v1/scatter/mod.rs b/src/query/service/src/servers/flight/v1/scatter/mod.rs index b5f5f900dab71..887b8539025b4 100644 --- a/src/query/service/src/servers/flight/v1/scatter/mod.rs +++ b/src/query/service/src/servers/flight/v1/scatter/mod.rs @@ -15,7 +15,9 @@ mod flight_scatter; mod flight_scatter_broadcast; mod flight_scatter_hash; +mod flight_scatter_row_fetch; pub use flight_scatter::FlightScatter; pub use flight_scatter_broadcast::BroadcastFlightScatter; pub use flight_scatter_hash::HashFlightScatter; +pub use flight_scatter_row_fetch::AdaptiveRowFetchFlightScatter; diff --git a/src/query/service/tests/it/sql/planner/optimizer/optimizer_test.rs b/src/query/service/tests/it/sql/planner/optimizer/optimizer_test.rs index 5f265988c1e30..7abd222247eee 100644 --- a/src/query/service/tests/it/sql/planner/optimizer/optimizer_test.rs +++ b/src/query/service/tests/it/sql/planner/optimizer/optimizer_test.rs @@ -32,7 +32,13 @@ use databend_common_sql_test_support::configure_optimizer_settings; use databend_common_sql_test_support::run_test_case_core; use databend_meta_client::types::NodeInfo; use databend_query::clusters::ClusterHelper; +use databend_query::physical_plans::ConstantTableScan; +use databend_query::physical_plans::ExchangeSink; use databend_query::physical_plans::PhysicalPlanBuilder; +use databend_query::physical_plans::PhysicalPlanCast; +use databend_query::schedulers::Fragmenter; +use databend_query::schedulers::QueryFragmentsActions; +use databend_query::servers::flight::v1::exchange::DataExchange; use databend_query::sessions::QueryContext; use databend_query::sessions::TableContextCluster; use databend_query::sessions::TableContextSettings; @@ -133,6 +139,96 @@ async fn test_optimizer() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_merge_dependent_shuffle_runs_on_coordinator() -> anyhow::Result<()> { + let fixture = TestFixture::setup().await?; + let ctx = fixture.new_query_ctx().await?; + let settings = ctx.get_settings(); + configure_optimizer_settings(settings.as_ref(), false)?; + settings.set_max_threads(16)?; + settings.set_setting("lazy_read_threshold".to_string(), "1000".to_string())?; + + let mut nodes = vec![ + create_node(&GlobalUniq::unique()), + create_node(&GlobalUniq::unique()), + ]; + let local_id = GlobalUniq::unique(); + nodes.push(create_node(&local_id)); + ctx.set_cluster(Cluster::create(nodes, local_id.clone())); + + execute_sql( + &ctx, + "CREATE TABLE lazy_row_fetch_fragment_test (a INT, b STRING) \ + ENGINE = FUSE STORAGE_FORMAT = 'PARQUET'", + ) + .await?; + + let runner = ServiceRunner(ctx.clone()); + let plan = runner + .bind_sql( + "SELECT a, b FROM lazy_row_fetch_fragment_test \ + WHERE a > 0 ORDER BY a DESC LIMIT 1000", + ) + .await?; + let optimized = runner.optimize_plan(plan).await?; + let Plan::Query { + metadata, + bind_context, + s_expr, + .. + } = &optimized + else { + unreachable!("expected query plan") + }; + + let mut builder = PhysicalPlanBuilder::new(metadata.clone(), ctx.clone(), false); + let physical = builder.build(s_expr, bind_context.column_set()).await?; + let fragments = Fragmenter::try_create(ctx.clone())?.build_fragment(&physical)?; + let mut fragments_actions = QueryFragmentsActions::create(ctx.clone()); + for fragment in fragments { + fragment.get_actions(ctx.clone(), &mut fragments_actions)?; + } + + let shuffle_actions = fragments_actions + .fragments_actions + .iter() + .filter(|actions| { + matches!( + actions.data_exchange.as_ref(), + Some(DataExchange::NodeToNodeExchange(_)) + ) + }) + .collect::>(); + assert_eq!(shuffle_actions.len(), 1); + assert_eq!(shuffle_actions[0].fragment_actions.len(), 3); + let Some(DataExchange::NodeToNodeExchange(exchange)) = + shuffle_actions[0].data_exchange.as_ref() + else { + unreachable!("expected adaptive RowFetch shuffle") + }; + let row_fetch = exchange + .row_fetch + .as_ref() + .expect("RowFetch shuffle must include adaptive routing"); + assert_eq!(row_fetch.local_block_threshold, 128); + + for action in &shuffle_actions[0].fragment_actions { + let exchange_sink = ExchangeSink::from_physical_plan(&action.physical_plan) + .expect("shuffle action must contain an ExchangeSink"); + + if action.executor == local_id { + assert!(!ConstantTableScan::check_physical_plan( + &exchange_sink.input + )); + } else { + assert!(ConstantTableScan::check_physical_plan(&exchange_sink.input)); + } + } + + execute_sql(&ctx, "DROP TABLE lazy_row_fetch_fragment_test").await?; + Ok(()) +} + fn create_node(local_id: &str) -> Arc { let mut node_info = NodeInfo::create( local_id.to_string(), diff --git a/src/query/sql/src/executor/physical_plans/common.rs b/src/query/sql/src/executor/physical_plans/common.rs index 3023ae01d72b2..e56035a1afe0b 100644 --- a/src/query/sql/src/executor/physical_plans/common.rs +++ b/src/query/sql/src/executor/physical_plans/common.rs @@ -68,6 +68,11 @@ pub enum FragmentKind { Init, // Partitioned by hash Normal, + // RowFetch chooses local or hash routing from the runtime block dispersion. + RowFetch { + row_id_col_offset: usize, + local_block_threshold: usize, + }, // Broadcast Expansive, Merge, diff --git a/src/query/sql/test-support/data/cases/lazy_row_fetch/distributed_row_fetch.yaml b/src/query/sql/test-support/data/cases/lazy_row_fetch/distributed_row_fetch.yaml new file mode 100644 index 0000000000000..4f99d13bc3268 --- /dev/null +++ b/src/query/sql/test-support/data/cases/lazy_row_fetch/distributed_row_fetch.yaml @@ -0,0 +1,11 @@ +name: distributed_row_fetch +description: Distribute lazy RowFetch by physical block while preserving TopK order +node_num: 3 +sql: | + SELECT a, b, c + FROM lazy_row_fetch + WHERE a > 0 + ORDER BY a DESC + LIMIT 1000 +tables: + lazy_row_fetch: lazy_row_fetch/lazy_row_fetch.sql diff --git a/src/query/sql/test-support/data/cases/lazy_row_fetch/local_row_fetch_small_limit.yaml b/src/query/sql/test-support/data/cases/lazy_row_fetch/local_row_fetch_small_limit.yaml new file mode 100644 index 0000000000000..6d3108141adde --- /dev/null +++ b/src/query/sql/test-support/data/cases/lazy_row_fetch/local_row_fetch_small_limit.yaml @@ -0,0 +1,11 @@ +name: local_row_fetch_small_limit +description: Keep small lazy RowFetch local in a cluster +node_num: 3 +sql: | + SELECT a, b, c + FROM lazy_row_fetch + WHERE a > 0 + ORDER BY a DESC + LIMIT 10 +tables: + lazy_row_fetch: lazy_row_fetch/lazy_row_fetch.sql diff --git a/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_optimized.txt b/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_optimized.txt new file mode 100644 index 0000000000000..0d926bac05197 --- /dev/null +++ b/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_optimized.txt @@ -0,0 +1,16 @@ +Limit +├── limit: [1000] +├── offset: [0] +└── Sort + ├── sort keys: [lazy_row_fetch.a (#0) DESC NULLS LAST] + ├── limit: [1000] + └── Exchange(MergeSort) + └── Sort + ├── sort keys: [lazy_row_fetch.a (#0) DESC NULLS LAST] + ├── limit: [1000] + └── Scan + ├── table: default.lazy_row_fetch (#0) + ├── filters: [gt(lazy_row_fetch.a (#0), 0)] + ├── order by: [lazy_row_fetch.a (#0) DESC] + └── limit: 1000 + diff --git a/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_physical.txt b/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_physical.txt new file mode 100644 index 0000000000000..3192010674182 --- /dev/null +++ b/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_physical.txt @@ -0,0 +1,41 @@ +Sort(Single) +├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), lazy_row_fetch.b (#1), lazy_row_fetch.c (#2)] +├── sort keys: [a DESC NULLS LAST] +├── estimated rows: 0.00 +└── Exchange + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), lazy_row_fetch.b (#1), lazy_row_fetch.c (#2)] + ├── exchange type: Merge + └── RowFetch + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), lazy_row_fetch.b (#1), lazy_row_fetch.c (#2)] + ├── columns to fetch: [b, c] + ├── estimated rows: 0.00 + └── Exchange + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3)] + ├── exchange type: Adaptive RowFetch(Hash(bit_shift_right(3, 31)), local blocks <= 128) + └── Limit + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3)] + ├── limit: 1000 + ├── offset: 0 + ├── estimated rows: 0.00 + └── Sort(Final) + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3)] + ├── sort keys: [a DESC NULLS LAST] + ├── estimated rows: 0.00 + └── Exchange + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), #_order_col] + ├── exchange type: Merge + └── Sort(Partial) + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), #_order_col] + ├── sort keys: [a DESC NULLS LAST] + ├── estimated rows: 0.00 + └── TableScan + ├── table: default.default.lazy_row_fetch + ├── scan id: 0 + ├── output columns: [a (#0), _row_id (#3)] + ├── read rows: 0 + ├── read size: 0 + ├── partitions total: 0 + ├── partitions scanned: 0 + ├── push downs: [filters: [is_true(lazy_row_fetch.a (#0) > 0)], limit: 1000] + └── estimated rows: 0.00 + diff --git a/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_raw.txt b/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_raw.txt new file mode 100644 index 0000000000000..74433f0b7de56 --- /dev/null +++ b/src/query/sql/test-support/data/results/lazy_row_fetch/distributed_row_fetch_raw.txt @@ -0,0 +1,16 @@ +Limit +├── limit: [1000] +├── offset: [0] +└── Sort + ├── sort keys: [lazy_row_fetch.a (#0) DESC NULLS LAST] + ├── limit: [NONE] + └── EvalScalar + ├── scalars: [lazy_row_fetch.a (#0) AS (#0), lazy_row_fetch.b (#1) AS (#1), lazy_row_fetch.c (#2) AS (#2)] + └── Filter + ├── filters: [gt(lazy_row_fetch.a (#0), 0)] + └── Scan + ├── table: default.lazy_row_fetch (#0) + ├── filters: [] + ├── order by: [] + └── limit: NONE + diff --git a/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_optimized.txt b/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_optimized.txt new file mode 100644 index 0000000000000..c3b3089fe39a9 --- /dev/null +++ b/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_optimized.txt @@ -0,0 +1,16 @@ +Limit +├── limit: [10] +├── offset: [0] +└── Sort + ├── sort keys: [lazy_row_fetch.a (#0) DESC NULLS LAST] + ├── limit: [10] + └── Exchange(MergeSort) + └── Sort + ├── sort keys: [lazy_row_fetch.a (#0) DESC NULLS LAST] + ├── limit: [10] + └── Scan + ├── table: default.lazy_row_fetch (#0) + ├── filters: [gt(lazy_row_fetch.a (#0), 0)] + ├── order by: [lazy_row_fetch.a (#0) DESC] + └── limit: 10 + diff --git a/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_physical.txt b/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_physical.txt new file mode 100644 index 0000000000000..a6bad32f117f9 --- /dev/null +++ b/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_physical.txt @@ -0,0 +1,31 @@ +RowFetch +├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), lazy_row_fetch.b (#1), lazy_row_fetch.c (#2)] +├── columns to fetch: [b, c] +├── estimated rows: 0.00 +└── Limit + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3)] + ├── limit: 10 + ├── offset: 0 + ├── estimated rows: 0.00 + └── Sort(Final) + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3)] + ├── sort keys: [a DESC NULLS LAST] + ├── estimated rows: 0.00 + └── Exchange + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), #_order_col] + ├── exchange type: Merge + └── Sort(Partial) + ├── output columns: [lazy_row_fetch.a (#0), lazy_row_fetch._row_id (#3), #_order_col] + ├── sort keys: [a DESC NULLS LAST] + ├── estimated rows: 0.00 + └── TableScan + ├── table: default.default.lazy_row_fetch + ├── scan id: 0 + ├── output columns: [a (#0), _row_id (#3)] + ├── read rows: 0 + ├── read size: 0 + ├── partitions total: 0 + ├── partitions scanned: 0 + ├── push downs: [filters: [is_true(lazy_row_fetch.a (#0) > 0)], limit: 10] + └── estimated rows: 0.00 + diff --git a/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_raw.txt b/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_raw.txt new file mode 100644 index 0000000000000..6a4c3f773ff9b --- /dev/null +++ b/src/query/sql/test-support/data/results/lazy_row_fetch/local_row_fetch_small_limit_raw.txt @@ -0,0 +1,16 @@ +Limit +├── limit: [10] +├── offset: [0] +└── Sort + ├── sort keys: [lazy_row_fetch.a (#0) DESC NULLS LAST] + ├── limit: [NONE] + └── EvalScalar + ├── scalars: [lazy_row_fetch.a (#0) AS (#0), lazy_row_fetch.b (#1) AS (#1), lazy_row_fetch.c (#2) AS (#2)] + └── Filter + ├── filters: [gt(lazy_row_fetch.a (#0), 0)] + └── Scan + ├── table: default.lazy_row_fetch (#0) + ├── filters: [] + ├── order by: [] + └── limit: NONE + diff --git a/src/query/sql/test-support/data/tables/lazy_row_fetch/lazy_row_fetch.sql b/src/query/sql/test-support/data/tables/lazy_row_fetch/lazy_row_fetch.sql new file mode 100644 index 0000000000000..a2f852529d982 --- /dev/null +++ b/src/query/sql/test-support/data/tables/lazy_row_fetch/lazy_row_fetch.sql @@ -0,0 +1,5 @@ +CREATE TABLE lazy_row_fetch ( + a INT, + b STRING, + c STRING +) ENGINE = FUSE STORAGE_FORMAT = 'PARQUET'; diff --git a/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs b/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs index 98b4000a8c210..f215c529d992c 100644 --- a/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs +++ b/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs @@ -68,6 +68,7 @@ pub fn row_fetch_processor( .map(|field| DataType::from(field.data_type())) .collect::>(); let block_reader = fuse_table.create_block_reader(ctx.clone(), projection.clone(), true)?; + let query_id = ctx.get_id(); match &fuse_table.storage_format { FuseStorageFormat::Parquet => { @@ -92,6 +93,7 @@ pub fn row_fetch_processor( read_settings, max_threads, io_semaphore.clone(), + query_id.clone(), ), need_wrap_nullable, fetched_data_types.clone(), @@ -115,10 +117,14 @@ impl RowsFetchMetadata for Arc { #[async_trait::async_trait] pub trait RowsFetcher { - type Metadata: RowsFetchMetadata; + type Metadata: RowsFetchMetadata + Clone; async fn initialize(&mut self) -> Result<()>; async fn fetch_metadata(&mut self, _block_id: u64) -> Result; + async fn fetch_metadata_batch( + &mut self, + block_ids: &[u64], + ) -> Result>; async fn fetch( &mut self, @@ -237,13 +243,23 @@ impl Processor for TransformRowsFetcher< match entry.data_type() { DataType::Number(NumberDataType::UInt64) => { let row_id_column = self.get_row_id_column(&data); + let block_ids = row_id_column + .iter() + .map(|row_id| split_row_id(*row_id).0) + .collect::>() + .into_iter() + .collect::>(); + let prefetched_metadata = self.fetcher.fetch_metadata_batch(&block_ids).await?; // Process the row id column in block batch // Ensure that the same block would be processed in the same batch and threads for (idx, row_id) in row_id_column.iter().enumerate() { let (prefix, _) = split_row_id(*row_id); if !self.metadata.contains_key(&prefix) { - let metadata = self.fetcher.fetch_metadata(prefix).await?; + let metadata = match prefetched_metadata.get(&prefix) { + Some(metadata) => metadata.clone(), + None => self.fetcher.fetch_metadata(prefix).await?, + }; self.metadata.insert(prefix, metadata); } @@ -265,6 +281,15 @@ impl Processor for TransformRowsFetcher< let column = entry.to_column(); let value = column.into_nullable().unwrap(); let row_id_column = value.column.into_number().unwrap().into_u_int64().unwrap(); + let block_ids = row_id_column + .iter() + .zip(value.validity.iter()) + .filter(|(_, is_valid)| *is_valid) + .map(|(row_id, _)| split_row_id(*row_id).0) + .collect::>() + .into_iter() + .collect::>(); + let prefetched_metadata = self.fetcher.fetch_metadata_batch(&block_ids).await?; for (idx, (row_id, is_valid)) in row_id_column.iter().zip(value.validity.iter()).enumerate() @@ -280,7 +305,10 @@ impl Processor for TransformRowsFetcher< let (prefix, _) = split_row_id(*row_id); if !self.metadata.contains_key(&prefix) { - let metadata = self.fetcher.fetch_metadata(prefix).await?; + let metadata = match prefetched_metadata.get(&prefix) { + Some(metadata) => metadata.clone(), + None => self.fetcher.fetch_metadata(prefix).await?, + }; self.metadata.insert(prefix, metadata); } diff --git a/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs b/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs index 375acb536a130..47362bd276807 100644 --- a/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs +++ b/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs @@ -17,6 +17,8 @@ use std::collections::hash_map::Entry; use std::future::Future; use std::sync::Arc; +use databend_common_base::runtime::profile::Profile; +use databend_common_base::runtime::profile::ProfileStatisticsName; use databend_common_base::runtime::spawn; use databend_common_catalog::plan::Projection; use databend_common_catalog::plan::block_id_in_segment; @@ -30,6 +32,10 @@ use databend_common_exception::Result; use databend_common_expression::ColumnId; use databend_common_expression::DataBlock; use databend_common_expression::TableSchemaRef; +use databend_common_metrics::storage::metrics_inc_row_fetch_batches; +use databend_common_metrics::storage::metrics_inc_row_fetch_blocks; +use databend_common_metrics::storage::metrics_inc_row_fetch_estimated_bytes; +use databend_common_metrics::storage::metrics_inc_row_fetch_rows; use databend_common_storage::ColumnNodes; use databend_storages_common_cache::CacheAccessor; use databend_storages_common_cache::CacheValue; @@ -42,7 +48,10 @@ use databend_storages_common_table_meta::meta::Compression; use databend_storages_common_table_meta::meta::TableSnapshot; use futures_util::stream::FuturesUnordered; use futures_util::stream::StreamExt; +use futures_util::stream::TryStreamExt; use itertools::Itertools; +use log::debug; +use log::info; use tokio::sync::Semaphore; use tokio::task::JoinHandle; @@ -106,6 +115,7 @@ impl RowsFetchMetadata for RowsFetchMetadataImpl { } pub(super) struct ParquetRowsFetcher { + query_id: String, snapshot: Option>, table: Arc, projection: Projection, @@ -180,12 +190,155 @@ impl RowsFetcher for ParquetRowsFetcher { .unwrap()) } + async fn fetch_metadata_batch( + &mut self, + block_ids: &[u64], + ) -> Result> { + let started = std::time::Instant::now(); + let mut result = HashMap::with_capacity(block_ids.len()); + let mut missing_by_segment = HashMap::>::new(); + for block_id in block_ids { + if let Some(metadata) = self.block_meta_lru_cache.get(block_id.to_string()) { + result.insert(*block_id, metadata.clone()); + } else { + missing_by_segment + .entry(split_prefix(*block_id).0) + .or_default() + .push(*block_id); + } + } + + if missing_by_segment.is_empty() { + return Ok(result); + } + let cached_blocks = result.len(); + let segments_to_read = missing_by_segment.len(); + + let snapshot = self + .snapshot + .as_ref() + .expect("ParquetRowsFetcher must be initialized before fetching metadata") + .clone(); + let operator = self.table.operator.clone(); + let schema = self.schema.clone(); + let requests = futures_util::stream::iter(missing_by_segment.into_iter().map( + |(segment, block_ids)| { + let snapshot = snapshot.clone(); + let segment_reader = + MetaReaders::segment_info_reader(operator.clone(), schema.clone()); + async move { + let (location, ver) = snapshot + .segments + .get(segment as usize) + .cloned() + .ok_or_else(|| { + ErrorCode::Internal(format!( + "RowFetch segment ID {segment} is out of bounds for {} segments", + snapshot.segments.len() + )) + })?; + let compact_segment_info = segment_reader + .read(&LoadParams { + ver, + location, + len_hint: None, + put_cache: true, + }) + .await?; + Ok::<_, ErrorCode>((block_ids, compact_segment_info)) + } + }, + )) + .buffer_unordered(self.max_threads.max(1)) + .try_collect::>() + .await?; + + for (block_ids, compact_segment_info) in requests { + let blocks = compact_segment_info.block_metas()?; + let mut prefixes = Vec::with_capacity(block_ids.len()); + let mut block_metas = Vec::with_capacity(block_ids.len()); + for block_id in block_ids { + let (_, block) = split_prefix(block_id); + if block as usize >= blocks.len() { + return Err(ErrorCode::Internal(format!( + "RowFetch block ID {block} is out of bounds for {} blocks", + blocks.len() + ))); + } + let block_idx = block_idx_in_segment(blocks.len(), block as usize); + let block_meta = blocks[block_idx].clone(); + prefixes.push(block_id); + block_metas.push(block_meta); + } + + for (prefix, metadata) in prefixes.into_iter().zip(self.build_metadata(&block_metas)?) { + self.block_meta_lru_cache + .insert(prefix.to_string(), metadata); + let metadata = self + .block_meta_lru_cache + .get(prefix.to_string()) + .clone() + .expect("inserted RowFetch metadata must be cached"); + result.insert(prefix, metadata); + } + } + + let message = format!( + "RowFetch metadata batch query_id={} blocks={} cached_blocks={} segments_read={} elapsed_ms={} max_concurrency={}", + self.query_id, + block_ids.len(), + cached_blocks, + segments_to_read, + started.elapsed().as_millis(), + self.max_threads + ); + if segments_to_read >= 8 { + info!("{}", message); + } else { + debug!("{}", message); + } + + Ok(result) + } + #[async_backtrace::framed] async fn fetch( &mut self, row_ids: &[u64], metadata: HashMap, ) -> Result { + let block_count = metadata.len(); + let estimated_bytes = metadata + .values() + .map(|metadata| metadata.block_bytes) + .sum::(); + + Profile::record_usize_profile(ProfileStatisticsName::RowFetchRows, row_ids.len()); + Profile::record_usize_profile(ProfileStatisticsName::RowFetchBlocks, block_count); + Profile::record_usize_profile( + ProfileStatisticsName::RowFetchEstimatedBytes, + estimated_bytes, + ); + Profile::record_usize_profile(ProfileStatisticsName::RowFetchBatches, 1); + metrics_inc_row_fetch_rows(row_ids.len() as u64); + metrics_inc_row_fetch_blocks(block_count as u64); + metrics_inc_row_fetch_estimated_bytes(estimated_bytes as u64); + metrics_inc_row_fetch_batches(1); + + let message = format!( + "RowFetch batch query_id={} rows={} distinct_blocks={} estimated_bytes={} max_concurrency={}", + self.query_id, + row_ids.len(), + block_count, + estimated_bytes, + self.max_threads + ); + if block_count >= 128 || estimated_bytes >= 256 * 1024 * 1024 { + info!("{}", message); + } else { + debug!("{}", message); + } + let final_block_index = metadata .keys() .enumerate() @@ -288,11 +441,13 @@ impl ParquetRowsFetcher { settings: ReadSettings, max_threads: usize, io_semaphore: Arc, + query_id: String, ) -> Self { let schema = table.schema(); let operator = table.operator.clone(); let segment_reader = MetaReaders::segment_info_reader(operator, schema.clone()); ParquetRowsFetcher { + query_id, table, snapshot: None, segment_reader, diff --git a/tests/sqllogictests/suites/mode/cluster/lazy_read.test b/tests/sqllogictests/suites/mode/cluster/lazy_read.test index aa382477e09ed..86b8dae330aa8 100644 --- a/tests/sqllogictests/suites/mode/cluster/lazy_read.test +++ b/tests/sqllogictests/suites/mode/cluster/lazy_read.test @@ -181,3 +181,33 @@ Limit statement ok drop table t_lazy + +# Exercise adaptive distributed RowFetch with more than the 128-block local threshold. +statement ok +set lazy_read_threshold = 200 + +statement ok +create or replace table t_adaptive_row_fetch (id uint64 not null, payload string not null) row_per_block = 1 + +statement ok +insert into t_adaptive_row_fetch select number, to_string(number) from numbers(300) + +query I +select count(distinct _block_name) from t_adaptive_row_fetch +---- +300 + +# Validate that fetched payloads remain paired with their row IDs after shuffle and merge. +query IIII +select count(), sum(if(payload = to_string(id), 1, 0)), min(id), max(id) +from ( + select id, payload + from t_adaptive_row_fetch + order by id desc + limit 150 +) +---- +150 150 150 299 + +statement ok +drop table t_adaptive_row_fetch