chore(query): add setting min_skip_partial_aggregation_capacity (#17040)

* chore: min_skip_partial_aggregation_capacity

* update

* update

Author: sundy-li

.cargo/audit.toml

@@ -47,7 +47,7 @@ ignore = [
     "RUSTSEC-2024-0351",
     # gix-path: improperly resolves configuration path reported by Git
     "RUSTSEC-2024-0371",
-    # Remotely exploitable Denial of Service in Tonic, ignored temporarily 
+    # Remotely exploitable Denial of Service in Tonic, ignored temporarily
     "RUSTSEC-2024-0376",
     #rustls network-reachable panic in `Acceptor::accept`
     "RUSTSEC-2024-0399",

src/query/expression/src/aggregate/mod.rs

@@ -89,24 +89,37 @@ impl HashTableConfig {
         self
     }
 
-    pub fn with_partial(mut self, partial_agg: bool, active_threads: usize) -> Self {
+    pub fn with_partial(
+        mut self,
+        partial_agg: bool,
+        active_threads: usize,
+        min_skip_partial_aggregation_capacity: usize,
+    ) -> Self {
         self.partial_agg = partial_agg;
 
         // init max_partial_capacity
         let total_shared_cache_size = active_threads * L3_CACHE_SIZE;
         let cache_per_active_thread =
             L1_CACHE_SIZE + L2_CACHE_SIZE + total_shared_cache_size / active_threads;
         let size_per_entry = (8_f64 * LOAD_FACTOR) as usize;
-        let capacity = (cache_per_active_thread / size_per_entry).next_power_of_two();
+        let capacity = (cache_per_active_thread / size_per_entry)
+            .next_power_of_two()
+            .max(min_skip_partial_aggregation_capacity);
         self.max_partial_capacity = capacity;
 
         self
     }
 
-    pub fn cluster_with_partial(mut self, partial_agg: bool, node_nums: usize) -> Self {
+    pub fn cluster_with_partial(
+        mut self,
+        partial_agg: bool,
+        node_nums: usize,
+        min_skip_partial_aggregation_capacity: usize,
+    ) -> Self {
         self.partial_agg = partial_agg;
         self.repartition_radix_bits_incr = 4;
-        self.max_partial_capacity = 131072 * (2 << node_nums);
+        self.max_partial_capacity =
+            (131072 * (2 << node_nums)).max(min_skip_partial_aggregation_capacity);
 
         self
     }

src/query/service/src/pipelines/builders/builder_aggregate.rs

@@ -127,13 +127,22 @@ impl PipelineBuilder {
         let schema_before_group_by = params.input_schema.clone();
         let sample_block = DataBlock::empty_with_schema(schema_before_group_by);
         let method = DataBlock::choose_hash_method(&sample_block, group_cols, efficiently_memory)?;
+        let min_skip_partial_aggregation_capacity =
+            self.settings.get_min_skip_partial_aggregation_capacity()? as usize;
 
         // Need a global atomic to read the max current radix bits hint
-        let partial_agg_config = if self.ctx.get_cluster().is_empty() {
-            HashTableConfig::default().with_partial(true, max_threads as usize)
+        let partial_agg_config = if !self.is_exchange_neighbor {
+            HashTableConfig::default().with_partial(
+                true,
+                max_threads as usize,
+                min_skip_partial_aggregation_capacity,
+            )
         } else {
-            HashTableConfig::default()
-                .cluster_with_partial(true, self.ctx.get_cluster().nodes.len())
+            HashTableConfig::default().cluster_with_partial(
+                true,
+                self.ctx.get_cluster().nodes.len(),
+                min_skip_partial_aggregation_capacity,
+            )
         };
 
         self.main_pipeline.add_transform(|input, output| {

src/query/settings/src/settings_default.rs

@@ -713,6 +713,12 @@ impl DefaultSettings {
                     mode: SettingMode::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("min_skip_partial_aggregation_capacity", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(0),
+                    desc: "The min size of aggregate hashtable skip partial aggregation capacity",
+                    mode: SettingMode::Both,
+                    range: Some(SettingRange::Numeric(0..=u64::MAX)),
+                }),
                 ("numeric_cast_option", DefaultSettingValue {
                     value: UserSettingValue::String("rounding".to_string()),
                     desc: "Set numeric cast mode as \"rounding\" or \"truncating\".",

src/query/settings/src/settings_getter_setter.rs

@@ -423,6 +423,10 @@ impl Settings {
         Ok(self.try_get_u64("enable_experimental_aggregate_hashtable")? == 1)
     }
 
+    pub fn get_min_skip_partial_aggregation_capacity(&self) -> Result<u64> {
+        self.try_get_u64("min_skip_partial_aggregation_capacity")
+    }
+
     pub fn get_lazy_read_threshold(&self) -> Result<u64> {
         self.try_get_u64("lazy_read_threshold")
     }