gene-bordegaray commented on code in PR #23231:
URL: https://github.com/apache/datafusion/pull/23231#discussion_r3537111722
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -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 - after this returns,
`indices[p]` contains the row indices from `arrays`
Review Comment:
nit:
```rust
/// 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`.
```
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2324,34 +2734,43 @@ mod tests {
}
#[tokio::test]
- async fn unsupported_range_partitioning() -> Result<()> {
+ async fn range_repartition_routes_string_rows() -> Result<()> {
Review Comment:
nit: could we move this guy up with the other tests
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -1652,12 +1758,13 @@ impl RepartitionExec {
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"
- );
+ Partitioning::Range(range_partitioning) => {
Review Comment:
yes I dont see why not
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2241,6 +2533,124 @@ mod tests {
Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
}
+ fn nullable_u32_schema() -> Arc<Schema> {
Review Comment:
Thanks for adding the coverag, this is in much better shape now.
One thing on the test helpers. I think they may be a bit more than needed.
The important cases to cover are:
- single-key and compund-key routing
- ASC / DESC behavior
- null ordering
- non-primitive/string routing
For the compound-key test, I don’t think it needs to use `Int32` and this
causes us to have an extra helper here. I think we can just switch that guy to
`UInt32` and eliminate helpers.
I think we can also simplifying the `u32` collectors into one nullable-aware
helper and either dropping `assert_partition_contains_only_u32_values` or
making the assertions more direct.
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -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 - after this returns,
`indices[p]` contains 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 hash and range routers first fills one index vector per output
partition. This method
Review Comment:
nit: maybe just say router
--
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]
