From e0639ca32cca798910c5e206e9b1479a2bfa4683 Mon Sep 17 00:00:00 2001 From: Nishchaya Sharma <129059557+shinzoxD@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:21:15 +0530 Subject: [PATCH 1/2] fix: preserve range partitions when scaling file scans It is not valid to claim a listing scan is hash-partitioned. The join tests now declare Range output partitioning on the hive-layout testdata instead. file_scan_config_scale_up_leaf_node no longer flattens those file groups. FileScanConfig::repartitioned already refuses to rearrange them, because moving a file to another slot breaks the range (or hash) claim and can produce wrong partitioned joins and SinglePartitioned aggregates. Each group stays in its original slot and is assigned to one task. scale_partitioning keeps Range only when the partition count is unchanged. Concatenating several range-partitioned tasks, as NetworkCoalesceExec does, is not a valid Range of the new size, so it degrades to UnknownPartitioning. Stage display treats Range like Hash: every task shares the same partition indices. Also escape Windows paths in snapshot filters so cwd is not parsed as a regex. Fixes #628 --- src/events/defaults/file_scan_config.rs | 218 +++++++++++++++++++++--- src/execution_plans/common.rs | 55 +++++- src/stage.rs | 18 +- src/test_utils/insta.rs | 20 ++- tests/join.rs | 113 +++++++----- tests/range_partitioned.rs | 175 +++++++++++++++++++ 6 files changed, 525 insertions(+), 74 deletions(-) create mode 100644 tests/range_partitioned.rs diff --git a/src/events/defaults/file_scan_config.rs b/src/events/defaults/file_scan_config.rs index f81bd1e60..52cf952d7 100644 --- a/src/events/defaults/file_scan_config.rs +++ b/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::()?; 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::>(); - 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::>() } 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::>(); - 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(items: Vec, target_groups: usize) -> Vec> { - let mut groups = (0..target_groups) - .map(|_| Vec::new()) - .collect::>>(); - 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> { + 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::() + .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::() + .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(items: Vec, target_groups: usize) -> Vec> { + let mut groups = (0..target_groups) + .map(|_| Vec::new()) + .collect::>>(); + 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>> { + leaf.variants() + .iter() + .map(|variant| { + let dse = variant.downcast_ref::().unwrap(); + let file_scan = dse.data_source().downcast_ref::().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, DataFusionError> { + let files = (0..range_partitions) + .map(|i| vec![format!("part-{i}")]) + .collect::>(); + let file_refs = files + .iter() + .map(|files| files.iter().map(|s| s.as_str()).collect::>()) + .collect::>(); + let file_refs = file_refs.iter().map(|v| v.as_slice()).collect::>(); + make_range_data_source_exec_with_files(&file_refs) + } + + fn make_range_data_source_exec_with_files( + files_per_partition: &[&[&str]], + ) -> Result, 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::>(); + let split_points = (1..files_per_partition.len()) + .map(|i| SplitPoint::new(vec![ScalarValue::Int32(Some(i as i32 * 10))])) + .collect::>(); + 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, DataFusionError> { let ctx = SessionContext::new(); register_parquet_tables(&ctx).await?; diff --git a/src/execution_plans/common.rs b/src/execution_plans/common.rs index 8af615d98..2a2d1f7a5 100644 --- a/src/execution_plans/common.rs +++ b/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))); } } diff --git a/src/stage.rs b/src/stage.rs index cd054c575..a60496632 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -611,13 +611,17 @@ fn metrics_by_task_id(metrics: &MetricsSet) -> HashMap { fn format_tasks_for_stage(n_tasks: usize, head: &Arc) -> 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}") } diff --git a/src/test_utils/insta.rs b/src/test_utils/insta.rs index 67d2a7977..ccaf622c7 100644 --- a/src/test_utils/insta.rs +++ b/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(®ex_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 +} diff --git a/tests/join.rs b/tests/join.rs index 096d7e82f..8da4694d9 100644 --- a/tests/join.rs +++ b/tests/join.rs @@ -6,24 +6,25 @@ mod tests { util::pretty::{self, pretty_format_batches}, }; use datafusion::{ + common::{ScalarValue, SplitPoint}, + datasource::{ + file_format::parquet::ParquetFormat, + listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, + }, error::Result, + logical_expr::{Partitioning as LogicalPartitioning, RangePartitioning}, physical_plan::collect, - prelude::{ParquetReadOptions, SessionContext, col}, + prelude::{SessionContext, col}, }; use datafusion_distributed::{ DefaultSessionBuilder, assert_snapshot, display_plan_ascii, test_utils::localhost::start_localhost_context, }; + use std::sync::Arc; fn set_configs(ctx: &mut SessionContext) { - // Preserve hive-style file partitions. - ctx.state_ref() - .write() - .config_mut() - .options_mut() - .optimizer - .preserve_file_partitions = 1; - // Read data from 4 hive-style partitions. + // Read data from 4 range partitions. Declared range output partitioning + // uses the range partition count, not target_partitions, for file groups. ctx.state_ref() .write() .config_mut() @@ -45,22 +46,54 @@ mod tests { .hash_join_single_partition_threshold_rows = 0; } - async fn register_tables(ctx: &SessionContext) -> Result<()> { - // Register hive-style partitioning for the dim table. - let dim_options = ParquetReadOptions::default() - .table_partition_cols(vec![("d_dkey".to_string(), DataType::Utf8)]); - ctx.register_parquet("dim", "testdata/join/parquet/dim", dim_options) - .await?; + fn range_on(column: &str) -> Result { + // testdata/join hive folders are A,B,C,D in path order, which matches + // these split points: A < B, B < C, C < D. + Ok(LogicalPartitioning::Range(RangePartitioning::try_new( + vec![col(column).sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Utf8(Some("B".to_string()))]), + SplitPoint::new(vec![ScalarValue::Utf8(Some("C".to_string()))]), + SplitPoint::new(vec![ScalarValue::Utf8(Some("D".to_string()))]), + ], + )?)) + } + + async fn register_range_table( + ctx: &SessionContext, + name: &str, + path: &str, + partition_col: &str, + file_sort_order: Vec>, + ) -> Result<()> { + let table_path = ListingTableUrl::parse(path)?; + let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_table_partition_cols(vec![(partition_col.to_string(), DataType::Utf8)]) + .with_file_sort_order(file_sort_order) + .with_output_partitioning(Some(range_on(partition_col)?)); + let schema = options.infer_schema(&ctx.state(), &table_path).await?; + let table = ListingTable::try_new( + ListingTableConfig::new(table_path) + .with_listing_options(options) + .with_schema(schema), + )?; + ctx.register_table(name, Arc::new(table))?; + Ok(()) + } - // Register hive-style partitioning for the fact table. - let fact_options = ParquetReadOptions::default() - .table_partition_cols(vec![("f_dkey".to_string(), DataType::Utf8)]) - .file_sort_order(vec![vec![ + async fn register_tables(ctx: &SessionContext) -> Result<()> { + register_range_table(ctx, "dim", "testdata/join/parquet/dim", "d_dkey", vec![]).await?; + register_range_table( + ctx, + "fact", + "testdata/join/parquet/fact", + "f_dkey", + vec![vec![ col("f_dkey").sort(true, false), col("timestamp").sort(true, false), - ]]); - ctx.register_parquet("fact", "testdata/join/parquet/fact", fact_options) - .await?; + ]], + ) + .await?; Ok(()) } @@ -96,7 +129,7 @@ mod tests { "#; // Execute the query using distributed datafusion, 2 workers, - // and hive-style partitioned data. + // and range-partitioned data. let (mut distributed_ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; set_configs(&mut distributed_ctx); @@ -105,7 +138,7 @@ mod tests { execute_query(&distributed_ctx, query).await?; // Ensure the distributed plan matches our target plan, registering - // hive-style partitioning and avoiding data-shuffling repartitions. + // range partitioning and avoiding data-shuffling repartitions. assert_snapshot!(&distributed_plan, @" ┌───── DistributedExec @@ -116,11 +149,11 @@ mod tests { │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@3, f_dkey@2)], projection=[f_dkey@6, timestamp@4, value@5, env@0, service@1, host@2] │ FilterExec: service@1 = log │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [], []]}, projection=[env, service, host, d_dkey], output_partitioning=Hash([d_dkey@3], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] - │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet], [], []]}, projection=[env, service, host, d_dkey], output_partitioning=Hash([d_dkey@3], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], []]}, projection=[env, service, host, d_dkey], output_partitioning=Range([d_dkey@3 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t1: DataSourceExec: file_groups={4 groups: [[], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, host, d_dkey], output_partitioning=Range([d_dkey@3 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 4), file_type=parquet - │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet], [], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 4), file_type=parquet + │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Range([f_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={4 groups: [[], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Range([f_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); @@ -176,7 +209,7 @@ mod tests { "#; // Execute the query using distributed datafusion, 2 workers, - // and hive-style partitioned data. + // and range-partitioned data. let (mut distributed_ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; set_configs(&mut distributed_ctx); @@ -185,7 +218,7 @@ mod tests { execute_query(&distributed_ctx, query).await?; // Ensure the distributed plan matches our target plan, registering - // hive-style partitioning and avoiding data-shuffling repartitions. + // range partitioning and avoiding data-shuffling repartitions. assert_snapshot!(&distributed_plan, @r#" ┌───── DistributedExec │ SortPreservingMergeExec: [f_dkey@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST] @@ -197,11 +230,11 @@ mod tests { │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@2)], projection=[f_dkey@4, env@0, timestamp@2, value@3] │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] - │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], []]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t1: DataSourceExec: file_groups={4 groups: [[], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 4), file_type=parquet - │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet], [], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 4), file_type=parquet + │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Range([f_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={4 groups: [[], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Range([f_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "#); @@ -253,7 +286,7 @@ mod tests { "#; // Execute the query using distributed datafusion, 2 workers, - // and hive-style partitioned data. + // and range-partitioned data. let (mut distributed_ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; set_configs(&mut distributed_ctx); @@ -262,7 +295,7 @@ mod tests { execute_query(&distributed_ctx, query).await?; // Ensure the distributed plan matches our target plan, registering - // hive-style partitioning and avoiding data-shuffling repartitions. + // range partitioning and avoiding data-shuffling repartitions. assert_snapshot!(&distributed_plan, @r#" ┌───── DistributedExec │ SortPreservingMergeExec: [env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST] @@ -279,11 +312,11 @@ mod tests { │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@2)], projection=[f_dkey@4, env@0, timestamp@2, value@3] │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] - │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], []]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t1: DataSourceExec: file_groups={4 groups: [[], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 4), file_type=parquet - │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet], [], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 4), file_type=parquet + │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], []]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Range([f_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={4 groups: [[], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Range([f_dkey@2 ASC], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "#); diff --git a/tests/range_partitioned.rs b/tests/range_partitioned.rs new file mode 100644 index 000000000..8db40f7e5 --- /dev/null +++ b/tests/range_partitioned.rs @@ -0,0 +1,175 @@ +#[cfg(all(feature = "integration", test))] +mod tests { + use arrow::{ + array::{Int32Array, RecordBatch}, + datatypes::{DataType, Field, Schema}, + util::pretty::pretty_format_batches, + }; + use datafusion::{ + common::{ScalarValue, SplitPoint}, + datasource::{ + file_format::parquet::ParquetFormat, + listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, + }, + error::Result, + logical_expr::{Partitioning as LogicalPartitioning, RangePartitioning, col}, + physical_plan::collect, + prelude::SessionContext, + }; + use datafusion_distributed::{ + DefaultSessionBuilder, assert_snapshot, display_plan_ascii, + test_utils::localhost::start_localhost_context, + }; + use parquet::arrow::ArrowWriter; + use std::fs::{self, File}; + use std::path::PathBuf; + use std::sync::Arc; + use uuid::Uuid; + + /// 8 range partitions with `target_partitions=2`. DataFusion must keep the + /// declared range partition count rather than collapsing to 2 file groups. + const RANGE_PARTITIONS: usize = 8; + const TARGET_PARTITIONS: usize = 2; + + #[tokio::test] + async fn test_range_scan_keeps_declared_partition_count() + -> Result<(), Box> { + let table_dir = write_range_table(RANGE_PARTITIONS, 1)?; + struct Cleanup(PathBuf); + impl Drop for Cleanup { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + let _cleanup = Cleanup(table_dir.clone()); + let (mut ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; + set_target_partitions(&mut ctx, TARGET_PARTITIONS); + register_range_table(&ctx, "t", &table_dir, RANGE_PARTITIONS)?; + + let (plan, results) = execute_query( + &ctx, + "SELECT range_key, SUM(value) AS total FROM t GROUP BY range_key ORDER BY range_key", + ) + .await?; + + assert_snapshot!(&plan, @r#" + ┌───── DistributedExec + │ SortPreservingMergeExec: [range_key@0 ASC NULLS LAST] + │ [Stage 1] => NetworkCoalesceExec: output_partitions=16, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=8 + │ ProjectionExec: expr=[range_key@0 as range_key, sum(t.value)@1 as total] + │ SortExec: expr=[range_key@0 ASC NULLS LAST], preserve_partitioning=[true] + │ AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(t.value)] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={8 groups: [[/dfd-range-UUID/part-0.parquet], [], [/dfd-range-UUID/part-2.parquet], [], [/dfd-range-UUID/part-4.parquet], ...]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30), (40), (50), (60), (70)], 8), file_type=parquet + │ t1: DataSourceExec: file_groups={8 groups: [[], [/dfd-range-UUID/part-1.parquet], [], [/dfd-range-UUID/part-3.parquet], [], ...]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30), (40), (50), (60), (70)], 8), file_type=parquet + └────────────────────────────────────────────────── + "#); + + let pretty_results = pretty_format_batches(&results)?; + assert_snapshot!(pretty_results, @" + +-----------+-------+ + | range_key | total | + +-----------+-------+ + | 0 | 0 | + | 10 | 10 | + | 20 | 20 | + | 30 | 30 | + | 40 | 40 | + | 50 | 50 | + | 60 | 60 | + | 70 | 70 | + +-----------+-------+ + "); + + Ok(()) + } + + fn set_target_partitions(ctx: &mut SessionContext, target_partitions: usize) { + ctx.state_ref() + .write() + .config_mut() + .options_mut() + .execution + .target_partitions = target_partitions; + } + + fn range_partitioning(range_partitions: usize) -> Result { + let split_points = (1..range_partitions) + .map(|i| SplitPoint::new(vec![ScalarValue::Int32(Some(i as i32 * 10))])) + .collect(); + Ok(LogicalPartitioning::Range(RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + split_points, + )?)) + } + + fn register_range_table( + ctx: &SessionContext, + name: &str, + table_dir: &PathBuf, + range_partitions: usize, + ) -> Result<()> { + let table_path = ListingTableUrl::parse(format!( + "{}/", + table_dir.to_str().expect("table path should be utf8") + ))?; + let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_output_partitioning(Some(range_partitioning(range_partitions)?)); + let schema = range_schema(); + let table = ListingTable::try_new( + ListingTableConfig::new(table_path) + .with_listing_options(options) + .with_schema(schema), + )?; + ctx.register_table(name, Arc::new(table))?; + Ok(()) + } + + async fn execute_query( + ctx: &SessionContext, + query: &str, + ) -> Result<(String, Vec)> { + let df = ctx.sql(query).await?; + let (state, logical_plan) = df.into_parts(); + let physical_plan = state.create_physical_plan(&logical_plan).await?; + let distributed_plan = display_plan_ascii(physical_plan.as_ref(), false); + let results = collect(physical_plan, state.task_ctx()).await?; + Ok((distributed_plan, results)) + } + + fn range_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("range_key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])) + } + + fn write_range_table(range_partitions: usize, files_per_partition: usize) -> Result { + let table_dir = + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(format!("dfd-range-{}", Uuid::new_v4())); + fs::create_dir_all(&table_dir)?; + let schema = range_schema(); + let mut file_idx = 0; + for part in 0..range_partitions { + let key = part as i32 * 10; + for _ in 0..files_per_partition { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![key])), + Arc::new(Int32Array::from(vec![key])), + ], + )?; + let path = table_dir.join(format!("part-{file_idx}.parquet")); + let file = File::create(&path)?; + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None)?; + writer.write(&batch)?; + writer.close()?; + file_idx += 1; + } + } + Ok(table_dir) + } +} From 8c6d2e39943bd8cf35441eca7f043d8d7676d3d8 Mon Sep 17 00:00:00 2001 From: Nishchaya Sharma <129059557+shinzoxD@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:30:25 +0530 Subject: [PATCH 2/2] fix: take &Path in range_partitioned test helper Clippy -D warnings failed on ptr_arg: writing &PathBuf instead of &Path. --- tests/range_partitioned.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/range_partitioned.rs b/tests/range_partitioned.rs index 8db40f7e5..08fdd1b59 100644 --- a/tests/range_partitioned.rs +++ b/tests/range_partitioned.rs @@ -22,7 +22,7 @@ mod tests { }; use parquet::arrow::ArrowWriter; use std::fs::{self, File}; - use std::path::PathBuf; + use std::path::{Path, PathBuf}; use std::sync::Arc; use uuid::Uuid; @@ -108,7 +108,7 @@ mod tests { fn register_range_table( ctx: &SessionContext, name: &str, - table_dir: &PathBuf, + table_dir: &Path, range_partitions: usize, ) -> Result<()> { let table_path = ListingTableUrl::parse(format!(