milenkovicm commented on code in PR #19668:
URL: https://github.com/apache/datafusion/pull/19668#discussion_r2666191787
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -437,30 +437,58 @@ impl BatchPartitioner {
/// Create a new [`BatchPartitioner`] with the provided [`Partitioning`]
///
/// The time spent repartitioning will be recorded to `timer`
+ /// Create a new [`BatchPartitioner`] for hash partitioning
+ pub fn try_new_hash(
+ exprs: Vec<Arc<dyn PhysicalExpr>>,
+ num_partitions: usize,
+ timer: metrics::Time,
+ ) -> Result<Self> {
+ Ok(Self {
+ state: BatchPartitionerState::Hash {
+ exprs,
+ num_partitions,
+ hash_buffer: vec![],
+ },
+ timer,
+ })
+ }
+
+ /// Create a new [`BatchPartitioner`] for round-robin partitioning
+ pub fn try_new_round_robin(
+ num_partitions: usize,
+ timer: metrics::Time,
+ input_partition: usize,
+ num_input_partitions: usize,
+ ) -> Result<Self> {
+ Ok(Self {
+ state: BatchPartitionerState::RoundRobin {
+ num_partitions,
+ next_idx: (input_partition * num_partitions) /
num_input_partitions,
+ },
+ timer,
+ })
+ }
+
Review Comment:
can we please add documentation for try new as well?
--
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]
