refactor(cubestore): collapse batch flag into strategy selector, concurrent download in merge

Remove CUBESTORE_BATCH_REPARTITION / batch_repartition_enabled. The repartition
strategy (per_chunk / per_partition / range) is now the sole selector:
- per_chunk: one job per chunk -> repartition_chunk
- per_partition: one anchor job per partition -> merge
- range: RepartitionRange jobs -> merge

The job handler and scheduler dispatch on the strategy directly. The dead
per_chunk+batch hybrid (per-chunk loop and its producer/consumer prefetch) is
gone, along with its two now-obsolete tests.

Prefetch is reworked: drop the byte-budget producer/consumer and
CUBESTORE_REPARTITION_PREFETCH_BUDGET; add a plain bool
CUBESTORE_REPARTITION_CONCURRENT_DOWNLOAD (default off) that downloads a merge
group's chunk parquets concurrently before building the merge inputs. It applies
to both per_partition and range; the group is already bounded by
repartition_merge_max_input_files and the pool by download_concurrency, so no
extra budget is needed.

Author: waralexrom

rust/cubestore/cubestore/src/cluster/ingestion/job_processor.rs

@@ -1,5 +1,5 @@
 use crate::config::injection::DIService;
-use crate::config::{Config, ConfigObj};
+use crate::config::{Config, ConfigObj, RepartitionStrategy};
 use crate::import::ImportService;
 use crate::metastore::job::{Job, JobType};
 use crate::metastore::table::Table;
@@ -222,17 +222,14 @@ impl JobIsolatedProcessor {
                     }
                     let data_loaded_size = DataLoadedSize::new();
                     app_metrics::JOBS_REPARTITION_CHUNK.add(1);
-                    // FIXME: a RepartitionChunk job whose chunk_id is used as an
-                    // anchor for the whole partition (batch_repartition_enabled).
-                    // We deliberately overload the existing RepartitionChunk type
-                    // instead of introducing a dedicated per-partition JobType:
-                    // a new JobType variant cannot be deserialized by an older
-                    // binary, which would make the whole job shard unreadable when
-                    // a deployment switches between the `latest` and `release`
-                    // channels (version can move both ways at any time). Reusing a
-                    // known type keeps both directions safe — an old binary just
-                    // repartitions the single anchor chunk. See repartition_partition_chunks.
-                    let r = if self.config_obj.batch_repartition_enabled() {
+                    // PerPartition reuses the RepartitionChunk job with the smallest chunk
+                    // as the anchor (one job per partition) and merges all its chunks;
+                    // PerChunk gets one job per chunk and repartitions just that chunk. We
+                    // overload the existing RepartitionChunk type (no new JobType) so an
+                    // older binary stays able to deserialize it across channel switches.
+                    let r = if self.config_obj.repartition_strategy()
+                        == RepartitionStrategy::PerPartition
+                    {
                         let partition_id = chunk.get_row().get_partition_id();
                         let time_budget = Duration::from_secs(
                             self.config_obj.repartition_chunks_time_budget_secs(),

rust/cubestore/cubestore/src/cluster/mod.rs

@@ -1002,20 +1002,14 @@ impl Cluster for ClusterImpl {
                     }
                 }
             }
-        } else if self.config_obj.repartition_strategy() == RepartitionStrategy::PerPartition
-            || self.config_obj.batch_repartition_enabled()
-        {
-            // One job per partition that batches all persisted chunks, but keyed on a
-            // chunk (RepartitionChunk), not the partition. We reuse the existing job
-            // type instead of a dedicated per-partition JobType so an older binary stays
-            // able to deserialize it across `latest`/`release` channel switches (a new
-            // variant would make its whole job shard unreadable). The anchor is the
-            // smallest persisted chunk id so add_job dedups to a single job per
-            // partition; the worker resolves the partition from it and processes the
-            // anchor last (see repartition_partition_chunks). The PerPartition strategy
-            // requires this single-job form: a per-chunk job under it would re-merge the
-            // whole partition. An old binary just repartitions this one chunk and drains
-            // the rest via its own per-chunk path.
+        } else if self.config_obj.repartition_strategy() == RepartitionStrategy::PerPartition {
+            // One job per partition, keyed on a chunk (RepartitionChunk) rather than a
+            // dedicated per-partition JobType so an older binary can still deserialize it
+            // across `latest`/`release` channel switches. The anchor is the smallest
+            // persisted chunk id so add_job dedups to a single job per partition; the
+            // worker resolves the partition from it and merges all its chunks (see
+            // repartition_partition_chunks). An old binary just repartitions this one
+            // chunk and drains the rest via its own per-chunk path.
             if let Some(anchor_chunk_id) = chunks.iter().map(|c| c.get_id()).min() {
                 let node = self.node_name_by_partition(p);
                 let job = self

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

@@ -540,11 +540,6 @@ pub trait ConfigObj: DIService {
     /// default (hash placement); kept as a flag for rollout/rollback.
     fn load_aware_import_placement_enabled(&self) -> bool;
 
-    /// When enabled, an inactive parent partition is repartitioned by a single
-    /// per-partition job that loops over its persisted chunks, instead of one
-    /// job per chunk. Kept as a flag for emergency rollback.
-    fn batch_repartition_enabled(&self) -> bool;
-
     /// Time budget for a single batch repartition job. The job yields its runner
     /// slot once the budget is exceeded (after the current chunk); the remainder
     /// is drained by a follow-up job via the cascade. Kept separate from
@@ -567,13 +562,11 @@ pub trait ConfigObj: DIService {
     /// Off by default; when disabled, chunks are sent whole and filtered only in the subprocess.
     fn prefilter_in_memory_chunks_enabled(&self) -> bool;
 
-    /// Byte budget for prefetching persisted chunk parquets ahead of a batch
-    /// repartition job. A sequential producer downloads upcoming chunks (anchor
-    /// last) while the job processes the current one, bounded so that no more
-    /// than this many bytes of fetched-but-unprocessed data sit on local disk.
-    /// The env value accepts size suffixes (e.g. `512MB`). `None` (env unset) or
-    /// `Some(0)` disables prefetching.
-    fn repartition_prefetch_budget_bytes(&self) -> Option<u64>;
+    /// When enabled, a merge group's chunk parquets (PerPartition / Range strategies) are
+    /// downloaded concurrently before building the merge inputs, instead of one at a time.
+    /// The group is already bounded by repartition_merge_max_input_files and the download
+    /// pool by download_concurrency, so no extra budget is needed. Off by default.
+    fn repartition_concurrent_download(&self) -> bool;
 
     /// Which repartition strategy to use for an inactive parent's persisted chunks.
     /// Defaults to PerChunk.
@@ -733,12 +726,11 @@ pub struct ConfigObjImpl {
     pub upload_to_remote: bool,
     pub enable_topk: bool,
     pub load_aware_import_placement_enabled: bool,
-    pub batch_repartition_enabled: bool,
     pub repartition_chunks_time_budget_secs: u64,
     pub push_partial_aggregate_below_merge_enabled: bool,
     pub compaction_split_by_total_file_size_enabled: bool,
     pub prefilter_in_memory_chunks_enabled: bool,
-    pub repartition_prefetch_budget_bytes: Option<u64>,
+    pub repartition_concurrent_download: bool,
     pub repartition_strategy: RepartitionStrategy,
     pub repartition_merge_max_input_files: usize,
     pub repartition_merge_max_rows: u64,
@@ -1054,9 +1046,6 @@ impl ConfigObj for ConfigObjImpl {
     fn load_aware_import_placement_enabled(&self) -> bool {
         self.load_aware_import_placement_enabled
     }
-    fn batch_repartition_enabled(&self) -> bool {
-        self.batch_repartition_enabled
-    }
     fn repartition_chunks_time_budget_secs(&self) -> u64 {
         self.repartition_chunks_time_budget_secs
     }
@@ -1069,8 +1058,8 @@ impl ConfigObj for ConfigObjImpl {
     fn prefilter_in_memory_chunks_enabled(&self) -> bool {
         self.prefilter_in_memory_chunks_enabled
     }
-    fn repartition_prefetch_budget_bytes(&self) -> Option<u64> {
-        self.repartition_prefetch_budget_bytes
+    fn repartition_concurrent_download(&self) -> bool {
+        self.repartition_concurrent_download
     }
     fn repartition_strategy(&self) -> RepartitionStrategy {
         self.repartition_strategy
@@ -1737,7 +1726,6 @@ impl Config {
                     "CUBESTORE_LOAD_AWARE_IMPORT_PLACEMENT",
                     false,
                 ),
-                batch_repartition_enabled: env_bool("CUBESTORE_BATCH_REPARTITION", false),
                 repartition_chunks_time_budget_secs: env_parse(
                     "CUBESTORE_REPARTITION_TIME_BUDGET_SECS",
                     60,
@@ -1754,12 +1742,10 @@ impl Config {
                     "CUBESTORE_PREFILTER_IN_MEMORY_CHUNKS",
                     false,
                 ),
-                repartition_prefetch_budget_bytes: env_parse_optional_size(
-                    "CUBESTORE_REPARTITION_PREFETCH_BUDGET",
-                    None,
-                    None,
-                )
-                .map(|n| n as u64),
+                repartition_concurrent_download: env_bool(
+                    "CUBESTORE_REPARTITION_CONCURRENT_DOWNLOAD",
+                    false,
+                ),
                 repartition_strategy: env_repartition_strategy(),
                 repartition_merge_max_input_files: env_parse(
                     "CUBESTORE_REPARTITION_MERGE_MAX_INPUT_FILES",
@@ -2007,14 +1993,13 @@ impl Config {
                 upload_to_remote: true,
                 enable_topk: true,
                 load_aware_import_placement_enabled: false,
-                batch_repartition_enabled: true,
                 repartition_chunks_time_budget_secs: 60,
                 push_partial_aggregate_below_merge_enabled: true,
                 compaction_split_by_total_file_size_enabled: false,
                 // Production default is off; kept on in tests so prefilter_chunks_shared_scan
                 // and the rest of the suite keep exercising the worker-side trim path.
                 prefilter_in_memory_chunks_enabled: true,
-                repartition_prefetch_budget_bytes: None,
+                repartition_concurrent_download: false,
                 repartition_strategy: RepartitionStrategy::PerChunk,
                 repartition_merge_max_input_files: 50,
                 repartition_merge_max_rows: 4_000_000,

rust/cubestore/cubestore/src/sql/mod.rs

@@ -3605,7 +3605,7 @@ mod tests {
             .update_config(|mut c| {
                 c.partition_split_threshold = 20;
                 c.compaction_chunks_count_threshold = 10;
-                c.batch_repartition_enabled = false;
+                c.repartition_strategy = crate::config::RepartitionStrategy::PerChunk;
                 c
             })
             .start_test(async move |services| {
@@ -3616,16 +3616,16 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn repartition_prefetch_keeps_data_consistent() -> Result<(), CubeError> {
-        // Batch repartition with chunk parquet prefetch enabled must drain and
-        // keep data consistent end-to-end (real downloads through the prefetch
-        // producer/consumer path).
-        Config::test("repartition_prefetch_keeps_data_consistent")
+    async fn repartition_concurrent_download_keeps_data_consistent() -> Result<(), CubeError> {
+        // PerPartition merge with concurrent chunk download enabled must drain and keep
+        // data consistent end-to-end (real concurrent downloads in the merge group build).
+        Config::test("repartition_concurrent_download_keeps_data_consistent")
             .update_config(|mut c| {
                 c.partition_split_threshold = 20;
                 c.compaction_chunks_count_threshold = 10;
-                c.batch_repartition_enabled = true;
-                c.repartition_prefetch_budget_bytes = Some(64 * 1024 * 1024);
+                c.repartition_strategy = crate::config::RepartitionStrategy::PerPartition;
+                c.repartition_merge_max_input_files = 4;
+                c.repartition_concurrent_download = true;
                 c
             })
             .start_test(async move |services| {
@@ -3721,7 +3721,8 @@ mod tests {
             .update_config(|mut c| {
                 c.partition_split_threshold = 20;
                 c.compaction_chunks_count_threshold = 10;
-                c.batch_repartition_enabled = true;
+                c.repartition_strategy = crate::config::RepartitionStrategy::PerPartition;
+                c.repartition_merge_max_input_files = 2;
                 c.repartition_chunks_time_budget_secs = 1;
                 c
             })