This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23231-9dc0f4b50c629726f4573dc2b024844429b227f0 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 50dcf0f47a648bae12a4e00033800edd0fbd0b7c Author: Saad Tajwar <[email protected]> AuthorDate: Wed Jul 8 12:48:16 2026 -0700 feat: physical execution for range partitioning (#23231) ## Which issue does this PR close? - Closes #23137 ## Rationale for this change Range repartitioning was already planned and serialized into physical plans, but `RepartitionExec` could not execute it. This PR completes the core execution path so rows in an input batch are routed to the correct output partition based on range split points and the ordering defined on the partitioning scheme. ## What changes are included in this PR? This PR adds a `Range` variant to `BatchPartitioner` that evaluates the ordering expressions on each input batch, compares each row's key against split points using `compare_rows` (respecting ASC/DESC and null ordering), and assigns row indices to output partitions via binary search. The partitioned row indices are then materialized into sub-batches using the same `partition_grouped_take` path as hash repartitioning. `pull_from_input` is wired to construct a range partitioner for `Partitioning::Range`, replacing the previous `not_impl_err!` at execution time. Optimizer-related paths remain intentionally unimplemented and are tracked in [#23230](https://github.com/apache/datafusion/issues/23230): projection pushdown through `RepartitionExec` (`try_swapping_with_projection`), sort pushdown (`try_pushdown_sort`), and changing partition counts via `repartitioned()`. ## Are these changes tested? Yes! ## Are there any user-facing changes? No public API changes --- datafusion/physical-plan/src/repartition/mod.rs | 590 +++++++++++++++++++----- 1 file changed, 483 insertions(+), 107 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 1617af3a68..cb55a19ee8 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -19,10 +19,11 @@ //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. +use std::cmp::Ordering; use std::fmt::{Debug, Formatter}; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use std::vec; @@ -46,20 +47,22 @@ use crate::{ check_if_same_properties, }; -use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions}; +use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions}; use arrow::compute::take_arrays; use arrow::datatypes::{SchemaRef, UInt32Type}; +use arrow_schema::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; -use datafusion_common::utils::transpose; +use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, assert_or_internal_err, internal_err, + ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, + assert_or_internal_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; -use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr, RangePartitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use crate::filter_pushdown::{ @@ -256,7 +259,7 @@ impl SharedCoalescer { /// sender, finalize the coalescer and return its residual batches; if /// other senders are still active, return `Ok(None)`. fn finalize(&self) -> Result<Vec<RecordBatch>> { - let was_last = self.active_senders.fetch_sub(1, Ordering::AcqRel) == 1; + let was_last = self.active_senders.fetch_sub(1, AtomicOrdering::AcqRel) == 1; if !was_last { return Ok(vec![]); } @@ -570,6 +573,18 @@ enum BatchPartitionerState { num_partitions: usize, next_idx: usize, }, + Range { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Sort options from the `LexOrdering` + sort_options: Vec<SortOptions>, + /// Boundaries between adjacent partitions. + split_points: Vec<SplitPoint>, + /// Row indices grouped by output partition + indices: Vec<Vec<u32>>, + /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points + partition_buffer: Vec<ScalarValue>, + }, } /// Fixed RandomState used for hash repartitioning to ensure consistent behavior across @@ -706,13 +721,40 @@ impl BatchPartitioner { timer, } } + + /// Create a new [`BatchPartitioner`] for range-based repartitioning. + /// + /// # Parameters + /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions + /// - `timer`: Metric used to record time spent during repartitioning. + pub fn new_range_partitioner( + range_partitioning: &RangePartitioning, + timer: metrics::Time, + ) -> Self { + let ordering = range_partitioning.ordering().clone(); + let split_points = range_partitioning.split_points().to_vec(); + let num_partitions = range_partitioning.partition_count(); + let sort_options: Vec<SortOptions> = ordering.iter().map(|e| e.options).collect(); + + Self { + state: BatchPartitionerState::Range { + partition_buffer: Vec::with_capacity(ordering.len()), + ordering, + sort_options, + split_points, + indices: vec![vec![]; num_partitions], + }, + timer, + } + } + /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. /// /// This is a convenience constructor that delegates to the specialized - /// hash or round-robin constructors depending on the partitioning variant. + /// hash, round-robin, or range constructors depending on the partitioning variant. /// /// # Parameters - /// - `partitioning`: Partitioning scheme to apply (hash or round-robin). + /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range). /// - `timer`: Metric used to record time spent during repartitioning. /// - `input_partition`: Index of the current input partition. /// - `num_input_partitions`: Total number of input partitions. @@ -738,12 +780,8 @@ impl BatchPartitioner { num_input_partitions, )) } - Partitioning::Range(_) => { - // Range repartition execution is tracked in - // https://github.com/apache/datafusion/issues/22397 - not_impl_err!( - "Range partitioning execution is not implemented by RepartitionExec" - ) + Partitioning::Range(range_repartitioning) => { + Ok(Self::new_range_partitioner(&range_repartitioning, timer)) } other => { not_impl_err!("Unsupported repartitioning scheme {other:?}") @@ -831,22 +869,95 @@ impl BatchPartitioner { Box::new(partitioned_batches.into_iter()) } + BatchPartitionerState::Range { + ordering, + sort_options, + split_points, + indices, + partition_buffer, + } => { + // Tracking time required for distributing indexes across output partitions + let timer = self.timer.timer(); + if split_points.is_empty() { + timer.done(); + Box::new(std::iter::once(Ok((0, batch)))) + } else { + let arrays = evaluate_expressions_to_arrays( + ordering.iter().map(|e| &e.expr), + &batch, + )?; + + indices.iter_mut().for_each(|v| v.clear()); + + Self::partition_range_indices( + &arrays, + split_points, + sort_options, + partition_buffer, + indices, + )?; + + // Finished building index-arrays for output partitions + timer.done(); + + let partitioned_batches = + Self::partition_grouped_take(&batch, indices, &self.timer)?; + + Box::new(partitioned_batches.into_iter()) + } + } }; Ok(it) } + /// Groups input row indices by range partition. This populates `indices[p]` with the + /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`. + fn partition_range_indices( + arrays: &[Arc<dyn Array>], + split_points: &[SplitPoint], + sort_options: &[SortOptions], + row_key_buffer: &mut Vec<ScalarValue>, + indices: &mut [Vec<u32>], + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row + extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?; + + let mut low = 0; + let mut high = split_points.len(); + while low < high { + let mid = low + (high - low) / 2; + let comparison = compare_rows( + row_key_buffer, + split_points[mid].values(), + sort_options, + )?; + match comparison { + Ordering::Less => high = mid, + Ordering::Equal | Ordering::Greater => low = mid + 1, + } + } + + indices[low].push(row_idx as u32) + } + + Ok(()) + } + // return the number of output partitions fn num_partitions(&self) -> usize { match &self.state { BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions, - BatchPartitionerState::Hash { indices, .. } => indices.len(), + BatchPartitionerState::Hash { indices, .. } + | BatchPartitionerState::Range { indices, .. } => indices.len(), } } - /// Build repartitioned hash output batches using one `take` per input batch. + /// Build repartitioned hash/range output batches using one `take` per input batch. /// - /// The hash router first fills one index vector per output partition. This method + /// The routers first fills one index vector per output partition. This method /// concatenates those index vectors, performs one grouped `take_arrays`, and /// then returns each output partition as a slice of the reordered batch. /// @@ -1447,10 +1558,8 @@ impl ExecutionPlan for RepartitionExec { } Partitioning::Range(_) => { // Range partitioning optimizer propagation is tracked in - // https://github.com/apache/datafusion/issues/22395 - return not_impl_err!( - "Projection pushdown through RepartitionExec with range partitioning is not implemented" - ); + // https://github.com/apache/datafusion/issues/23230 + return Ok(None); } others => others.clone(), }; @@ -1491,10 +1600,8 @@ impl ExecutionPlan for RepartitionExec { match self.partitioning() { Partitioning::Range(_) => { // Range partitioning optimizer propagation is tracked in - // https://github.com/apache/datafusion/issues/22395 - return not_impl_err!( - "Sort pushdown through RepartitionExec with range partitioning is not implemented" - ); + // https://github.com/apache/datafusion/issues/23230 + return Ok(SortOrderPushdownResult::Unsupported); } Partitioning::RoundRobinBatch(_) | Partitioning::Hash(_, _) @@ -1524,11 +1631,9 @@ impl ExecutionPlan for RepartitionExec { Hash(hash, _) => Hash(hash, target_partitions), UnknownPartitioning(_) => UnknownPartitioning(target_partitions), Range(_) => { - // Range repartition execution is tracked in - // https://github.com/apache/datafusion/issues/22397 - return not_impl_err!( - "Changing RepartitionExec partition counts with range partitioning is not implemented" - ); + // Range repartition optimizations are tracked in + // https://github.com/apache/datafusion/issues/23230 + return Ok(None); } }; Ok(Some(Arc::new(Self { @@ -1642,33 +1747,12 @@ impl RepartitionExec { input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = match &partitioning { - Partitioning::Hash(exprs, num_partitions) => { - BatchPartitioner::new_hash_partitioner( - exprs.clone(), - *num_partitions, - metrics.repartition_time.clone(), - )? - } - Partitioning::RoundRobinBatch(num_partitions) => { - BatchPartitioner::new_round_robin_partitioner( - *num_partitions, - metrics.repartition_time.clone(), - input_partition, - num_input_partitions, - ) - } - Partitioning::Range(_) => { - // Range repartition execution is tracked in - // https://github.com/apache/datafusion/issues/22397 - return not_impl_err!( - "Range partitioning execution is not implemented by RepartitionExec" - ); - } - other => { - return not_impl_err!("Unsupported repartitioning scheme {other:?}"); - } - }; + let mut partitioner = BatchPartitioner::try_new( + partitioning, + metrics.repartition_time.clone(), + input_partition, + num_input_partitions, + )?; // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); @@ -2024,7 +2108,7 @@ mod tests { use arrow::array::{ArrayRef, StringArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::ScalarValue; - use datafusion_common::cast::as_string_array; + use datafusion_common::cast::{as_string_array, as_uint32_array}; use datafusion_common::exec_err; use datafusion_common::test_util::batches_to_sort_string; use datafusion_common_runtime::JoinSet; @@ -2128,7 +2212,7 @@ mod tests { #[tokio::test] async fn one_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition]; @@ -2151,7 +2235,7 @@ mod tests { #[tokio::test] async fn many_to_one_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2168,7 +2252,7 @@ mod tests { #[tokio::test] async fn many_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2189,7 +2273,7 @@ mod tests { #[tokio::test] async fn many_to_many_hash_partition() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2211,9 +2295,267 @@ mod tests { Ok(()) } + #[tokio::test] + async fn many_to_many_range_partition() -> Result<()> { + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + // create_batch values are [1, 2, 3, 4, 5, 6, 7, 8]; split at 3 and 6 yields + // 2, 3, and 3 rows per batch respectively + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![3, 6])?; + + let output_partitions = repartition(&schema, partitions, partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!(300, partition_row_count(&output_partitions[0])); + assert_eq!(450, partition_row_count(&output_partitions[1])); + assert_eq!(450, partition_row_count(&output_partitions[2])); + assert_eq!( + collect_partition_u32_values(&output_partitions[0]) + .into_iter() + .flatten() + .collect::<HashSet<_>>(), + HashSet::from([1, 2]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[1]) + .into_iter() + .flatten() + .collect::<HashSet<_>>(), + HashSet::from([3, 4, 5]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[2]) + .into_iter() + .flatten() + .collect::<HashSet<_>>(), + HashSet::from([6, 7, 8]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_compound_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![5, 10, 10, 10, 10, 15])), + Arc::new(UInt32Array::from(vec![1, 1, 3, 5, 7, 0])), + ], + )?; + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(1)), + ]), + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(5)), + ]), + ], + )?); + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![(5, 1)], + collect_partition_u32_pairs(&output_partitions[0]) + ); + assert_eq!( + vec![(10, 1), (10, 3)], + collect_partition_u32_pairs(&output_partitions[1]) + ); + assert_eq!( + vec![(10, 5), (10, 7), (15, 0)], + collect_partition_u32_pairs(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_last() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, false), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![None, Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_first() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, true), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![None, Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_asc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![10, 20])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_desc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 20, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(true, false), vec![20, 10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(15), Some(20)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(5), Some(10)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_string_rows() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["bar", "baz", "foo", "qux"])) as ArrayRef, + )])?; + + let schema = batch.schema(); + let expr = col("my_awesome_field", &schema)?; + let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema)); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new_default(expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "foo".to_string(), + ))])], + )?); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?; + + let mut partition_0 = Vec::new(); + let mut stream = exec.execute(0, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + partition_0.push(result?); + } + + let mut partition_1 = Vec::new(); + let mut stream = exec.execute(1, task_ctx)?; + while let Some(result) = stream.next().await { + partition_1.push(result?); + } + + assert_eq!( + vec!["bar", "baz"], + collect_partition_string_values(&partition_0) + ); + assert_eq!( + vec!["foo", "qux"], + collect_partition_string_values(&partition_1) + ); + + Ok(()) + } + #[tokio::test] async fn test_repartition_with_coalescing() -> Result<()> { - let schema = test_schema(); + let schema = test_schema(false); // create 50 batches, each having 8 rows let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone()]; @@ -2237,8 +2579,76 @@ mod tests { Ok(()) } - fn test_schema() -> Arc<Schema> { - Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) + fn test_schema(nullable: bool) -> Arc<Schema> { + Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::UInt32, + nullable, + )])) + } + + fn u32_range_partitioning( + schema: &SchemaRef, + sort_options: SortOptions, + split_values: Vec<u32>, + ) -> Result<Partitioning> { + let expr = col("c0", schema)?; + Ok(Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(expr, sort_options)].into(), + split_values + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect(), + )?)) + } + + fn partition_row_count(batches: &[RecordBatch]) -> usize { + batches.iter().map(|batch| batch.num_rows()).sum() + } + + fn collect_partition_u32_values(batches: &[RecordBatch]) -> Vec<Option<u32>> { + batches + .iter() + .flat_map(|batch| { + let array = + as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + (0..array.len()) + .map(|idx| { + if array.is_null(idx) { + None + } else { + Some(array.value(idx)) + } + }) + .collect::<Vec<_>>() + }) + .collect() + } + + fn collect_partition_u32_pairs(batches: &[RecordBatch]) -> Vec<(u32, u32)> { + batches + .iter() + .flat_map(|batch| { + let a = as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + let b = as_uint32_array(batch.column(1)).expect("expected UInt32 column"); + (0..a.len()) + .map(|idx| (a.value(idx), b.value(idx))) + .collect::<Vec<_>>() + }) + .collect() + } + + fn collect_partition_string_values(batches: &[RecordBatch]) -> Vec<&str> { + batches + .iter() + .flat_map(|batch| { + let array = + as_string_array(batch.column(0)).expect("expected Utf8 column"); + (0..array.len()) + .map(|idx| array.value(idx)) + .collect::<Vec<_>>() + }) + .collect() } async fn repartition( @@ -2271,7 +2681,7 @@ mod tests { let handle: SpawnedTask<Result<Vec<Vec<RecordBatch>>>> = SpawnedTask::spawn(async move { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2323,40 +2733,6 @@ mod tests { ); } - #[tokio::test] - async fn unsupported_range_partitioning() -> Result<()> { - let task_ctx = Arc::new(TaskContext::default()); - let batch = RecordBatch::try_from_iter(vec![( - "my_awesome_field", - Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef, - )])?; - - let schema = batch.schema(); - let expr = col("my_awesome_field", &schema)?; - let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema)); - let partitioning = Partitioning::Range(RangePartitioning::new( - [PhysicalSortExpr::new_default(expr)].into(), - vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( - "foo".to_string(), - ))])], - )); - let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?; - let output_stream = exec.execute(0, task_ctx)?; - - let result_string = crate::common::collect(output_stream) - .await - .unwrap_err() - .to_string(); - assert!( - result_string.contains( - "Range partitioning execution is not implemented by RepartitionExec" - ), - "actual: {result_string}" - ); - - Ok(()) - } - #[tokio::test] async fn error_for_input_exec() { // This generates an error on a call to execute. The error @@ -2665,7 +3041,7 @@ mod tests { #[tokio::test] async fn repartition_with_spilling() -> Result<()> { // Test that repartition successfully spills to disk when memory is constrained - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2727,7 +3103,7 @@ mod tests { #[tokio::test] async fn repartition_with_partial_spilling() -> Result<()> { // Test that repartition can handle partial spilling (some batches in memory, some spilled) - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2797,7 +3173,7 @@ mod tests { #[tokio::test] async fn repartition_without_spilling() -> Result<()> { // Test that repartition does not spill when there's ample memory - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2859,7 +3235,7 @@ mod tests { use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; // Test that repartition fails with OOM when disk manager is disabled - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2902,7 +3278,7 @@ mod tests { /// Create batch fn create_batch() -> RecordBatch { - let schema = test_schema(); + let schema = test_schema(false); RecordBatch::try_new( schema, vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))], @@ -2912,7 +3288,7 @@ mod tests { /// Create batches with sequential values for ordering tests fn create_ordered_batches(num_batches: usize) -> Vec<RecordBatch> { - let schema = test_schema(); + let schema = test_schema(false); (0..num_batches) .map(|i| { let start = (i * 8) as u32; @@ -2933,7 +3309,7 @@ mod tests { // This tests the state machine fix where we must block on spill_stream // when a Spilled marker is received, rather than continuing to poll the channel - let schema = test_schema(); + let schema = test_schema(false); // Create batches with sequential values: batch 0 has [0,1,2,3,4,5,6,7], // batch 1 has [8,9,10,11,12,13,14,15], etc. let partition = create_ordered_batches(20); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
