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..08fdd1b59 --- /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::{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: &Path, + 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) + } +}