You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/06 07:52:33 UTC
(datafusion) branch gh-readonly-queue/main/pr-24598-3266eaa91714cb1ffd0ba8f17f51880e05289084 created (now 8642154323)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24598-3266eaa91714cb1ffd0ba8f17f51880e05289084
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 8642154323 Range partition on primitives. (#24598)
This branch includes the following new commits:
new 8642154323 Range partition on primitives. (#24598)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Range partition on primitives. (#24598)
Posted by gi...@apache.org.
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-24598-3266eaa91714cb1ffd0ba8f17f51880e05289084
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 8642154323c3800948f495af9ba2ee1539399105
Author: Stu Hood <st...@gmail.com>
AuthorDate: Sun Sep 6 07:52:09 2026 +0000
Range partition on primitives. (#24598)
## Which issue does this PR close?
- Closes #24597.
- Related to https://github.com/apache/datafusion/issues/22395.
## Rationale for this change
In query plans that use range re-partitioning (such as matching the
range partitioning of underlying base data, or evaluating dynamic range
filters in hash joins), range re-partitioning can represent a
significant fraction of total query execution time.
As detailed in the issue, the existing implementation in
[`RepartitionExec`](https://github.com/apache/datafusion/blob/6fe001d4baf5c6089d24b5634a23b1345797fb14/datafusion/physical-plan/src/repartition/mod.rs#L100-L150)
/
[`BatchPartitioner`](https://github.com/apache/datafusion/blob/6fe001d4baf5c6089d24b5634a23b1345797fb14/datafusion/physical-plan/src/repartition/mod.rs#L623-L640)
and
[`RangeExpr`](https://github.com/apache/datafusion/blob/6fe001d4baf5c6089d24b5634a23b1345797fb14/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs#L180-L240)
assigns rows to partitions on a row-by-row basis using
[`extract_row_at_idx_to_buf`](https://github.com/apache/datafusion/blob/6fe001d4baf5c6089d24b5634a23b1345797fb14/datafusion-common/src/utils/mod.rs)
and dynamic
[`ScalarValue`](https://github.com/apache/datafusion/blob/6fe001d4baf5c6089d24b5634a23b1345797fb14/datafusion-common/src/scalar/mod.rs)
comparisons. This allocates per row and has dynamic dispatch overhead,
making range re-partitioning 15x–30x slower than hash re-partitioning on
standard integer keys.
This PR adds a zero-allocation, vectorizable implementation, achieving
10x–16x speedup on primitive numeric keys and 2x–3x speedup on string
and composite keys.
## What changes are included in this PR?
- Adds an internal `RangeRouter` with specialized routing paths:
- A zero-allocation path for single primitive numeric, temporal, and
float columns that binary searches directly over Arrow buffer slices.
- A (more) vectorizable path using Arrow's `RowConverter` storing flat
`Rows` buffers for strings, decimals, dictionary arrays, and composite
keys.
- Updates `RepartitionExec` to construct and share an `Arc<RangeRouter>`
across input partitions, avoiding redundant split-point encoding.
- Updates `BatchPartitioner` and `RangeExpr` to delegate partition
routing to `RangeRouter`.
- Adds `BatchPartitioner::try_new_range_partitioner` and deprecates
`BatchPartitioner::new_range_partitioner`.
- Adds a Criterion suite.
### Benchmark Results (vs `main` baseline)
| Benchmark | Partitions | Baseline (`main`) | This PR | Speedup |
| :--- | :--- | :--- | :--- | :--- |
| **`range_repartition_i64_uniform`** | 8 | 275.3 µs (29.8 Melem/s) |
23.4 µs (349.5 Melem/s) | **11.7x** |
| | 64 | 459.8 µs (17.8 Melem/s) | 40.9 µs (200.4 Melem/s) | **11.2x** |
| | 512 | 1,194.6 µs (6.86 Melem/s) | 89.9 µs (91.2 Melem/s) | **12.2x**
|
| **`range_repartition_i64_sequential`** | 8 | 291.8 µs (28.1 Melem/s) |
23.2 µs (353.6 Melem/s) | **12.6x** |
| | 64 | 478.6 µs (17.1 Melem/s) | 38.5 µs (212.9 Melem/s) | **12.4x** |
| | 512 | 989.9 µs (8.28 Melem/s) | 83.9 µs (97.6 Melem/s) | **11.8x** |
| **`range_repartition_utf8_uniform`** | 8 | 543.5 µs (15.1 Melem/s) |
184.2 µs (44.5 Melem/s) | **3.0x** |
| | 64 | 964.7 µs (8.49 Melem/s) | 330.5 µs (24.8 Melem/s) | **2.9x** |
| | 512 | 1,333.2 µs (6.14 Melem/s) | 551.7 µs (14.8 Melem/s) | **2.4x**
|
| **`range_repartition_composite_i64`** | 8 | 344.2 µs (23.8 Melem/s) |
114.3 µs (71.7 Melem/s) | **3.0x** |
| | 64 | 516.9 µs (15.8 Melem/s) | 249.7 µs (32.8 Melem/s) | **2.1x** |
| | 512 | 985.3 µs (8.31 Melem/s) | 473.6 µs (17.3 Melem/s) | **2.1x** |
## Are these changes tested?
Added unit and integration tests.
## Are there any user-facing changes?
- Added `BatchPartitioner::try_new_range_partitioner(&RangePartitioning,
&Schema, Time) -> Result<Self>`, which takes the input `&Schema` to
infer column data types and coerce split points.
- Deprecated `BatchPartitioner::new_range_partitioner` in favor of
`try_new_range_partitioner`.
- **`RangeExpr`**:
- Added `RangeExpr::try_new_with_schema(Vec<PhysicalExprRef>,
&RangePartitioning, &Schema) -> Result<Self>`, ensuring routing matches
key column data types (precision, scale, timezone).
- Deprecated `RangeExpr::try_new(Vec<PhysicalExprRef>,
&RangePartitioning)` in favor of `try_new_with_schema`.
---
datafusion/physical-plan/Cargo.toml | 4 +
.../physical-plan/benches/range_repartition.rs | 425 ++++++++++
.../src/joins/hash_join/shared_bounds.rs | 3 +-
datafusion/physical-plan/src/repartition/mod.rs | 632 +++++++++++---
datafusion/physical-plan/src/repartition/range.rs | 908 +++++++++++++++++++++
datafusion/proto/tests/cases/plans/exprs.rs | 3 +-
6 files changed, 1857 insertions(+), 118 deletions(-)
diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml
index 1adac74ecb..0aa22653b4 100644
--- a/datafusion/physical-plan/Cargo.toml
+++ b/datafusion/physical-plan/Cargo.toml
@@ -161,3 +161,7 @@ name = "bounded_window"
[[bench]]
harness = false
name = "window_filter"
+
+[[bench]]
+harness = false
+name = "range_repartition"
diff --git a/datafusion/physical-plan/benches/range_repartition.rs b/datafusion/physical-plan/benches/range_repartition.rs
new file mode 100644
index 0000000000..b90e313a5e
--- /dev/null
+++ b/datafusion/physical-plan/benches/range_repartition.rs
@@ -0,0 +1,425 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use arrow::compute::SortOptions;
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use datafusion_common::ScalarValue;
+use datafusion_physical_expr::expressions::col;
+use datafusion_physical_expr::{
+ LexOrdering, PhysicalExpr, PhysicalSortExpr, RangePartitioning, SplitPoint,
+};
+use datafusion_physical_plan::metrics::Time;
+use datafusion_physical_plan::repartition::{BatchPartitioner, RangeExpr};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+
+const BATCH_SIZE: usize = 8192;
+const PARTITION_COUNTS: [usize; 7] = [8, 16, 32, 64, 128, 256, 512];
+const SEED: u64 = 42;
+
+fn create_i64_uniform_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch {
+ let mut rng = StdRng::seed_from_u64(SEED);
+ let key_values: Vec<i64> = (0..BATCH_SIZE)
+ .map(|_| rng.random_range(0..max_val))
+ .collect();
+ let payload_values: Vec<i64> = (0..BATCH_SIZE).map(|i| i as i64).collect();
+
+ RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![
+ Arc::new(Int64Array::from(key_values)) as ArrayRef,
+ Arc::new(Int64Array::from(payload_values)) as ArrayRef,
+ ],
+ )
+ .unwrap()
+}
+
+fn create_i64_sequential_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch {
+ let key_values: Vec<i64> = (0..BATCH_SIZE)
+ .map(|i| ((i as i64) * max_val) / (BATCH_SIZE as i64))
+ .collect();
+ let payload_values: Vec<i64> = (0..BATCH_SIZE).map(|i| i as i64).collect();
+
+ RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![
+ Arc::new(Int64Array::from(key_values)) as ArrayRef,
+ Arc::new(Int64Array::from(payload_values)) as ArrayRef,
+ ],
+ )
+ .unwrap()
+}
+
+fn create_utf8_uniform_batch(schema: &SchemaRef, max_val: usize) -> RecordBatch {
+ let mut rng = StdRng::seed_from_u64(SEED);
+ let key_strings: Vec<String> = (0..BATCH_SIZE)
+ .map(|_| format!("key_{:010}", rng.random_range(0..max_val)))
+ .collect();
+ let payload_values: Vec<i64> = (0..BATCH_SIZE).map(|i| i as i64).collect();
+
+ RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![
+ Arc::new(StringArray::from_iter_values(
+ key_strings.iter().map(String::as_str),
+ )) as ArrayRef,
+ Arc::new(Int64Array::from(payload_values)) as ArrayRef,
+ ],
+ )
+ .unwrap()
+}
+
+fn create_composite_i64_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch {
+ let mut rng = StdRng::seed_from_u64(SEED);
+ let key1_values: Vec<i64> = (0..BATCH_SIZE)
+ .map(|_| rng.random_range(0..max_val))
+ .collect();
+ let key2_values: Vec<i64> = (0..BATCH_SIZE)
+ .map(|_| rng.random_range(0..max_val))
+ .collect();
+ let payload_values: Vec<i64> = (0..BATCH_SIZE).map(|i| i as i64).collect();
+
+ RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![
+ Arc::new(Int64Array::from(key1_values)) as ArrayRef,
+ Arc::new(Int64Array::from(key2_values)) as ArrayRef,
+ Arc::new(Int64Array::from(payload_values)) as ArrayRef,
+ ],
+ )
+ .unwrap()
+}
+
+fn bench_range_repartition_i64_uniform(c: &mut Criterion) {
+ let mut group = c.benchmark_group("range_repartition_i64_uniform");
+ group.throughput(Throughput::Elements(BATCH_SIZE as u64));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int64, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ let max_val = 1_000_000i64;
+ let batch = create_i64_uniform_batch(&schema, max_val);
+
+ for &num_partitions in &PARTITION_COUNTS {
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("key", &schema).unwrap(),
+ SortOptions::default(),
+ )])
+ .unwrap();
+
+ let split_points: Vec<SplitPoint> = (1..num_partitions)
+ .map(|i| {
+ let val = (i as i64 * max_val) / (num_partitions as i64);
+ SplitPoint::new(vec![ScalarValue::Int64(Some(val))])
+ })
+ .collect();
+
+ let range_part = RangePartitioning::try_new(ordering, split_points).unwrap();
+
+ group.bench_with_input(
+ BenchmarkId::new("partitions", num_partitions),
+ &num_partitions,
+ |b, _| {
+ let mut partitioner = BatchPartitioner::try_new_range_partitioner(
+ &range_part,
+ &schema,
+ Time::default(),
+ )
+ .unwrap();
+ b.iter(|| {
+ partitioner
+ .partition(batch.clone(), |p, b| {
+ black_box((p, b));
+ Ok(())
+ })
+ .unwrap();
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+fn bench_range_repartition_i64_sequential(c: &mut Criterion) {
+ let mut group = c.benchmark_group("range_repartition_i64_sequential");
+ group.throughput(Throughput::Elements(BATCH_SIZE as u64));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int64, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ let max_val = 1_000_000i64;
+ let batch = create_i64_sequential_batch(&schema, max_val);
+
+ for &num_partitions in &PARTITION_COUNTS {
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("key", &schema).unwrap(),
+ SortOptions::default(),
+ )])
+ .unwrap();
+
+ let split_points: Vec<SplitPoint> = (1..num_partitions)
+ .map(|i| {
+ let val = (i as i64 * max_val) / (num_partitions as i64);
+ SplitPoint::new(vec![ScalarValue::Int64(Some(val))])
+ })
+ .collect();
+
+ let range_part = RangePartitioning::try_new(ordering, split_points).unwrap();
+
+ group.bench_with_input(
+ BenchmarkId::new("partitions", num_partitions),
+ &num_partitions,
+ |b, _| {
+ let mut partitioner = BatchPartitioner::try_new_range_partitioner(
+ &range_part,
+ &schema,
+ Time::default(),
+ )
+ .unwrap();
+ b.iter(|| {
+ partitioner
+ .partition(batch.clone(), |p, b| {
+ black_box((p, b));
+ Ok(())
+ })
+ .unwrap();
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+fn bench_range_repartition_utf8_uniform(c: &mut Criterion) {
+ let mut group = c.benchmark_group("range_repartition_utf8_uniform");
+ group.throughput(Throughput::Elements(BATCH_SIZE as u64));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ let max_val = 1_000_000usize;
+ let batch = create_utf8_uniform_batch(&schema, max_val);
+
+ for &num_partitions in &PARTITION_COUNTS {
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("key", &schema).unwrap(),
+ SortOptions::default(),
+ )])
+ .unwrap();
+
+ let split_points: Vec<SplitPoint> = (1..num_partitions)
+ .map(|i| {
+ let val = (i * max_val) / num_partitions;
+ SplitPoint::new(vec![ScalarValue::Utf8(Some(format!("key_{val:010}")))])
+ })
+ .collect();
+
+ let range_part = RangePartitioning::try_new(ordering, split_points).unwrap();
+
+ group.bench_with_input(
+ BenchmarkId::new("partitions", num_partitions),
+ &num_partitions,
+ |b, _| {
+ let mut partitioner = BatchPartitioner::try_new_range_partitioner(
+ &range_part,
+ &schema,
+ Time::default(),
+ )
+ .unwrap();
+ b.iter(|| {
+ partitioner
+ .partition(batch.clone(), |p, b| {
+ black_box((p, b));
+ Ok(())
+ })
+ .unwrap();
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+fn bench_range_repartition_composite_i64(c: &mut Criterion) {
+ let mut group = c.benchmark_group("range_repartition_composite_i64");
+ group.throughput(Throughput::Elements(BATCH_SIZE as u64));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key1", DataType::Int64, false),
+ Field::new("key2", DataType::Int64, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ let max_val = 1_000_000i64;
+ let batch = create_composite_i64_batch(&schema, max_val);
+
+ for &num_partitions in &PARTITION_COUNTS {
+ let ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new(col("key1", &schema).unwrap(), SortOptions::default()),
+ PhysicalSortExpr::new(col("key2", &schema).unwrap(), SortOptions::default()),
+ ])
+ .unwrap();
+
+ let split_points: Vec<SplitPoint> = (1..num_partitions)
+ .map(|i| {
+ let val1 = (i as i64 * max_val) / (num_partitions as i64);
+ let val2 = 0i64;
+ SplitPoint::new(vec![
+ ScalarValue::Int64(Some(val1)),
+ ScalarValue::Int64(Some(val2)),
+ ])
+ })
+ .collect();
+
+ let range_part = RangePartitioning::try_new(ordering, split_points).unwrap();
+
+ group.bench_with_input(
+ BenchmarkId::new("partitions", num_partitions),
+ &num_partitions,
+ |b, _| {
+ let mut partitioner = BatchPartitioner::try_new_range_partitioner(
+ &range_part,
+ &schema,
+ Time::default(),
+ )
+ .unwrap();
+ b.iter(|| {
+ partitioner
+ .partition(batch.clone(), |p, b| {
+ black_box((p, b));
+ Ok(())
+ })
+ .unwrap();
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+fn bench_range_expr_routing_i64(c: &mut Criterion) {
+ let mut group = c.benchmark_group("range_expr_routing_i64");
+ group.throughput(Throughput::Elements(BATCH_SIZE as u64));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int64, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ let max_val = 1_000_000i64;
+ let batch = create_i64_uniform_batch(&schema, max_val);
+
+ for &num_partitions in &PARTITION_COUNTS {
+ let col_expr = col("key", &schema).unwrap();
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::clone(&col_expr),
+ SortOptions::default(),
+ )])
+ .unwrap();
+
+ let split_points: Vec<SplitPoint> = (1..num_partitions)
+ .map(|i| {
+ let val = (i as i64 * max_val) / (num_partitions as i64);
+ SplitPoint::new(vec![ScalarValue::Int64(Some(val))])
+ })
+ .collect();
+
+ let range_part = RangePartitioning::try_new(ordering, split_points).unwrap();
+ let range_expr =
+ RangeExpr::try_new_with_schema(vec![col_expr], &range_part, &batch.schema())
+ .unwrap();
+
+ group.bench_with_input(
+ BenchmarkId::new("partitions", num_partitions),
+ &num_partitions,
+ |b, _| {
+ b.iter(|| {
+ let res = range_expr.evaluate(&batch).unwrap();
+ black_box(res);
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+fn bench_range_partitioner_construction(c: &mut Criterion) {
+ let mut group = c.benchmark_group("range_partitioner_construction");
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int64, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ let max_val = 1_000_000i64;
+
+ for &num_partitions in &PARTITION_COUNTS {
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("key", &schema).unwrap(),
+ SortOptions::default(),
+ )])
+ .unwrap();
+
+ let split_points: Vec<SplitPoint> = (1..num_partitions)
+ .map(|i| {
+ let val = (i as i64 * max_val) / (num_partitions as i64);
+ SplitPoint::new(vec![ScalarValue::Int64(Some(val))])
+ })
+ .collect();
+
+ let range_part = RangePartitioning::try_new(ordering, split_points).unwrap();
+
+ group.bench_with_input(
+ BenchmarkId::new("partitions", num_partitions),
+ &num_partitions,
+ |b, _| {
+ b.iter(|| {
+ let partitioner = BatchPartitioner::try_new_range_partitioner(
+ black_box(&range_part),
+ black_box(&schema),
+ Time::default(),
+ )
+ .unwrap();
+ black_box(partitioner);
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+criterion_group!(
+ benches,
+ bench_range_repartition_i64_uniform,
+ bench_range_repartition_i64_sequential,
+ bench_range_repartition_utf8_uniform,
+ bench_range_repartition_composite_i64,
+ bench_range_expr_routing_i64,
+ bench_range_partitioner_construction
+);
+criterion_main!(benches);
diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs
index a4b2a86983..62087c14c5 100644
--- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs
@@ -761,9 +761,10 @@ impl SharedBuildAccumulator {
partition_filters.len(),
range_partitioning.partition_count()
);
- Arc::new(RangeExpr::try_new(
+ Arc::new(RangeExpr::try_new_with_schema(
self.on_right.clone(),
range_partitioning,
+ &self.probe_schema,
)?) as Arc<dyn PhysicalExpr>
} else {
// Routes probe rows with the same `hash(keys) % partition_count` expression used
diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
index 7822c6facd..0bbc68641c 100644
--- a/datafusion/physical-plan/src/repartition/mod.rs
+++ b/datafusion/physical-plan/src/repartition/mod.rs
@@ -19,7 +19,6 @@
//! 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, Display, Formatter};
use std::pin::Pin;
use std::sync::Arc;
@@ -47,18 +46,18 @@ use crate::{
PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count,
};
-use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array};
+use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array};
use arrow::compute::take_arrays;
use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type};
use arrow_schema::SortOptions;
+use datafusion_common::ScalarValue;
use datafusion_common::config::ConfigOptions;
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
-use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose};
+use datafusion_common::utils::transpose;
use datafusion_common::{
- ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint,
- assert_or_internal_err, internal_datafusion_err, internal_err,
- validate_range_split_points,
+ ColumnStatistics, DataFusionError, HashMap, SplitPoint, assert_or_internal_err,
+ internal_datafusion_err, internal_err,
};
use datafusion_common::{Result, not_impl_err};
use datafusion_common_runtime::SpawnedTask;
@@ -89,6 +88,11 @@ use log::trace;
use parking_lot::Mutex;
mod distributor_channels;
+mod range;
+
+use range::RangeRouter;
+use std::hash::{Hash, Hasher};
+
use crate::repartition::distributor_channels::SendError;
use distributor_channels::{
DistributionReceiver, DistributionSender, channels, partition_aware_channels,
@@ -558,6 +562,24 @@ impl RepartitionExecState {
);
}
+ let range_router = if let Partitioning::Range(range_partitioning) = &partitioning
+ {
+ let ordering = range_partitioning.ordering();
+ let sort_options: Vec<SortOptions> =
+ ordering.iter().map(|e| e.options).collect();
+ let data_types = ordering
+ .iter()
+ .map(|e| e.expr.data_type(input.schema().as_ref()))
+ .collect::<Result<Vec<_>>>()?;
+ Some(Arc::new(RangeRouter::try_new_with_data_types(
+ &sort_options,
+ range_partitioning.split_points(),
+ &data_types,
+ )?))
+ } else {
+ None
+ };
+
// launch one async task per *input* partition
let mut spawned_tasks = Vec::with_capacity(num_input_partitions);
for (i, (stream, metrics)) in
@@ -591,6 +613,7 @@ impl RepartitionExecState {
stream,
txs,
partitioning.clone(),
+ range_router.clone(),
metrics,
// preserve_order depends on partition index to start from 0
if preserve_order { 0 } else { i },
@@ -634,14 +657,10 @@ enum BatchPartitionerState {
Range {
/// Ordered partitioning key.
ordering: LexOrdering,
- /// Sort options from the `LexOrdering`
- sort_options: Vec<SortOptions>,
- /// Boundaries between adjacent partitions.
- split_points: Vec<SplitPoint>,
+ /// Router for partition assignment.
+ router: Arc<RangeRouter>,
/// 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>,
},
}
@@ -653,48 +672,106 @@ pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_
///
/// This uses the same routing function as [`BatchPartitioner`], so dynamic
/// filtering and repartitioning agree for every [`ScalarValue`] comparison.
-#[derive(Debug, Hash, PartialEq, Eq)]
+#[derive(Debug, Clone)]
pub struct RangeExpr {
on_columns: Vec<PhysicalExprRef>,
- split_points: Vec<SplitPoint>,
- sort_options: Vec<SortOptions>,
+ router: RangeRouter,
+}
+
+impl PartialEq for RangeExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.on_columns == other.on_columns
+ && self.router.split_points() == other.router.split_points()
+ && self.router.sort_options() == other.router.sort_options()
+ }
+}
+
+impl Eq for RangeExpr {}
+
+impl Hash for RangeExpr {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.on_columns.hash(state);
+ self.router.split_points().hash(state);
+ self.router.sort_options().hash(state);
+ }
}
impl RangeExpr {
/// Creates a Range expression for `on_columns` using the supplied routing
- /// metadata.
- pub fn try_new(
+ /// metadata and schema.
+ pub fn try_new_with_schema(
on_columns: Vec<PhysicalExprRef>,
range_partitioning: &RangePartitioning,
+ schema: &Schema,
) -> Result<Self> {
- let sort_options = range_partitioning
+ let sort_options: Vec<SortOptions> = range_partitioning
.ordering()
.iter()
.map(|expr| expr.options)
.collect();
Self::try_new_parts(
on_columns,
- range_partitioning.split_points().to_vec(),
- sort_options,
+ range_partitioning.split_points(),
+ &sort_options,
+ schema,
)
}
+ /// Creates a Range expression for `on_columns` using the supplied routing
+ /// metadata.
+ ///
+ /// # Workaround for backwards compatibility
+ /// This method is preserved for backwards compatibility with DataFusion 55.0 and earlier,
+ /// where `RangeExpr` was constructed without a schema. Key data types are inferred from the
+ /// split points, which will fail to route batches if split points do not carry the exact
+ /// precision, scale, or timezone as the input schema columns.
+ ///
+ /// Prefer [`Self::try_new_with_schema`].
+ #[deprecated(since = "56.0.0", note = "Use RangeExpr::try_new_with_schema instead")]
+ pub fn try_new(
+ on_columns: Vec<PhysicalExprRef>,
+ range_partitioning: &RangePartitioning,
+ ) -> Result<Self> {
+ let sort_options: Vec<SortOptions> = range_partitioning
+ .ordering()
+ .iter()
+ .map(|expr| expr.options)
+ .collect();
+ assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key");
+ assert_or_internal_err!(
+ on_columns.len() == sort_options.len(),
+ "RangeExpr key count must match sort options"
+ );
+ let router =
+ RangeRouter::try_new(&sort_options, range_partitioning.split_points())?;
+ Ok(Self { on_columns, router })
+ }
+
fn try_new_parts(
on_columns: Vec<PhysicalExprRef>,
- split_points: Vec<SplitPoint>,
- sort_options: Vec<SortOptions>,
+ split_points: &[SplitPoint],
+ sort_options: &[SortOptions],
+ schema: &Schema,
) -> Result<Self> {
assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key");
assert_or_internal_err!(
on_columns.len() == sort_options.len(),
"RangeExpr key count must match sort options"
);
- validate_range_split_points(&split_points, &sort_options)?;
- Ok(Self {
- on_columns,
- split_points,
+ // Route on the key's actual types, not the split points' types: split
+ // points are not required to carry the key's exact type (`Decimal128`
+ // precision, timestamp timezone), and `RangeRouter` rejects any array
+ // whose type differs from the one it was built for.
+ let data_types = on_columns
+ .iter()
+ .map(|expr| expr.data_type(schema))
+ .collect::<Result<Vec<_>>>()?;
+ let router = RangeRouter::try_new_with_data_types(
sort_options,
- })
+ split_points,
+ &data_types,
+ )?;
+ Ok(Self { on_columns, router })
}
/// Get the columns used to compute Range partition IDs.
@@ -704,12 +781,12 @@ impl RangeExpr {
/// Returns the Range split points used for routing.
pub fn split_points(&self) -> &[SplitPoint] {
- &self.split_points
+ self.router.split_points()
}
/// Returns the per-key sort options used for routing.
pub fn sort_options(&self) -> &[SortOptions] {
- &self.sort_options
+ self.router.sort_options()
}
}
@@ -734,11 +811,10 @@ impl PhysicalExpr for RangeExpr {
self.on_columns.len(),
children.len()
);
- Ok(Arc::new(Self::try_new_parts(
- children,
- self.split_points.clone(),
- self.sort_options.clone(),
- )?))
+ Ok(Arc::new(Self {
+ on_columns: children,
+ router: self.router.clone(),
+ }))
}
fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
@@ -750,17 +826,14 @@ impl PhysicalExpr for RangeExpr {
}
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ if self.router.split_points().is_empty() {
+ return Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(0))));
+ }
+
let arrays = evaluate_expressions_to_arrays(self.on_columns.iter(), batch)?;
- let mut row_key_buffer = Vec::with_capacity(arrays.len());
let mut partition_ids = Vec::with_capacity(batch.num_rows());
- for row_idx in 0..batch.num_rows() {
- extract_row_at_idx_to_buf(&arrays, row_idx, &mut row_key_buffer)?;
- partition_ids.push(range_partition_id(
- &row_key_buffer,
- &self.split_points,
- &self.sort_options,
- )? as u64);
- }
+ self.router
+ .route_partition_ids(&arrays, &mut partition_ids)?;
Ok(ColumnarValue::Array(Arc::new(UInt64Array::from(
partition_ids,
))))
@@ -780,12 +853,13 @@ impl PhysicalExpr for RangeExpr {
let sort_exprs = self
.on_columns
.iter()
- .zip(&self.sort_options)
+ .zip(self.router.sort_options())
.map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr), *options))
.collect::<Vec<_>>();
let sort_expr = sort_exprs_try_to_proto(&sort_exprs, ctx)?;
let split_point = self
- .split_points
+ .router
+ .split_points()
.iter()
.map(|split_point| {
let value = split_point
@@ -822,10 +896,11 @@ impl RangeExpr {
return internal_err!("PhysicalExprNode is not a RangeExpr");
};
let sort_exprs = sort_exprs_try_from_proto(&range_expr.sort_expr, ctx)?;
- let (on_columns, sort_options) = sort_exprs
- .into_iter()
- .map(|sort_expr| (sort_expr.expr, sort_expr.options))
- .unzip();
+ let (on_columns, sort_options): (Vec<PhysicalExprRef>, Vec<SortOptions>) =
+ sort_exprs
+ .into_iter()
+ .map(|sort_expr| (sort_expr.expr, sort_expr.options))
+ .unzip();
let split_points = range_expr
.split_point
.iter()
@@ -840,29 +915,13 @@ impl RangeExpr {
.collect::<Result<Vec<_>>>()?;
Ok(Arc::new(Self::try_new_parts(
on_columns,
- split_points,
- sort_options,
+ &split_points,
+ &sort_options,
+ ctx.schema(),
)?))
}
}
-fn range_partition_id(
- row_key: &[ScalarValue],
- split_points: &[SplitPoint],
- sort_options: &[SortOptions],
-) -> Result<usize> {
- let mut low = 0;
- let mut high = split_points.len();
- while low < high {
- let mid = low + (high - low) / 2;
- match compare_rows(row_key, split_points[mid].values(), sort_options)? {
- Ordering::Less => high = mid,
- Ordering::Equal | Ordering::Greater => low = mid + 1,
- }
- }
- Ok(low)
-}
-
/// Computes `value % divisor` without division in the hot loop when `divisor`
/// is fixed for many values.
///
@@ -996,24 +1055,86 @@ impl BatchPartitioner {
/// Create a new [`BatchPartitioner`] for range-based repartitioning.
///
+ /// Key data types are inferred from `schema`, and split points are coerced to match.
+ ///
/// # Parameters
/// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions
+ /// - `schema`: Schema of the batches to be partitioned
/// - `timer`: Metric used to record time spent during repartitioning.
+ pub fn try_new_range_partitioner(
+ range_partitioning: &RangePartitioning,
+ schema: &Schema,
+ timer: metrics::Time,
+ ) -> Result<Self> {
+ let ordering = range_partitioning.ordering().clone();
+ let num_partitions = range_partitioning.partition_count();
+ let sort_options: Vec<SortOptions> = ordering.iter().map(|e| e.options).collect();
+ let data_types = ordering
+ .iter()
+ .map(|e| e.expr.data_type(schema))
+ .collect::<Result<Vec<_>>>()?;
+ let router = Arc::new(RangeRouter::try_new_with_data_types(
+ &sort_options,
+ range_partitioning.split_points(),
+ &data_types,
+ )?);
+
+ Ok(Self::new_range_partitioner_with_router(
+ ordering,
+ router,
+ num_partitions,
+ timer,
+ ))
+ }
+
+ /// Create a new [`BatchPartitioner`] for range-based repartitioning.
+ ///
+ /// # Panics
+ /// Panics if the range partitioning is invalid or cannot construct a range router.
+ ///
+ /// # Workaround for backwards compatibility
+ /// This method is preserved for backwards compatibility with DataFusion 55.0 and earlier,
+ /// but does not accept a schema. Key data types are inferred from the split points, which
+ /// will fail if split points do not carry the exact precision, scale, or timezone as the
+ /// underlying columns.
+ ///
+ /// Prefer [`Self::try_new_range_partitioner`] for fallible construction with schema validation.
+ #[deprecated(
+ since = "56.0.0",
+ note = "Use try_new_range_partitioner with schema instead"
+ )]
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();
+ let router = Arc::new(
+ RangeRouter::try_new(&sort_options, range_partitioning.split_points())
+ .expect("valid range partitioning"),
+ );
+
+ Self::new_range_partitioner_with_router(ordering, router, num_partitions, timer)
+ }
+ /// Create a new [`BatchPartitioner`] for range-based repartitioning using a pre-constructed [`RangeRouter`].
+ ///
+ /// # Parameters
+ /// - `ordering`: Lexical ordering expressions for range partitioning.
+ /// - `router`: Shared [`RangeRouter`] reference.
+ /// - `num_partitions`: Total number of output partitions.
+ /// - `timer`: Metric used to record time spent during repartitioning.
+ pub(crate) fn new_range_partitioner_with_router(
+ ordering: LexOrdering,
+ router: Arc<RangeRouter>,
+ num_partitions: usize,
+ timer: metrics::Time,
+ ) -> Self {
Self {
state: BatchPartitionerState::Range {
- partition_buffer: Vec::with_capacity(ordering.len()),
ordering,
- sort_options,
- split_points,
+ router,
indices: vec![vec![]; num_partitions],
},
timer,
@@ -1052,7 +1173,9 @@ impl BatchPartitioner {
num_input_partitions,
))
}
- Partitioning::Range(range_repartitioning) => {
+ Partitioning::Range(range_repartitioning) =>
+ {
+ #[expect(deprecated)]
Ok(Self::new_range_partitioner(&range_repartitioning, timer))
}
other => {
@@ -1145,14 +1268,12 @@ impl BatchPartitioner {
}
BatchPartitionerState::Range {
ordering,
- sort_options,
- split_points,
+ router,
indices,
- partition_buffer,
} => {
// Tracking time required for distributing indexes across output partitions
let timer = self.timer.timer();
- if split_points.is_empty() {
+ if router.num_split_points() == 0 {
timer.done();
Box::new(std::iter::once(Ok((0, batch))))
} else {
@@ -1165,13 +1286,7 @@ impl BatchPartitioner {
v.clear();
}
- Self::partition_range_indices(
- &arrays,
- split_points,
- sort_options,
- partition_buffer,
- indices,
- )?;
+ router.route_indices(&arrays, indices)?;
// Finished building index-arrays for output partitions
timer.done();
@@ -1187,28 +1302,6 @@ impl BatchPartitioner {
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 partition =
- range_partition_id(row_key_buffer, split_points, sort_options)?;
- indices[partition].push(row_idx as u32)
- }
-
- Ok(())
- }
-
// return the number of output partitions
fn num_partitions(&self) -> usize {
match &self.state {
@@ -1257,6 +1350,11 @@ impl BatchPartitioner {
return Ok(vec![]);
}
+ if partition_ranges.len() == 1 && partition_ranges[0].2 == batch.num_rows() {
+ let (partition, _, _) = partition_ranges[0];
+ return Ok(vec![Ok((partition, batch.clone()))]);
+ }
+
let batches = {
let _timer = timer.timer();
let indices_array: PrimitiveArray<UInt32Type> = reordered_indices.into();
@@ -2139,16 +2237,29 @@ impl RepartitionExec {
mut stream: SendableRecordBatchStream,
mut output_channels: HashMap<usize, OutputChannel>,
partitioning: Partitioning,
+ range_router: Option<Arc<RangeRouter>>,
metrics: RepartitionMetrics,
input_partition: usize,
num_input_partitions: usize,
) -> Result<()> {
- let mut partitioner = BatchPartitioner::try_new(
- partitioning,
- metrics.repartition_time.clone(),
- input_partition,
- num_input_partitions,
- )?;
+ let mut partitioner = match (partitioning, range_router) {
+ (Partitioning::Range(range_partitioning), Some(router)) => {
+ let ordering = range_partitioning.ordering().clone();
+ let num_partitions = range_partitioning.partition_count();
+ BatchPartitioner::new_range_partitioner_with_router(
+ ordering,
+ router,
+ num_partitions,
+ metrics.repartition_time.clone(),
+ )
+ }
+ (partitioning, _) => 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();
@@ -2502,8 +2613,11 @@ mod tests {
{collect, expressions::col},
};
- use arrow::array::{ArrayRef, StringArray, UInt32Array};
- use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::array::{
+ Array, ArrayRef, Decimal128Array, Int64Array, StringArray,
+ TimestampNanosecondArray, UInt32Array,
+ };
+ use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datafusion_common::ScalarValue;
use datafusion_common::cast::{as_string_array, as_uint32_array};
use datafusion_common::exec_err;
@@ -2511,7 +2625,9 @@ mod tests {
use datafusion_common_runtime::JoinSet;
use datafusion_execution::config::SessionConfig;
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
- use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint};
+ use datafusion_physical_expr::{
+ LexOrdering, PhysicalSortExpr, RangePartitioning, SplitPoint,
+ };
use insta::assert_snapshot;
#[derive(Debug)]
@@ -2554,9 +2670,10 @@ mod tests {
.into(),
split_points.clone(),
)?;
- let expr = Arc::new(RangeExpr::try_new(
+ let expr = Arc::new(RangeExpr::try_new_with_schema(
vec![col("a", &schema)?, col("b", &schema)?],
&range_partitioning,
+ &schema,
)?);
let remapped = col("a", &schema)?;
let rewritten =
@@ -3012,6 +3129,289 @@ mod tests {
Ok(())
}
+ /// Split points are not required to carry the key's exact type: `compare_rows`,
+ /// the routing function the range router replaced, compares `Decimal128` on
+ /// scale alone and ignores precision. Routing must not depend on the split
+ /// point's precision matching the column's.
+ #[tokio::test]
+ async fn range_repartition_routes_decimal_with_wider_column_precision() -> Result<()>
+ {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "k",
+ DataType::Decimal128(20, 2),
+ true,
+ )]));
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("k", &schema)?,
+ SortOptions::default(),
+ )])
+ .unwrap();
+ // Split point is Decimal128(10, 2); the column is Decimal128(20, 2).
+ let partitioning = Partitioning::Range(RangePartitioning::try_new(
+ ordering,
+ vec![SplitPoint::new(vec![ScalarValue::Decimal128(
+ Some(1000),
+ 10,
+ 2,
+ )])],
+ )?);
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(
+ Decimal128Array::from(vec![Some(500i128), Some(2000i128)])
+ .with_precision_and_scale(20, 2)?,
+ )],
+ )?;
+
+ let output_partitions =
+ repartition(&schema, vec![vec![batch]], partitioning).await?;
+
+ assert_eq!(2, output_partitions.len());
+ assert_eq!(1, partition_row_count(&output_partitions[0]));
+ assert_eq!(1, partition_row_count(&output_partitions[1]));
+
+ Ok(())
+ }
+
+ /// Same contract for timestamps: `compare_rows` ignores the timezone, since
+ /// the underlying values are UTC either way. A tz-less split point must
+ /// still route a `Timestamp(ns, "UTC")` column — including when the key is
+ /// compound, where the single-column fast path does not apply.
+ #[tokio::test]
+ async fn range_repartition_routes_compound_timestamp_key_ignoring_timezone()
+ -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new(
+ "t",
+ DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+ true,
+ ),
+ Field::new("i", DataType::Int64, true),
+ ]));
+ let ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new(col("t", &schema)?, SortOptions::default()),
+ PhysicalSortExpr::new(col("i", &schema)?, SortOptions::default()),
+ ])
+ .unwrap();
+ // Split point timestamp carries no timezone; the column carries "UTC".
+ let partitioning = Partitioning::Range(RangePartitioning::try_new(
+ ordering,
+ vec![SplitPoint::new(vec![
+ ScalarValue::TimestampNanosecond(Some(100), None),
+ ScalarValue::Int64(Some(0)),
+ ])],
+ )?);
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(
+ TimestampNanosecondArray::from(vec![Some(50i64), Some(200)])
+ .with_timezone("UTC"),
+ ),
+ Arc::new(Int64Array::from(vec![Some(1i64), Some(2)])),
+ ],
+ )?;
+
+ let output_partitions =
+ repartition(&schema, vec![vec![batch]], partitioning).await?;
+
+ assert_eq!(2, output_partitions.len());
+ assert_eq!(1, partition_row_count(&output_partitions[0]));
+ assert_eq!(1, partition_row_count(&output_partitions[1]));
+
+ Ok(())
+ }
+
+ #[test]
+ fn range_expr_routes_decimal_with_wider_column_precision() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "k",
+ DataType::Decimal128(20, 2),
+ true,
+ )]));
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("k", &schema)?,
+ SortOptions::default(),
+ )])
+ .unwrap();
+ // Split point is Decimal128(10, 2); the column is Decimal128(20, 2).
+ let range_part = RangePartitioning::try_new(
+ ordering,
+ vec![SplitPoint::new(vec![ScalarValue::Decimal128(
+ Some(1000),
+ 10,
+ 2,
+ )])],
+ )?;
+
+ let expr = RangeExpr::try_new_with_schema(
+ vec![col("k", &schema)?],
+ &range_part,
+ &schema,
+ )?;
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(
+ Decimal128Array::from(vec![Some(500i128), Some(2000i128)])
+ .with_precision_and_scale(20, 2)?,
+ )],
+ )?;
+
+ let res = expr.evaluate(&batch)?;
+ let ColumnarValue::Array(array) = res else {
+ panic!("expected array result");
+ };
+ let partition_ids = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .values();
+ assert_eq!(partition_ids, &[0, 1]);
+
+ Ok(())
+ }
+
+ #[test]
+ fn range_expr_routes_compound_timestamp_key_ignoring_timezone() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new(
+ "t",
+ DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+ true,
+ ),
+ Field::new("i", DataType::Int64, true),
+ ]));
+ let ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new(col("t", &schema)?, SortOptions::default()),
+ PhysicalSortExpr::new(col("i", &schema)?, SortOptions::default()),
+ ])
+ .unwrap();
+ // Split point timestamp carries no timezone; the column carries "UTC".
+ let range_part = RangePartitioning::try_new(
+ ordering,
+ vec![SplitPoint::new(vec![
+ ScalarValue::TimestampNanosecond(Some(100), None),
+ ScalarValue::Int64(Some(0)),
+ ])],
+ )?;
+
+ let expr = RangeExpr::try_new_with_schema(
+ vec![col("t", &schema)?, col("i", &schema)?],
+ &range_part,
+ &schema,
+ )?;
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(
+ TimestampNanosecondArray::from(vec![Some(50i64), Some(200)])
+ .with_timezone("UTC"),
+ ),
+ Arc::new(Int64Array::from(vec![Some(1i64), Some(2)])),
+ ],
+ )?;
+
+ let res = expr.evaluate(&batch)?;
+ let ColumnarValue::Array(array) = res else {
+ panic!("expected array result");
+ };
+ let partition_ids = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .values();
+ assert_eq!(partition_ids, &[0, 1]);
+
+ Ok(())
+ }
+
+ #[test]
+ #[expect(deprecated)]
+ fn range_expr_deprecated_try_new_routes_matching_types() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, true)]));
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("k", &schema)?,
+ SortOptions::default(),
+ )])
+ .unwrap();
+ let range_part = RangePartitioning::try_new(
+ ordering,
+ vec![SplitPoint::new(vec![ScalarValue::Int64(Some(100))])],
+ )?;
+
+ // Deprecated 2-argument constructor infers data types from split points
+ let expr = RangeExpr::try_new(vec![col("k", &schema)?], &range_part)?;
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int64Array::from(vec![Some(50), Some(200)]))],
+ )?;
+
+ let res = expr.evaluate(&batch)?;
+ let ColumnarValue::Array(array) = res else {
+ panic!("expected array result");
+ };
+ let partition_ids = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .values();
+ assert_eq!(partition_ids, &[0, 1]);
+
+ Ok(())
+ }
+
+ #[test]
+ fn batch_partitioner_routes_decimal_with_wider_column_precision() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "k",
+ DataType::Decimal128(20, 2),
+ true,
+ )]));
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("k", &schema)?,
+ SortOptions::default(),
+ )])
+ .unwrap();
+ // Split point is Decimal128(10, 2); the column is Decimal128(20, 2).
+ let range_part = RangePartitioning::try_new(
+ ordering,
+ vec![SplitPoint::new(vec![ScalarValue::Decimal128(
+ Some(1000),
+ 10,
+ 2,
+ )])],
+ )?;
+
+ let mut partitioner = BatchPartitioner::try_new_range_partitioner(
+ &range_part,
+ &schema,
+ metrics::Time::default(),
+ )?;
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(
+ Decimal128Array::from(vec![Some(500i128), Some(2000i128)])
+ .with_precision_and_scale(20, 2)?,
+ )],
+ )?;
+
+ let mut partitioned_batches = vec![];
+ partitioner.partition(batch, |partition, batch| {
+ partitioned_batches.push((partition, batch.num_rows()));
+ Ok(())
+ })?;
+
+ assert_eq!(partitioned_batches, vec![(0, 1), (1, 1)]);
+
+ Ok(())
+ }
+
#[test]
fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> {
// Three columns so the projection both narrows the schema (required for
diff --git a/datafusion/physical-plan/src/repartition/range.rs b/datafusion/physical-plan/src/repartition/range.rs
new file mode 100644
index 0000000000..c05f1e531e
--- /dev/null
+++ b/datafusion/physical-plan/src/repartition/range.rs
@@ -0,0 +1,908 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Routers for range partitioning.
+
+use std::cmp::Ordering;
+use std::sync::Arc;
+
+use arrow::array::*;
+use arrow::compute::SortOptions;
+use arrow::datatypes::*;
+use arrow::row::{Row, RowConverter, Rows, SortField};
+use datafusion_common::{
+ DataFusionError, Result, ScalarValue, exec_err, not_impl_err, plan_err,
+ validate_range_split_points,
+};
+use datafusion_physical_expr::SplitPoint;
+
+/// A router for assigning rows to range partitions.
+#[derive(Debug, Clone)]
+pub(crate) struct RangeRouter {
+ data_types: Vec<DataType>,
+ split_points: Vec<SplitPoint>,
+ sort_options: Vec<SortOptions>,
+ inner: RangeRouterInner,
+}
+
+#[derive(Debug, Clone)]
+enum RangeRouterInner {
+ /// Specialized fast path for a single primitive column with non-null split points.
+ Primitive(PrimitiveRangeRouter),
+ /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys.
+ Row(RowConverterRangeRouter),
+}
+
+impl RangeRouter {
+ /// Constructs the best router for the given sort options and split points,
+ /// inferring key data types from the split points.
+ pub(crate) fn try_new(
+ sort_options: &[SortOptions],
+ split_points: &[SplitPoint],
+ ) -> Result<Self> {
+ // Pass None so `try_new_with_optional_data_types` runs `validate_range_split_points`
+ // before indexing into split point columns, avoiding a panic on width mismatch.
+ Self::try_new_with_optional_data_types(sort_options, split_points, None)
+ }
+
+ /// Constructs the best router for the given sort options, split points, and target data types,
+ /// coercing split point values to the target data types.
+ pub(crate) fn try_new_with_data_types(
+ sort_options: &[SortOptions],
+ split_points: &[SplitPoint],
+ data_types: &[DataType],
+ ) -> Result<Self> {
+ Self::try_new_with_optional_data_types(
+ sort_options,
+ split_points,
+ Some(data_types),
+ )
+ }
+
+ fn try_new_with_optional_data_types(
+ sort_options: &[SortOptions],
+ split_points: &[SplitPoint],
+ data_types: Option<&[DataType]>,
+ ) -> Result<Self> {
+ validate_range_split_points(split_points, sort_options)?;
+
+ let (data_types, split_points) = if let Some(target_dts) = data_types {
+ if target_dts.len() != sort_options.len() {
+ return plan_err!(
+ "Range partitioning expected {} data types for sort options, but got {}",
+ sort_options.len(),
+ target_dts.len()
+ );
+ }
+ let coerced_split_points = split_points
+ .iter()
+ .map(|sp| {
+ let vals = sp
+ .values()
+ .iter()
+ .zip(target_dts)
+ .map(|(val, target_dt)| {
+ if val.data_type() == *target_dt {
+ Ok(val.clone())
+ } else {
+ val.cast_to(target_dt)
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+ Ok(SplitPoint::new(vals))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ (target_dts.to_vec(), coerced_split_points)
+ } else if !split_points.is_empty() {
+ let dts: Vec<DataType> = (0..sort_options.len())
+ .map(|col_idx| split_points[0].values()[col_idx].data_type())
+ .collect();
+ (dts, split_points.to_vec())
+ } else {
+ (vec![], vec![])
+ };
+
+ // Try single-column primitive fast path
+ if data_types.len() == 1
+ && !sort_options.is_empty()
+ && let Some(primitive_router) =
+ PrimitiveRangeRouter::try_new(&split_points, sort_options[0])
+ {
+ return Ok(Self {
+ data_types,
+ split_points,
+ sort_options: sort_options.to_vec(),
+ inner: RangeRouterInner::Primitive(primitive_router),
+ });
+ }
+
+ // Try RowConverter path
+ let row_router =
+ RowConverterRangeRouter::try_new(&data_types, sort_options, &split_points)?;
+ Ok(Self {
+ data_types,
+ split_points,
+ sort_options: sort_options.to_vec(),
+ inner: RangeRouterInner::Row(row_router),
+ })
+ }
+
+ /// Data types configured in this router.
+ #[cfg(test)]
+ pub(crate) fn data_types(&self) -> &[DataType] {
+ &self.data_types
+ }
+
+ /// Split points configured in this router.
+ pub(crate) fn split_points(&self) -> &[SplitPoint] {
+ &self.split_points
+ }
+
+ /// Sort options configured in this router.
+ pub(crate) fn sort_options(&self) -> &[SortOptions] {
+ &self.sort_options
+ }
+
+ /// Number of split points configured in this router.
+ pub(crate) fn num_split_points(&self) -> usize {
+ self.split_points.len()
+ }
+
+ /// Generic routing entry point that calls `emit(row_idx, partition)` for every row.
+ pub(crate) fn route_with<E>(&self, arrays: &[ArrayRef], mut emit: E) -> Result<()>
+ where
+ E: FnMut(usize, usize),
+ {
+ if self.split_points.is_empty() {
+ let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0);
+ for row_idx in 0..num_rows {
+ emit(row_idx, 0);
+ }
+ return Ok(());
+ }
+
+ if arrays.len() != self.data_types.len() {
+ return exec_err!(
+ "Range partitioning expected {} columns, but got {}",
+ self.data_types.len(),
+ arrays.len()
+ );
+ }
+
+ for (i, (arr, expected_dt)) in arrays.iter().zip(&self.data_types).enumerate() {
+ if arr.data_type() != expected_dt {
+ return exec_err!(
+ "Range partitioning expected column {i} to be of type {expected_dt:?}, but got {:?}",
+ arr.data_type()
+ );
+ }
+ }
+
+ match &self.inner {
+ RangeRouterInner::Primitive(r) => {
+ if let Some(first_col) = arrays.first() {
+ r.route_with(first_col.as_ref(), emit)
+ } else {
+ Ok(())
+ }
+ }
+ RangeRouterInner::Row(r) => r.route_with(arrays, emit),
+ }
+ }
+
+ /// Groups row indices from `arrays` into partition index buckets.
+ pub(crate) fn route_indices(
+ &self,
+ arrays: &[ArrayRef],
+ indices: &mut [Vec<u32>],
+ ) -> Result<()> {
+ self.route_with(arrays, |row_idx, partition| {
+ indices[partition].push(row_idx as u32);
+ })
+ }
+
+ /// Appends output partition IDs to `partition_ids`.
+ pub(crate) fn route_partition_ids(
+ &self,
+ arrays: &[ArrayRef],
+ partition_ids: &mut Vec<u64>,
+ ) -> Result<()> {
+ let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0);
+ partition_ids
+ .try_reserve(num_rows)
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ self.route_with(arrays, |_row_idx, partition| {
+ partition_ids.push(partition as u64);
+ })
+ }
+}
+
+macro_rules! define_primitive_router {
+ ($( ($variant:ident, $type:ty, $arrow_type:ty, $array:ident) ),* $(,)?) => {
+ /// Specialized router for primitive scalar types.
+ #[derive(Debug, Clone)]
+ enum PrimitiveRangeRouter {
+ $( $variant(PrimitiveValuesRouter<$type>), )*
+ Float32(FloatValuesRouter<f32>),
+ Float64(FloatValuesRouter<f64>),
+ }
+
+ impl PrimitiveRangeRouter {
+ fn try_new(split_points: &[SplitPoint], sort_options: SortOptions) -> Option<Self> {
+ if split_points.is_empty() {
+ return None;
+ }
+
+ let scalars = split_points.iter().map(|sp| sp.values()[0].clone());
+ let split_array = ScalarValue::iter_to_array(scalars).ok()?;
+ if split_array.null_count() > 0 {
+ return None;
+ }
+
+ macro_rules! make_primitive {
+ ($target_arrow_type:ty, $target_variant:ident) => {{
+ let arr = split_array
+ .as_any()
+ .downcast_ref::<PrimitiveArray<$target_arrow_type>>()?;
+ let vals = arr.values().to_vec();
+ Some(Self::$target_variant(PrimitiveValuesRouter::new(vals, sort_options)))
+ }};
+ }
+
+ match split_array.data_type() {
+ DataType::Int8 => make_primitive!(Int8Type, Int8),
+ DataType::Int16 => make_primitive!(Int16Type, Int16),
+ DataType::Int32 => make_primitive!(Int32Type, Int32),
+ DataType::Int64 => make_primitive!(Int64Type, Int64),
+ DataType::UInt8 => make_primitive!(UInt8Type, UInt8),
+ DataType::UInt16 => make_primitive!(UInt16Type, UInt16),
+ DataType::UInt32 => make_primitive!(UInt32Type, UInt32),
+ DataType::UInt64 => make_primitive!(UInt64Type, UInt64),
+ DataType::Date32 => make_primitive!(Date32Type, Date32),
+ DataType::Date64 => make_primitive!(Date64Type, Date64),
+ DataType::Time32(TimeUnit::Second) => make_primitive!(Time32SecondType, Time32Second),
+ DataType::Time32(TimeUnit::Millisecond) => make_primitive!(Time32MillisecondType, Time32Millisecond),
+ DataType::Time64(TimeUnit::Microsecond) => make_primitive!(Time64MicrosecondType, Time64Microsecond),
+ DataType::Time64(TimeUnit::Nanosecond) => make_primitive!(Time64NanosecondType, Time64Nanosecond),
+ DataType::Timestamp(TimeUnit::Second, _) => make_primitive!(TimestampSecondType, TimestampSecond),
+ DataType::Timestamp(TimeUnit::Millisecond, _) => make_primitive!(TimestampMillisecondType, TimestampMillisecond),
+ DataType::Timestamp(TimeUnit::Microsecond, _) => make_primitive!(TimestampMicrosecondType, TimestampMicrosecond),
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => make_primitive!(TimestampNanosecondType, TimestampNanosecond),
+ DataType::Float32 => {
+ let arr = split_array.as_any().downcast_ref::<Float32Array>()?;
+ let vals = arr.values().to_vec();
+ Some(Self::Float32(FloatValuesRouter::new(vals, sort_options)))
+ }
+ DataType::Float64 => {
+ let arr = split_array.as_any().downcast_ref::<Float64Array>()?;
+ let vals = arr.values().to_vec();
+ Some(Self::Float64(FloatValuesRouter::new(vals, sort_options)))
+ }
+ _ => None,
+ }
+ }
+
+ fn route_with<E>(&self, array: &dyn Array, emit: E) -> Result<()>
+ where
+ E: FnMut(usize, usize),
+ {
+ match self {
+ $(
+ Self::$variant(r) => {
+ let arr = array.as_any().downcast_ref::<$array>().ok_or_else(|| {
+ DataFusionError::Internal(format!("Expected {}", stringify!($array)))
+ })?;
+ r.route_with(arr, emit);
+ Ok(())
+ }
+ )*
+ Self::Float32(r) => {
+ let arr = array.as_any().downcast_ref::<Float32Array>().ok_or_else(|| {
+ DataFusionError::Internal("Expected Float32Array".to_string())
+ })?;
+ r.route_with(arr, emit);
+ Ok(())
+ }
+ Self::Float64(r) => {
+ let arr = array.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
+ DataFusionError::Internal("Expected Float64Array".to_string())
+ })?;
+ r.route_with(arr, emit);
+ Ok(())
+ }
+ }
+ }
+ }
+ };
+}
+
+define_primitive_router!(
+ (Int8, i8, Int8Type, Int8Array),
+ (Int16, i16, Int16Type, Int16Array),
+ (Int32, i32, Int32Type, Int32Array),
+ (Int64, i64, Int64Type, Int64Array),
+ (UInt8, u8, UInt8Type, UInt8Array),
+ (UInt16, u16, UInt16Type, UInt16Array),
+ (UInt32, u32, UInt32Type, UInt32Array),
+ (UInt64, u64, UInt64Type, UInt64Array),
+ (Date32, i32, Date32Type, Date32Array),
+ (Date64, i64, Date64Type, Date64Array),
+ (Time32Second, i32, Time32SecondType, Time32SecondArray),
+ (
+ Time32Millisecond,
+ i32,
+ Time32MillisecondType,
+ Time32MillisecondArray
+ ),
+ (
+ Time64Microsecond,
+ i64,
+ Time64MicrosecondType,
+ Time64MicrosecondArray
+ ),
+ (
+ Time64Nanosecond,
+ i64,
+ Time64NanosecondType,
+ Time64NanosecondArray
+ ),
+ (
+ TimestampSecond,
+ i64,
+ TimestampSecondType,
+ TimestampSecondArray
+ ),
+ (
+ TimestampMillisecond,
+ i64,
+ TimestampMillisecondType,
+ TimestampMillisecondArray
+ ),
+ (
+ TimestampMicrosecond,
+ i64,
+ TimestampMicrosecondType,
+ TimestampMicrosecondArray
+ ),
+ (
+ TimestampNanosecond,
+ i64,
+ TimestampNanosecondType,
+ TimestampNanosecondArray
+ ),
+);
+
+/// Generic router for primitive integer and temporal types.
+#[derive(Debug, Clone)]
+struct PrimitiveValuesRouter<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> {
+ split_points: Vec<T>,
+ sort_options: SortOptions,
+}
+
+impl<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> PrimitiveValuesRouter<T> {
+ fn new(split_points: Vec<T>, sort_options: SortOptions) -> Self {
+ Self {
+ split_points,
+ sort_options,
+ }
+ }
+
+ fn route_with<A: ArrowPrimitiveType<Native = T>, E: FnMut(usize, usize)>(
+ &self,
+ array: &PrimitiveArray<A>,
+ mut emit: E,
+ ) {
+ let split_points = &self.split_points;
+ let descending = self.sort_options.descending;
+ let nulls_first = self.sort_options.nulls_first;
+
+ if array.null_count() == 0 {
+ let values = array.values().as_ref();
+ if !descending {
+ for (idx, &val) in values.iter().enumerate() {
+ let p = split_points.partition_point(|&sp| sp <= val);
+ emit(idx, p);
+ }
+ } else {
+ for (idx, &val) in values.iter().enumerate() {
+ let p = split_points.partition_point(|&sp| sp >= val);
+ emit(idx, p);
+ }
+ }
+ } else {
+ let null_partition = if nulls_first { 0 } else { split_points.len() };
+ for idx in 0..array.len() {
+ if array.is_null(idx) {
+ emit(idx, null_partition);
+ } else {
+ let val = array.value(idx);
+ let p = if !descending {
+ split_points.partition_point(|&sp| sp <= val)
+ } else {
+ split_points.partition_point(|&sp| sp >= val)
+ };
+ emit(idx, p);
+ }
+ }
+ }
+ }
+}
+
+/// Generic router for floating point values using total ordering.
+#[derive(Debug, Clone)]
+struct FloatValuesRouter<T: Copy + Send + Sync + 'static> {
+ split_points: Vec<T>,
+ sort_options: SortOptions,
+}
+
+impl<T: Copy + Send + Sync + 'static> FloatValuesRouter<T> {
+ fn new(split_points: Vec<T>, sort_options: SortOptions) -> Self {
+ Self {
+ split_points,
+ sort_options,
+ }
+ }
+}
+
+macro_rules! impl_float_values_router {
+ ($t:ty, $arr:ty) => {
+ impl FloatValuesRouter<$t> {
+ fn route_with<E: FnMut(usize, usize)>(&self, array: &$arr, mut emit: E) {
+ let split_points = &self.split_points;
+ let descending = self.sort_options.descending;
+ let nulls_first = self.sort_options.nulls_first;
+
+ if array.null_count() == 0 {
+ let values = array.values().as_ref();
+ if !descending {
+ for (idx, &val) in values.iter().enumerate() {
+ let p = split_points.partition_point(|&sp| {
+ sp.total_cmp(&val) != Ordering::Greater
+ });
+ emit(idx, p);
+ }
+ } else {
+ for (idx, &val) in values.iter().enumerate() {
+ let p = split_points.partition_point(|&sp| {
+ sp.total_cmp(&val) != Ordering::Less
+ });
+ emit(idx, p);
+ }
+ }
+ } else {
+ let null_partition = if nulls_first { 0 } else { split_points.len() };
+ for idx in 0..array.len() {
+ if array.is_null(idx) {
+ emit(idx, null_partition);
+ } else {
+ let val = array.value(idx);
+ let p = if !descending {
+ split_points.partition_point(|&sp| {
+ sp.total_cmp(&val) != Ordering::Greater
+ })
+ } else {
+ split_points.partition_point(|&sp| {
+ sp.total_cmp(&val) != Ordering::Less
+ })
+ };
+ emit(idx, p);
+ }
+ }
+ }
+ }
+ }
+ };
+}
+
+impl_float_values_router!(f32, Float32Array);
+impl_float_values_router!(f64, Float64Array);
+
+/// Router backed by Arrow's RowConverter.
+#[derive(Debug, Clone)]
+struct RowConverterRangeRouter {
+ converter: Arc<RowConverter>,
+ split_point_rows: Option<Rows>,
+}
+
+impl RowConverterRangeRouter {
+ fn try_new(
+ data_types: &[DataType],
+ sort_options: &[SortOptions],
+ split_points: &[SplitPoint],
+ ) -> Result<Self> {
+ let sort_fields = data_types
+ .iter()
+ .zip(sort_options)
+ .map(|(dt, opt)| SortField::new_with_options(dt.clone(), *opt))
+ .collect::<Vec<_>>();
+
+ if !RowConverter::supports_fields(&sort_fields) {
+ return not_impl_err!(
+ "Range partitioning is not supported for data types: {:?}",
+ data_types
+ );
+ }
+
+ let row_converter = RowConverter::new(sort_fields)?;
+ let num_cols = data_types.len();
+
+ let split_point_rows = if split_points.is_empty() {
+ None
+ } else {
+ let split_point_arrays = (0..num_cols)
+ .map(|col_idx| {
+ let col_scalars =
+ split_points.iter().map(|sp| sp.values()[col_idx].clone());
+ ScalarValue::iter_to_array(col_scalars)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ Some(row_converter.convert_columns(&split_point_arrays)?)
+ };
+
+ Ok(Self {
+ converter: Arc::new(row_converter),
+ split_point_rows,
+ })
+ }
+
+ fn route_with<E: FnMut(usize, usize)>(
+ &self,
+ arrays: &[ArrayRef],
+ mut emit: E,
+ ) -> Result<()> {
+ let rows = self.converter.convert_columns(arrays)?;
+ if let Some(sp_rows) = &self.split_point_rows {
+ for (row_idx, row) in rows.iter().enumerate() {
+ let partition = partition_point_rows(sp_rows, &row);
+ emit(row_idx, partition);
+ }
+ } else {
+ for row_idx in 0..rows.num_rows() {
+ emit(row_idx, 0);
+ }
+ }
+ Ok(())
+ }
+}
+
+#[inline]
+fn partition_point_rows(sp_rows: &Rows, row: &Row<'_>) -> usize {
+ let mut left = 0;
+ let mut right = sp_rows.num_rows();
+ while left < right {
+ let mid = left + (right - left) / 2;
+ if sp_rows.row(mid) <= *row {
+ left = mid + 1;
+ } else {
+ right = mid;
+ }
+ }
+ left
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::Arc;
+
+ fn make_split_points_1d(scalars: Vec<ScalarValue>) -> Vec<SplitPoint> {
+ scalars
+ .into_iter()
+ .map(|s| SplitPoint::new(vec![s]))
+ .collect()
+ }
+
+ fn assert_routing(
+ router: &RangeRouter,
+ arrays: &[ArrayRef],
+ expected_partition_ids: &[u64],
+ expected_indices: Option<&[Vec<u32>]>,
+ ) -> Result<()> {
+ let mut partition_ids = Vec::new();
+ router.route_partition_ids(arrays, &mut partition_ids)?;
+ assert_eq!(partition_ids, expected_partition_ids);
+
+ if let Some(expected) = expected_indices {
+ let mut indices = vec![vec![]; expected.len()];
+ router.route_indices(arrays, &mut indices)?;
+ assert_eq!(indices, expected);
+ }
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_primitive_router_i64_asc() -> Result<()> {
+ let split_points = make_split_points_1d(vec![
+ ScalarValue::Int64(Some(10)),
+ ScalarValue::Int64(Some(20)),
+ ScalarValue::Int64(Some(30)),
+ ]);
+ let sort_options = vec![SortOptions {
+ descending: false,
+ nulls_first: true,
+ }];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+ assert!(matches!(router.inner, RangeRouterInner::Primitive(_)));
+
+ let input = Arc::new(Int64Array::from(vec![
+ Some(5),
+ Some(10),
+ Some(15),
+ Some(20),
+ Some(25),
+ Some(30),
+ Some(35),
+ None,
+ ])) as ArrayRef;
+
+ assert_routing(
+ &router,
+ &[input],
+ &[0, 1, 1, 2, 2, 3, 3, 0],
+ Some(&[vec![0, 7], vec![1, 2], vec![3, 4], vec![5, 6]]),
+ )
+ }
+
+ #[test]
+ fn test_primitive_router_i64_desc() -> Result<()> {
+ let split_points = make_split_points_1d(vec![
+ ScalarValue::Int64(Some(30)),
+ ScalarValue::Int64(Some(20)),
+ ScalarValue::Int64(Some(10)),
+ ]);
+ let sort_options = vec![SortOptions {
+ descending: true,
+ nulls_first: false,
+ }];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+ assert!(matches!(router.inner, RangeRouterInner::Primitive(_)));
+
+ let input = Arc::new(Int64Array::from(vec![
+ Some(35),
+ Some(30),
+ Some(25),
+ Some(20),
+ Some(15),
+ Some(10),
+ Some(5),
+ None,
+ ])) as ArrayRef;
+
+ assert_routing(
+ &router,
+ &[input],
+ &[0, 1, 1, 2, 2, 3, 3, 3],
+ Some(&[vec![0], vec![1, 2], vec![3, 4], vec![5, 6, 7]]),
+ )
+ }
+
+ #[test]
+ fn test_float_router() -> Result<()> {
+ let split_points = make_split_points_1d(vec![
+ ScalarValue::Float64(Some(0.0)),
+ ScalarValue::Float64(Some(100.0)),
+ ]);
+ let sort_options = vec![SortOptions {
+ descending: false,
+ nulls_first: false,
+ }];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+ assert!(matches!(router.inner, RangeRouterInner::Primitive(_)));
+
+ let input = Arc::new(Float64Array::from(vec![
+ Some(-10.0),
+ Some(0.0),
+ Some(50.0),
+ Some(100.0),
+ Some(200.0),
+ None,
+ ])) as ArrayRef;
+
+ assert_routing(
+ &router,
+ &[input],
+ &[0, 1, 1, 2, 2, 2],
+ Some(&[vec![0], vec![1, 2], vec![3, 4, 5]]),
+ )
+ }
+
+ #[test]
+ fn test_row_converter_strings() -> Result<()> {
+ let split_points = make_split_points_1d(vec![
+ ScalarValue::Utf8(Some("d".to_string())),
+ ScalarValue::Utf8(Some("m".to_string())),
+ ScalarValue::Utf8(Some("s".to_string())),
+ ]);
+ let sort_options = vec![SortOptions {
+ descending: false,
+ nulls_first: true,
+ }];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+ assert!(matches!(router.inner, RangeRouterInner::Row(_)));
+
+ let input = Arc::new(StringArray::from(vec![
+ Some("apple"),
+ Some("d"),
+ Some("frog"),
+ Some("m"),
+ Some("orange"),
+ Some("s"),
+ Some("zebra"),
+ None,
+ ])) as ArrayRef;
+
+ assert_routing(
+ &router,
+ &[input],
+ &[0, 1, 1, 2, 2, 3, 3, 0],
+ Some(&[vec![0, 7], vec![1, 2], vec![3, 4], vec![5, 6]]),
+ )
+ }
+
+ #[test]
+ fn test_row_converter_composite_keys() -> Result<()> {
+ let split_points = vec![
+ SplitPoint::new(vec![
+ ScalarValue::Int64(Some(1)),
+ ScalarValue::Utf8(Some("b".to_string())),
+ ]),
+ SplitPoint::new(vec![
+ ScalarValue::Int64(Some(1)),
+ ScalarValue::Utf8(Some("d".to_string())),
+ ]),
+ SplitPoint::new(vec![
+ ScalarValue::Int64(Some(2)),
+ ScalarValue::Utf8(Some("a".to_string())),
+ ]),
+ ];
+ let sort_options = vec![
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ ];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+ assert!(matches!(router.inner, RangeRouterInner::Row(_)));
+
+ let col1 = Arc::new(Int64Array::from(vec![1, 1, 1, 1, 2, 2, 3])) as ArrayRef;
+ let col2 = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "a", "z", "a"]))
+ as ArrayRef;
+
+ assert_routing(
+ &router,
+ &[col1, col2],
+ &[0, 1, 1, 2, 3, 3, 3],
+ Some(&[vec![0], vec![1, 2], vec![3], vec![4, 5, 6]]),
+ )
+ }
+
+ #[test]
+ fn test_router_empty_split_points() -> Result<()> {
+ let split_points = vec![];
+ let sort_options = vec![SortOptions::default()];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+ assert_eq!(router.num_split_points(), 0);
+
+ let input = Arc::new(Int64Array::from(vec![10, 20, 30])) as ArrayRef;
+ assert_routing(&router, &[input], &[0, 0, 0], Some(&[vec![0, 1, 2]]))
+ }
+
+ #[test]
+ fn test_router_decimal_precision_widening() -> Result<()> {
+ // Split point defined with precision 10, scale 2
+ let split_points = vec![SplitPoint::new(vec![ScalarValue::Decimal128(
+ Some(1000),
+ 10,
+ 2,
+ )])];
+ let sort_options = vec![SortOptions::default()];
+ let target_data_types = vec![DataType::Decimal128(20, 2)];
+
+ let router = RangeRouter::try_new_with_data_types(
+ &sort_options,
+ &split_points,
+ &target_data_types,
+ )?;
+ assert_eq!(router.data_types(), &target_data_types);
+
+ // Column array is Decimal128(20, 2)
+ let array = Arc::new(
+ Decimal128Array::from(vec![Some(500i128), Some(1000i128), Some(2000i128)])
+ .with_precision_and_scale(20, 2)?,
+ ) as ArrayRef;
+
+ assert_routing(&router, &[array], &[0, 1, 1], Some(&[vec![0], vec![1, 2]]))
+ }
+
+ #[test]
+ fn test_router_timestamp_timezone_coercion() -> Result<()> {
+ // Split point defined without timezone
+ let split_points = vec![SplitPoint::new(vec![
+ ScalarValue::TimestampNanosecond(Some(100), None),
+ ScalarValue::Int64(Some(0)),
+ ])];
+ let sort_options = vec![SortOptions::default(), SortOptions::default()];
+ let target_data_types = vec![
+ DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+ DataType::Int64,
+ ];
+
+ let router = RangeRouter::try_new_with_data_types(
+ &sort_options,
+ &split_points,
+ &target_data_types,
+ )?;
+ assert_eq!(router.data_types(), &target_data_types);
+
+ let col1 = Arc::new(
+ TimestampNanosecondArray::from(vec![Some(50i64), Some(200i64)])
+ .with_timezone("UTC"),
+ ) as ArrayRef;
+ let col2 = Arc::new(Int64Array::from(vec![Some(1i64), Some(2i64)])) as ArrayRef;
+
+ assert_routing(&router, &[col1, col2], &[0, 1], Some(&[vec![0], vec![1]]))
+ }
+
+ #[test]
+ fn test_router_type_mismatch_error() -> Result<()> {
+ let split_points = make_split_points_1d(vec![ScalarValue::Int64(Some(10))]);
+ let sort_options = vec![SortOptions::default()];
+
+ let router = RangeRouter::try_new(&sort_options, &split_points)?;
+
+ // Pass Float64Array instead of Int64Array
+ let invalid_input = Arc::new(Float64Array::from(vec![5.0, 15.0])) as ArrayRef;
+ let err = router.route_with(&[invalid_input], |_, _| {}).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("Range partitioning expected column 0 to be of type Int64")
+ );
+
+ // Pass wrong column count
+ let err = router.route_with(&[], |_, _| {}).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("Range partitioning expected 1 columns, but got 0")
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_router_width_mismatch_error() {
+ let split_points = vec![SplitPoint::new(vec![ScalarValue::Int64(Some(10))])];
+ // 2 sort options but split point only has 1 column
+ let sort_options = vec![SortOptions::default(), SortOptions::default()];
+
+ let err = RangeRouter::try_new(&sort_options, &split_points).unwrap_err();
+ assert!(err.to_string().contains(
+ "Range partitioning split point 0 has width 1, but ordering has width 2"
+ ));
+ }
+}
diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs
index 5fce0231a6..f2b14b0439 100644
--- a/datafusion/proto/tests/cases/plans/exprs.rs
+++ b/datafusion/proto/tests/cases/plans/exprs.rs
@@ -196,11 +196,12 @@ fn roundtrip_range_expr() -> Result<()> {
ScalarValue::Float64(Some(1.0)),
])],
)?;
- let range_expr: Arc<dyn PhysicalExpr> = Arc::new(RangeExpr::try_new(
+ let range_expr: Arc<dyn PhysicalExpr> = Arc::new(RangeExpr::try_new_with_schema(
// Expression remapping may produce duplicate children. Preserve both
// so their sort options stay aligned with the split-point values.
vec![col("a", &schema)?, col("a", &schema)?],
&range_partitioning,
+ &schema,
)?);
let filter_expr = binary(range_expr, Operator::Eq, lit(0u64), &schema)?;
let plan = Arc::new(FilterExec::try_new(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org