feat(query): adaptively distribute lazy row fetch

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<BTreeMap<ProfileStatisticsName, ProfileDesc>
                 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<BTreeMap<ProfileStatisticsName, ProfileDesc>
                 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()
 }

src/common/metrics/src/metrics/storage.rs

@@ -147,6 +147,22 @@ static REMOTE_IO_READ_MILLISECONDS: LazyLock<Histogram> =
     LazyLock::new(|| register_histogram_in_milliseconds("fuse_remote_io_read_milliseconds"));
 static REMOTE_IO_DESERIALIZE_MILLISECONDS: LazyLock<Histogram> =
     LazyLock::new(|| register_histogram_in_milliseconds("fuse_remote_io_deserialize_milliseconds"));
+static ROW_FETCH_ROWS: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_rows"));
+static ROW_FETCH_BLOCKS: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_blocks"));
+static ROW_FETCH_ESTIMATED_BYTES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_estimated_bytes"));
+static ROW_FETCH_BATCHES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_batches"));
+static ROW_FETCH_LOCAL_BATCHES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_local_batches"));
+static ROW_FETCH_DISTRIBUTED_BATCHES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_distributed_batches"));
+static ROW_FETCH_INPUT_BATCHES: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_input_batches"));
+static ROW_FETCH_AFFINITY_REASSIGNED_BLOCKS: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("fuse_row_fetch_affinity_reassigned_blocks"));
 static BLOCK_WRITE_NUMS: LazyLock<Counter> =
     LazyLock::new(|| register_counter("fuse_block_write_nums"));
 static BLOCK_WRITE_BYTES: LazyLock<Counter> =
@@ -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);

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

@@ -57,6 +57,19 @@ impl<'a> PhysicalFormat for ExchangeFormatter<'a> {
                         .collect::<Vec<_>>()
                         .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::<Vec<_>>()
+                        .join(", "),
+                    local_block_threshold
+                ),
                 FragmentKind::Expansive => "Broadcast".to_string(),
                 FragmentKind::Merge => "Merge".to_string(),
                 FragmentKind::GlobalShuffle => format!(

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<databend_common_expression::RemoteExpr> {
+        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<IndexType> = 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)