fix: preserve range partitions when scaling file scans

src/events/defaults/file_scan_config.rs

@@ -40,33 +40,38 @@ pub(crate) fn file_scan_config_scale_up_leaf_node(
     let file_scan = dse.data_source().downcast_ref::<FileScanConfig>()?;
     let partition_count = ev.plan.output_partitioning().partition_count();
 
-    let rebalanced = if file_scan.output_partitioning.is_some() {
-        let all_partitioned_files = file_scan
-            .file_groups
-            .iter()
-            .flat_map(|file_group| file_group.iter().cloned())
-            .collect::<Vec<_>>();
-        rebalance_round_robin(all_partitioned_files, partition_count * ev.task_count)
+    let file_scans = if file_scan.output_partitioning.is_some() {
+        // File groups are the declared hash/range partitions. Moving a file to a
+        // different group would invalidate that mapping — DataFusion's
+        // FileScanConfig::repartitioned also refuses to do so. Keep group i as
+        // partition i on every task, and assign each group to a single task so
+        // co-partitioned joins and single-partitioned aggregates stay correct.
+        scale_declared_partition_file_groups(&file_scan.file_groups, ev.task_count)
             .into_iter()
-            .map(FileGroup::new)
+            .map(|file_groups| {
+                let mut cfg = file_scan.clone();
+                cfg.file_groups = file_groups;
+                cfg
+            })
             .collect::<Vec<_>>()
     } else {
-        FileGroupPartitioner::new()
+        let rebalanced = FileGroupPartitioner::new()
             .with_target_partitions(partition_count * ev.task_count)
             .with_repartition_file_min_size(0)
             .with_preserve_order_within_groups(!file_scan.output_ordering.is_empty())
             .repartition_file_groups(&file_scan.file_groups)
             .unwrap_or_else(|| file_scan.file_groups.clone())
             .into_iter()
-            .collect()
-    };
+            .collect::<Vec<_>>();
 
-    let mut file_scan_template = file_scan.clone();
-    file_scan_template.file_groups.clear();
-    let mut file_scans = vec![file_scan_template; ev.task_count];
-    for (i, file_group) in rebalanced.into_iter().enumerate() {
-        file_scans[i % ev.task_count].file_groups.push(file_group);
-    }
+        let mut file_scan_template = file_scan.clone();
+        file_scan_template.file_groups.clear();
+        let mut file_scans = vec![file_scan_template; ev.task_count];
+        for (i, file_group) in rebalanced.into_iter().enumerate() {
+            file_scans[i % ev.task_count].file_groups.push(file_group);
+        }
+        file_scans
+    };
 
     let distributed_leaf_result = DistributedLeafExec::try_new(
         Arc::clone(ev.plan),
@@ -84,14 +89,18 @@ pub(crate) fn file_scan_config_scale_up_leaf_node(
     ))))
 }
 
-fn rebalance_round_robin<T>(items: Vec<T>, target_groups: usize) -> Vec<Vec<T>> {
-    let mut groups = (0..target_groups)
-        .map(|_| Vec::new())
-        .collect::<Vec<Vec<T>>>();
-    for (idx, item) in items.into_iter().enumerate() {
-        groups[idx % target_groups].push(item);
+fn scale_declared_partition_file_groups(
+    file_groups: &[FileGroup],
+    task_count: usize,
+) -> Vec<Vec<FileGroup>> {
+    let mut per_task = vec![vec![FileGroup::new(vec![]); file_groups.len()]; task_count];
+    for (part_idx, group) in file_groups.iter().enumerate() {
+        if group.is_empty() {
+            continue;
+        }
+        per_task[part_idx % task_count][part_idx] = group.clone();
     }
-    groups
+    per_task
 }
 
 #[cfg(test)]
@@ -156,6 +165,81 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_scale_up_preserves_range_file_group_slots() -> Result<(), DataFusionError> {
+        let plan = make_range_data_source_exec(4)?;
+        let cfg = SessionConfig::new();
+        let response = file_scan_config_scale_up_leaf_node(ScaleUpLeafNodeEvent {
+            plan: &plan,
+            task_count: 2,
+            session_config: &cfg,
+        })
+        .expect("a file scan should be recognized")?;
+        let leaf = response
+            .plan
+            .downcast_ref::<DistributedLeafExec>()
+            .expect("scale-up should return DistributedLeafExec");
+
+        assert_eq!(leaf.variants().len(), 2);
+        let slots = variant_file_slots(leaf);
+        // Each original range partition stays in its own slot and is assigned
+        // to exactly one task. Flattening files and rebalancing would put
+        // part-2 in slot 1 of task 0.
+        assert_eq!(
+            slots,
+            vec![
+                vec![vec!["part-0"], vec![], vec!["part-2"], vec![]],
+                vec![vec![], vec!["part-1"], vec![], vec!["part-3"]],
+            ]
+        );
+        assert!(matches!(
+            leaf.variants()[0].output_partitioning(),
+            datafusion::physical_expr::Partitioning::Range(range)
+                if range.partition_count() == 4
+        ));
+        Ok(())
+    }
+
+    #[test]
+    fn test_scale_up_keeps_all_files_of_a_range_partition_on_one_task()
+    -> Result<(), DataFusionError> {
+        let plan = make_range_data_source_exec_with_files(&[
+            &["p0-a", "p0-b"],
+            &["p1-a"],
+            &["p2-a", "p2-b", "p2-c"],
+        ])?;
+        let cfg = SessionConfig::new();
+        let response = file_scan_config_scale_up_leaf_node(ScaleUpLeafNodeEvent {
+            plan: &plan,
+            task_count: 2,
+            session_config: &cfg,
+        })
+        .expect("a file scan should be recognized")?;
+        let leaf = response
+            .plan
+            .downcast_ref::<DistributedLeafExec>()
+            .expect("scale-up should return DistributedLeafExec");
+
+        assert_eq!(
+            variant_file_slots(leaf),
+            vec![
+                vec![vec!["p0-a", "p0-b"], vec![], vec!["p2-a", "p2-b", "p2-c"]],
+                vec![vec![], vec!["p1-a"], vec![]],
+            ]
+        );
+        Ok(())
+    }
+
+    fn rebalance_round_robin<T>(items: Vec<T>, target_groups: usize) -> Vec<Vec<T>> {
+        let mut groups = (0..target_groups)
+            .map(|_| Vec::new())
+            .collect::<Vec<Vec<T>>>();
+        for (idx, item) in items.into_iter().enumerate() {
+            groups[idx % target_groups].push(item);
+        }
+        groups
+    }
+
     #[test]
     fn test_rebalance_round_robin_fixes_group_boundary_skew() {
         let groups = rebalance_round_robin((0..8).collect(), 5);
@@ -185,6 +269,92 @@ mod tests {
             .sum()
     }
 
+    fn variant_file_slots(leaf: &DistributedLeafExec) -> Vec<Vec<Vec<String>>> {
+        leaf.variants()
+            .iter()
+            .map(|variant| {
+                let dse = variant.downcast_ref::<DataSourceExec>().unwrap();
+                let file_scan = dse.data_source().downcast_ref::<FileScanConfig>().unwrap();
+                file_scan
+                    .file_groups
+                    .iter()
+                    .map(|group| {
+                        group
+                            .files()
+                            .iter()
+                            .map(|file| {
+                                file.object_meta
+                                    .location
+                                    .filename()
+                                    .unwrap_or(file.object_meta.location.as_ref())
+                                    .to_string()
+                            })
+                            .collect()
+                    })
+                    .collect()
+            })
+            .collect()
+    }
+
+    fn make_range_data_source_exec(
+        range_partitions: usize,
+    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        let files = (0..range_partitions)
+            .map(|i| vec![format!("part-{i}")])
+            .collect::<Vec<_>>();
+        let file_refs = files
+            .iter()
+            .map(|files| files.iter().map(|s| s.as_str()).collect::<Vec<_>>())
+            .collect::<Vec<_>>();
+        let file_refs = file_refs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
+        make_range_data_source_exec_with_files(&file_refs)
+    }
+
+    fn make_range_data_source_exec_with_files(
+        files_per_partition: &[&[&str]],
+    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        use datafusion::common::{ScalarValue, SplitPoint};
+        use datafusion::datasource::listing::PartitionedFile;
+        use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource};
+        use datafusion::execution::object_store::ObjectStoreUrl;
+        use datafusion::physical_expr::{PhysicalSortExpr, RangePartitioning, expressions::col};
+        use datafusion::physical_expr_common::sort_expr::LexOrdering;
+
+        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+            arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, false),
+        ]));
+        let file_groups = files_per_partition
+            .iter()
+            .map(|files| {
+                FileGroup::new(
+                    files
+                        .iter()
+                        .map(|name| PartitionedFile::new(name.to_string(), 1024))
+                        .collect(),
+                )
+            })
+            .collect::<Vec<_>>();
+        let split_points = (1..files_per_partition.len())
+            .map(|i| SplitPoint::new(vec![ScalarValue::Int32(Some(i as i32 * 10))]))
+            .collect::<Vec<_>>();
+        let ordering =
+            LexOrdering::new([PhysicalSortExpr::new_default(col("k", schema.as_ref())?)]);
+        let Some(ordering) = ordering else {
+            return Err(DataFusionError::Internal(
+                "range ordering must not be empty".to_string(),
+            ));
+        };
+        let range = RangePartitioning::try_new(ordering, split_points)?;
+        let config = FileScanConfigBuilder::new(
+            ObjectStoreUrl::local_filesystem(),
+            Arc::new(ParquetSource::new(Arc::clone(&schema))),
+        )
+        .with_file_groups(file_groups)
+        .with_output_partitioning(Some(datafusion::physical_expr::Partitioning::Range(range)))
+        .build();
+        Ok(DataSourceExec::from_data_source(config))
+    }
+
     async fn make_data_source_exec() -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         let ctx = SessionContext::new();
         register_parquet_tables(&ctx).await?;

src/execution_plans/common.rs

@@ -22,6 +22,59 @@ pub(super) fn scale_partitioning(
         Partitioning::RoundRobinBatch(p) => Partitioning::RoundRobinBatch(f(*p)),
         Partitioning::Hash(hash, p) => Partitioning::Hash(hash.clone(), f(*p)),
         Partitioning::UnknownPartitioning(p) => Partitioning::UnknownPartitioning(f(*p)),
-        Partitioning::Range(range) => Partitioning::Range(range.clone()),
+        Partitioning::Range(range) => {
+            // Range partition count is defined by the split points. Changing the
+            // count (for example concatenating several range-partitioned tasks in
+            // NetworkCoalesceExec) does not produce a valid Range of the new size,
+            // so keep Range only when the count is unchanged.
+            let new_count = f(range.partition_count());
+            if new_count == range.partition_count() {
+                Partitioning::Range(range.clone())
+            } else {
+                Partitioning::UnknownPartitioning(new_count)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::common::{ScalarValue, SplitPoint};
+    use datafusion::physical_expr::{PhysicalSortExpr, RangePartitioning, expressions::col};
+    use datafusion::physical_expr_common::sort_expr::LexOrdering;
+
+    fn range_partitioning(split_count: usize) -> Partitioning {
+        let schema = Schema::new(vec![Field::new("k", DataType::Int32, false)]);
+        let ordering =
+            LexOrdering::new([PhysicalSortExpr::new_default(col("k", &schema).unwrap())]).unwrap();
+        let split_points = (1..=split_count)
+            .map(|i| SplitPoint::new(vec![ScalarValue::Int32(Some(i as i32 * 10))]))
+            .collect();
+        Partitioning::Range(RangePartitioning::try_new(ordering, split_points).unwrap())
+    }
+
+    #[test]
+    fn scale_partitioning_preserves_range_when_count_is_unchanged() {
+        let range = range_partitioning(2);
+        let scaled = scale_partitioning(&range, |p| p);
+        assert_eq!(scaled, range);
+        assert_eq!(scaled.partition_count(), 3);
+    }
+
+    #[test]
+    fn scale_partitioning_degrades_range_when_count_changes() {
+        let range = range_partitioning(2);
+        let scaled = scale_partitioning(&range, |p| p * 2);
+        assert!(matches!(scaled, Partitioning::UnknownPartitioning(6)));
+    }
+
+    #[test]
+    fn scale_partitioning_still_scales_hash() {
+        let schema = Schema::new(vec![Field::new("k", DataType::Int32, false)]);
+        let hash = Partitioning::Hash(vec![col("k", &schema).unwrap()], 4);
+        let scaled = scale_partitioning(&hash, |p| p * 2);
+        assert!(matches!(scaled, Partitioning::Hash(_, 8)));
     }
 }

src/stage.rs

@@ -611,13 +611,17 @@ fn metrics_by_task_id(metrics: &MetricsSet) -> HashMap<usize, MetricsSet> {
 fn format_tasks_for_stage(n_tasks: usize, head: &Arc<dyn ExecutionPlan>) -> String {
     let partitioning = head.properties().output_partitioning();
     let input_partitions = partitioning.partition_count();
-    let hash_shuffle = matches!(partitioning, Partitioning::Hash(_, _));
-    // In a hash shuffle every task reads the same partition range, so the stage spans
-    // `input_partitions` distinct partitions. Otherwise each task owns its own slice, for a total
-    // of `n_tasks * input_partitions`.
-    let partitions = match hash_shuffle {
-        true => input_partitions,
-        false => n_tasks * input_partitions,
+    // Hash and Range both assign the same partition index to the same key across
+    // tasks, so the stage spans `input_partitions` distinct partitions. Otherwise
+    // each task owns its own slice, for a total of `n_tasks * input_partitions`.
+    let aligned_partitions = matches!(
+        partitioning,
+        Partitioning::Hash(_, _) | Partitioning::Range(_)
+    );
+    let partitions = if aligned_partitions {
+        input_partitions
+    } else {
+        n_tasks * input_partitions
     };
     format!("tasks={n_tasks}, partitions={partitions}")
 }

src/test_utils/insta.rs

@@ -16,8 +16,10 @@ pub fn settings() -> insta::Settings {
     unsafe { env::set_var("INSTA_WORKSPACE_ROOT", env!("CARGO_MANIFEST_DIR")) };
     let mut settings = insta::Settings::clone_current();
     let cwd = env::current_dir().unwrap();
-    let cwd = cwd.to_str().unwrap();
-    settings.add_filter(cwd.trim_start_matches("/"), "");
+    // Plans use forward slashes even on Windows. Escape the path so drive
+    // letters and backslashes are not treated as regex syntax.
+    let cwd = cwd.to_string_lossy().replace('\\', "/");
+    settings.add_filter(&regex_escape(cwd.trim_start_matches('/')), "");
     settings.add_filter(
         r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
         "UUID",
@@ -32,3 +34,17 @@ pub fn settings() -> insta::Settings {
     );
     settings
 }
+
+fn regex_escape(input: &str) -> String {
+    let mut escaped = String::with_capacity(input.len());
+    for c in input.chars() {
+        if matches!(
+            c,
+            '\\' | '.' | '+' | '*' | '?' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '^' | '$'
+        ) {
+            escaped.push('\\');
+        }
+        escaped.push(c);
+    }
+    escaped
+}