jayzhan211 commented on code in PR #24598:
URL: https://github.com/apache/datafusion/pull/24598#discussion_r3872391420
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -999,25 +1009,31 @@ impl BatchPartitioner {
/// # 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(
+ pub fn try_new_range_partitioner(
range_partitioning: &RangePartitioning,
timer: metrics::Time,
- ) -> Self {
+ ) -> Result<Self> {
let ordering = range_partitioning.ordering().clone();
- let split_points = range_partitioning.split_points().to_vec();
+ let split_points = range_partitioning.split_points();
let num_partitions = range_partitioning.partition_count();
let sort_options: Vec<SortOptions> = ordering.iter().map(|e|
e.options).collect();
+ let data_types: Vec<DataType> = if !split_points.is_empty() {
+ (0..ordering.len())
+ .map(|col_idx| split_points[0].values()[col_idx].data_type())
+ .collect()
+ } else {
+ vec![]
+ };
Review Comment:
```rs
/// 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(())
}
```
Here are the edge cases
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
