feat(cubestore): top-k merge strategies via CUBESTORE_TOPK_STRATEGY (cube-js#11152)

Author: waralexrom

rust/cubestore/cubestore/src/config/mod.rs

@@ -600,6 +600,9 @@ pub trait ConfigObj: DIService {
     /// partition coalesce (the hash aggregate ignores input order, so the per-row merge is wasted).
     fn coalesce_under_hash_aggregate(&self) -> bool;
 
+    /// Router-side merge strategy for distributed value-ordered top-k.
+    fn topk_aggregate_strategy(&self) -> TopKAggregateStrategy;
+
     fn allow_decimal128(&self) -> bool;
 
     fn enable_remove_orphaned_remote_files(&self) -> bool;
@@ -762,6 +765,7 @@ pub struct ConfigObjImpl {
     pub group_by_limit_factor: usize,
     pub group_by_limit_per_partition: bool,
     pub coalesce_under_hash_aggregate: bool,
+    pub topk_aggregate_strategy: TopKAggregateStrategy,
     pub allow_decimal128: bool,
     pub enable_remove_orphaned_remote_files: bool,
     pub enable_startup_warmup: bool,
@@ -1116,6 +1120,10 @@ impl ConfigObj for ConfigObjImpl {
         self.coalesce_under_hash_aggregate
     }
 
+    fn topk_aggregate_strategy(&self) -> TopKAggregateStrategy {
+        self.topk_aggregate_strategy
+    }
+
     fn allow_decimal128(&self) -> bool {
         self.allow_decimal128
     }
@@ -1291,6 +1299,43 @@ pub async fn init_test_logger() {
     }
 }
 
+/// Router-side merge strategy for distributed value-ordered top-k (`SELECT ... GROUP BY x ORDER BY
+/// agg(...) LIMIT k`). Selected by `CUBESTORE_TOPK_STRATEGY`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum TopKAggregateStrategy {
+    /// Original streaming NRA merge with per-row state (default).
+    Streaming,
+    /// Same streaming NRA merge (keeps early termination, bounded router memory), but vectorized.
+    VectorizedStreaming,
+    /// ClickHouse-style full re-aggregation on the router + fetch-limited sort. Drops early
+    /// termination, so the router materializes every distinct group.
+    FullMerge,
+}
+
+/// Parses [`TopKAggregateStrategy`] from an env var. Lenient: an unset or unrecognized value falls
+/// back to the default (`Streaming`) with a warning rather than panicking -- a malformed perf toggle
+/// must never take a node down.
+fn env_topk_strategy(name: &str) -> TopKAggregateStrategy {
+    match env::var(name) {
+        Err(_) => TopKAggregateStrategy::Streaming,
+        Ok(v) => match v.trim().to_ascii_lowercase().as_str() {
+            "" | "streaming" | "default" | "v1" => TopKAggregateStrategy::Streaming,
+            "vectorized" | "vectorized_streaming" | "v2" => {
+                TopKAggregateStrategy::VectorizedStreaming
+            }
+            "full_merge" | "full-merge" | "fullmerge" => TopKAggregateStrategy::FullMerge,
+            other => {
+                log::warn!(
+                    "unknown {} value '{}', using default (streaming)",
+                    name,
+                    other
+                );
+                TopKAggregateStrategy::Streaming
+            }
+        },
+    }
+}
+
 fn env_bool(name: &str, default: bool) -> bool {
     env::var(name)
         .ok()
@@ -1846,6 +1891,7 @@ impl Config {
                 group_by_limit_factor: env_parse_lenient("CUBESTORE_GROUP_BY_LIMIT_FACTOR", 0),
                 group_by_limit_per_partition: env_flag("CUBESTORE_GROUP_BY_LIMIT_PER_PARTITION"),
                 coalesce_under_hash_aggregate: env_flag("CUBESTORE_COALESCE_UNDER_HASH_AGGREGATE"),
+                topk_aggregate_strategy: env_topk_strategy("CUBESTORE_TOPK_STRATEGY"),
                 allow_decimal128: env_bool("CUBESTORE_ALLOW_DECIMAL128", false),
                 enable_remove_orphaned_remote_files: env_bool(
                     "CUBESTORE_ENABLE_REMOVE_ORPHANED_REMOTE_FILES",
@@ -2105,6 +2151,7 @@ impl Config {
                 group_by_limit_factor: 2,
                 group_by_limit_per_partition: false,
                 coalesce_under_hash_aggregate: false,
+                topk_aggregate_strategy: TopKAggregateStrategy::Streaming,
                 allow_decimal128: false,
                 enable_remove_orphaned_remote_files: false,
                 enable_startup_warmup: true,

rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs

@@ -8,6 +8,7 @@ mod trace_data_loaded;
 
 use super::serialized_plan::PreSerializedPlan;
 use crate::cluster::{Cluster, WorkerPlanningParams};
+use crate::config::TopKAggregateStrategy;
 use crate::queryplanner::optimizations::distributed_partial_aggregate::{
     add_limit_to_workers, drop_sort_merge_under_global_aggregate, ensure_partition_merge,
     push_aggregate_to_workers, push_sorted_partial_aggregate_below_merge,
@@ -45,6 +46,7 @@ pub struct CubeQueryPlanner {
     data_loaded_size: Option<Arc<DataLoadedSize>>,
     group_by_limit_factor: usize,
     group_by_limit_per_partition: bool,
+    topk_strategy: TopKAggregateStrategy,
 }
 
 impl CubeQueryPlanner {
@@ -54,6 +56,7 @@ impl CubeQueryPlanner {
         memory_handler: Arc<dyn MemoryHandler>,
         group_by_limit_factor: usize,
         group_by_limit_per_partition: bool,
+        topk_strategy: TopKAggregateStrategy,
     ) -> CubeQueryPlanner {
         CubeQueryPlanner {
             cluster: Some(cluster),
@@ -63,6 +66,7 @@ impl CubeQueryPlanner {
             data_loaded_size: None,
             group_by_limit_factor,
             group_by_limit_per_partition,
+            topk_strategy,
         }
     }
 
@@ -73,6 +77,7 @@ impl CubeQueryPlanner {
         data_loaded_size: Option<Arc<DataLoadedSize>>,
         group_by_limit_factor: usize,
         group_by_limit_per_partition: bool,
+        topk_strategy: TopKAggregateStrategy,
     ) -> CubeQueryPlanner {
         CubeQueryPlanner {
             serialized_plan,
@@ -82,6 +87,7 @@ impl CubeQueryPlanner {
             data_loaded_size,
             group_by_limit_factor,
             group_by_limit_per_partition,
+            topk_strategy,
         }
     }
 }
@@ -104,6 +110,7 @@ impl QueryPlanner for CubeQueryPlanner {
                 cluster: self.cluster.clone(),
                 worker_planning_params: self.worker_partition_count,
                 serialized_plan: self.serialized_plan.clone(),
+                topk_strategy: self.topk_strategy,
             }),
             Arc::new(RollingWindowPlanner {}),
         ])

rust/cubestore/cubestore/src/queryplanner/planning.rs

@@ -33,6 +33,7 @@ use std::any::Any;
 use std::fmt::Formatter;
 
 use crate::cluster::{Cluster, WorkerPlanningParams};
+use crate::config::TopKAggregateStrategy;
 use crate::metastore::multi_index::MultiPartition;
 use crate::metastore::table::{Table, TablePath};
 use crate::metastore::{
@@ -1970,6 +1971,7 @@ pub struct CubeExtensionPlanner {
     // Set on the workers.
     pub worker_planning_params: Option<WorkerPlanningParams>,
     pub serialized_plan: Arc<PreSerializedPlan>,
+    pub topk_strategy: TopKAggregateStrategy,
 }
 
 #[async_trait]

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -569,6 +569,7 @@ impl QueryExecutorImpl {
             self.memory_handler.clone(),
             self.config.group_by_limit_factor(),
             self.config.group_by_limit_per_partition(),
+            self.config.topk_aggregate_strategy(),
         ))
     }
 
@@ -586,6 +587,7 @@ impl QueryExecutorImpl {
             data_loaded_size.clone(),
             self.config.group_by_limit_factor(),
             self.config.group_by_limit_per_partition(),
+            self.config.topk_aggregate_strategy(),
         ))
     }