feat: add automatic block-level partition shuffle

Add heuristic-based block-level shuffle for better load balancing when
tables have few segments relative to cluster size.

Changes:
- Add BlockMod shuffle kind for block-level distribution
- Add auto_block_shuffle_threshold setting (default=5, 0 to disable)
- When segment_count < nodes * threshold, use block-level shuffle
- Each executor filters blocks by block_idx % num_executors == executor_idx
- Add info logging for shuffle strategy selection
- Preserve partition kind during reshuffle to prevent data duplication

Author: dantengsky

src/query/catalog/src/plan/partition.rs

@@ -103,6 +103,8 @@ pub enum PartitionsShuffleKind {
     ConsistentHash,
     // Bind the Partition to executor by partition.rand() order.
     Rand,
+    // Bind the Partition to executor by block-level modulo (block_idx % num_executors)
+    BlockMod,
     // Bind the Partition to executor by broadcast
     BroadcastCluster,
     // Bind the Partition to warehouse executor by broadcast
@@ -194,13 +196,16 @@ impl Partitions {
                 parts
             }
             // the executors will be all nodes in the warehouse if a query is BroadcastWarehouse.
-            PartitionsShuffleKind::BroadcastCluster | PartitionsShuffleKind::BroadcastWarehouse => {
+            PartitionsShuffleKind::BlockMod
+            | PartitionsShuffleKind::BroadcastCluster
+            | PartitionsShuffleKind::BroadcastWarehouse => {
                 return Ok(executors_sorted
                     .into_iter()
                     .map(|executor| {
                         (
                             executor.id.clone(),
-                            Partitions::create(PartitionsShuffleKind::Seq, self.partitions.clone()),
+                            // Preserve the original kind so executors know how to handle it
+                            Partitions::create(self.kind.clone(), self.partitions.clone()),
                         )
                     })
                     .collect());

src/query/catalog/tests/it/partitions.rs

@@ -254,14 +254,14 @@ fn test_partition_reshuffle() {
         writeln!(file, "{:?}", e2_parts.len()).unwrap();
     }
 
-    // Broadcast.
+    // BlockMod.
     {
-        let partitions = gen_parts(PartitionsShuffleKind::BroadcastCluster, 3);
+        let partitions = gen_parts(PartitionsShuffleKind::BlockMod, 3);
         let shuffle = partitions.reshuffle(executors_2.clone()).unwrap();
 
         writeln!(
             file,
-            "PartitionsShuffleKind::Broadcast: 3 partitions of 2 executors"
+            "PartitionsShuffleKind::BlockMod: 3 partitions of 2 executors"
         )
         .unwrap();
         let e1_parts = shuffle.get(&executors_2[0].id).unwrap();

src/query/catalog/tests/it/testdata/partition-reshuffle.txt

@@ -27,6 +27,6 @@ Partitions { kind: Seq, partitions: [{"type":"fuse_lazy","loc":"4"}, {"type":"fu
 PartitionsShuffleKind::Rand: 11 partitions of 2 executors
 5
 6
-PartitionsShuffleKind::Broadcast: 3 partitions of 2 executors
-Partitions { kind: Seq, partitions: [{"type":"fuse_lazy","loc":"0"}, {"type":"fuse_lazy","loc":"1"}, {"type":"fuse_lazy","loc":"2"}] }
-Partitions { kind: Seq, partitions: [{"type":"fuse_lazy","loc":"0"}, {"type":"fuse_lazy","loc":"1"}, {"type":"fuse_lazy","loc":"2"}] }
+PartitionsShuffleKind::BlockMod: 3 partitions of 2 executors
+Partitions { kind: BlockMod, partitions: [{"type":"fuse_lazy","loc":"0"}, {"type":"fuse_lazy","loc":"1"}, {"type":"fuse_lazy","loc":"2"}] }
+Partitions { kind: BlockMod, partitions: [{"type":"fuse_lazy","loc":"0"}, {"type":"fuse_lazy","loc":"1"}, {"type":"fuse_lazy","loc":"2"}] }

src/query/service/tests/it/storages/fuse/pruning_pipeline.rs

@@ -94,6 +94,7 @@ async fn apply_snapshot_pruning(
         cache_key,
         segment_locs.len(),
         0,
+        false, // use_block_level_shuffle
     )?;
     prune_pipeline.set_max_threads(1);
     prune_pipeline.set_on_init(move || {

src/query/settings/src/settings_default.rs

@@ -1314,6 +1314,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("auto_block_shuffle_threshold", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(5),
+                    desc: "Threshold for automatic block-level shuffle. When segment_count < cluster_nodes * threshold, uses block-level distribution (block_idx % num_executors) instead of segment-level. Set to 0 to disable.",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Both,
+                    range: Some(SettingRange::Numeric(0..=100)),
+                }),
                 ("enable_prune_pipeline", DefaultSettingValue {
                     value: UserSettingValue::UInt64(1),
                     desc: "Enable pruning pipeline",

src/query/settings/src/settings_getter_setter.rs

@@ -976,6 +976,10 @@ impl Settings {
         Ok(self.try_get_u64("enable_distributed_pruning")? == 1)
     }
 
+    pub fn get_auto_block_shuffle_threshold(&self) -> Result<usize> {
+        Ok(self.try_get_u64("auto_block_shuffle_threshold")? as usize)
+    }
+
     pub fn get_persist_materialized_cte(&self) -> Result<bool> {
         Ok(self.try_get_u64("persist_materialized_cte")? != 0)
     }

src/query/storages/basic/src/memory_table.rs

@@ -208,7 +208,7 @@ impl Table for MemoryTable {
         let parts = vec![MemoryPartInfo::create()];
         return Ok((
             statistics,
-            Partitions::create(PartitionsShuffleKind::BroadcastCluster, parts),
+            Partitions::create(PartitionsShuffleKind::BlockMod, parts),
         ));
     }
 

src/query/storages/fuse/src/operations/read_partitions.rs

@@ -55,6 +55,7 @@ use databend_storages_common_index::NgramArgs;
 use databend_storages_common_pruner::BlockMetaIndex;
 use databend_storages_common_pruner::TopNPruner;
 use databend_storages_common_table_meta::meta::BlockMeta;
+use databend_storages_common_table_meta::meta::BlockSlotDescription;
 use databend_storages_common_table_meta::meta::ColumnStatistics;
 use databend_storages_common_table_meta::meta::column_oriented_segment::BLOCK_SIZE;
 use databend_storages_common_table_meta::meta::column_oriented_segment::BLOOM_FILTER_INDEX_LOCATION;
@@ -169,6 +170,27 @@ impl FuseTable {
                         segments.push(FuseLazyPartInfo::create(idx, segment_location))
                     }
 
+                    // Determine shuffle kind based on heuristic:
+                    // Use block-level shuffle when segment count is small relative to cluster size
+                    let threshold = ctx.get_settings().get_auto_block_shuffle_threshold()?;
+                    let shuffle_kind =
+                        if threshold > 0 && nodes_num > 1 && segment_len < nodes_num * threshold {
+                            // Block-level shuffle: broadcast all segments to all executors,
+                            // each executor filters blocks by block_idx % num_executors == executor_idx
+                            info!(
+                                "Using BlockMod shuffle: segments={}, nodes={}, threshold={}",
+                                segment_len, nodes_num, threshold
+                            );
+                            PartitionsShuffleKind::BlockMod
+                        } else {
+                            // Default: use Mod shuffle kind for cache affinity at segment level
+                            info!(
+                                "Using Mod shuffle: segments={}, nodes={}, threshold={}",
+                                segment_len, nodes_num, threshold
+                            );
+                            PartitionsShuffleKind::Mod
+                        };
+
                     return Ok((
                         PartStatistics::new_estimated(
                             Some(snapshot_loc),
@@ -177,7 +199,7 @@ impl FuseTable {
                             segment_len,
                             segment_len,
                         ),
-                        Partitions::create(PartitionsShuffleKind::Mod, segments),
+                        Partitions::create(shuffle_kind, segments),
                     ));
                 }
 
@@ -284,6 +306,9 @@ impl FuseTable {
 
         let (segment_tx, segment_rx) = async_channel::bounded(max_io_requests);
 
+        // Check if we should use block-level shuffle based on the partition kind
+        let use_block_level_shuffle = plan.parts.kind == PartitionsShuffleKind::BlockMod;
+
         match segment_format {
             FuseSegmentFormat::Row => {
                 self.prune_segments_with_pipeline(
@@ -295,6 +320,7 @@ impl FuseTable {
                     derterministic_cache_key.clone(),
                     lazy_init_segments.len(),
                     plan_id,
+                    use_block_level_shuffle,
                 )?;
             }
             FuseSegmentFormat::Column => {
@@ -420,6 +446,7 @@ impl FuseTable {
         derterministic_cache_key: Option<String>,
         partitions_total: usize,
         plan_id: u32,
+        use_block_level_shuffle: bool,
     ) -> Result<()> {
         let max_threads = ctx.get_settings().get_max_threads()? as usize;
         prune_pipeline.add_source(
@@ -524,6 +551,26 @@ impl FuseTable {
             .filter(|p| p.order_by.is_empty() && p.filters.is_none())
             .and_then(|p| p.limit);
         let enable_prune_cache = ctx.get_settings().get_enable_prune_cache()?;
+
+        // Compute block slot for block-level shuffle
+        let block_slot = if use_block_level_shuffle {
+            let cluster = ctx.get_cluster();
+            if !cluster.is_empty() {
+                // Sort nodes by cache_id for deterministic ordering
+                let mut nodes = cluster.nodes.clone();
+                nodes.sort_by(|a, b| a.cache_id.cmp(&b.cache_id));
+                let num_slots = nodes.len();
+                let local_id = &cluster.local_id;
+                // Find the local node's position in the sorted list
+                let slot = nodes.iter().position(|n| &n.id == local_id).unwrap_or(0) as u32;
+                Some(BlockSlotDescription { num_slots, slot })
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
         let send_part_state = Arc::new(SendPartState::create(
             derterministic_cache_key,
             limit,
@@ -540,6 +587,7 @@ impl FuseTable {
                 pruner.table_schema.clone(),
                 send_part_state.clone(),
                 enable_prune_cache,
+                block_slot.clone(),
             )
         })?;
 

src/query/storages/fuse/src/pruning_pipeline/send_part_info_sink.rs

@@ -41,6 +41,7 @@ use databend_storages_common_cache::CacheAccessor;
 use databend_storages_common_cache::CachedObject;
 use databend_storages_common_pruner::BlockMetaIndex;
 use databend_storages_common_table_meta::meta::BlockMeta;
+use databend_storages_common_table_meta::meta::BlockSlotDescription;
 use parking_lot::Mutex;
 
 use crate::FuseTable;
@@ -134,6 +135,7 @@ pub struct SendPartInfoSink {
     statistics: PartStatistics,
     send_part_state: Arc<SendPartState>,
     enable_cache: bool,
+    block_slot: Option<BlockSlotDescription>,
 }
 
 impl SendPartInfoSink {
@@ -145,6 +147,7 @@ impl SendPartInfoSink {
         schema: TableSchemaRef,
         send_part_state: Arc<SendPartState>,
         enable_cache: bool,
+        block_slot: Option<BlockSlotDescription>,
     ) -> Result<ProcessorPtr> {
         let partitions = Partitions::default();
         let statistics = PartStatistics::default();
@@ -159,6 +162,7 @@ impl SendPartInfoSink {
                 statistics,
                 send_part_state,
                 enable_cache,
+                block_slot,
             },
         )))
     }
@@ -228,6 +232,13 @@ impl SendPartInfoSink {
         let mut parts = Vec::with_capacity(block_metas.len());
 
         for (block_meta_index, block_meta) in block_metas.iter() {
+            // Block-level partition shuffle: filter blocks by block_idx % num_slots == slot
+            if let Some(BlockSlotDescription { num_slots, slot }) = &self.block_slot {
+                if block_meta_index.block_idx % num_slots != *slot as usize {
+                    continue;
+                }
+            }
+
             let rows = block_meta.row_count as usize;
             let previous_limit = self.send_part_state.limit.fetch_sub(
                 rows.min(self.send_part_state.limit.load(Ordering::SeqCst)),
@@ -275,6 +286,13 @@ impl SendPartInfoSink {
         };
 
         for (block_meta_index, block_meta) in block_metas.iter() {
+            // Block-level partition shuffle: filter blocks by block_idx % num_slots == slot
+            if let Some(BlockSlotDescription { num_slots, slot }) = &self.block_slot {
+                if block_meta_index.block_idx % num_slots != *slot as usize {
+                    continue;
+                }
+            }
+
             let rows = block_meta.row_count as usize;
             let previous_limit = self.send_part_state.limit.fetch_sub(
                 rows.min(self.send_part_state.limit.load(Ordering::SeqCst)),