Dandandan commented on code in PR #24034:
URL: https://github.com/apache/datafusion/pull/24034#discussion_r3707000327
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -746,42 +754,66 @@ impl LinearSearch {
}
}
- /// Calculate indices of each partition (according to PARTITION BY
expression)
- /// `columns` contain partition by expression results.
- fn get_per_partition_indices(
+ /// Splits the rows of `batch` by partition, according to the PARTITION BY
+ /// expression results in `columns`. Returns the distinct partition keys
+ /// in first-appearance order, a permutation of the row indices of
+ /// `batch` that groups each partition's rows together, and the
+ /// boundaries of each partition's run of rows within that permutation:
+ /// partition `p` occupies `permutation[bounds[p]..bounds[p + 1]]`, and
+ /// its indices are in ascending (stream) order.
+ fn compute_partition_permutation(
&mut self,
columns: &[ArrayRef],
batch: &RecordBatch,
- ) -> Result<Vec<(PartitionKey, Vec<u32>)>> {
- let mut batch_hashes = vec![0; batch.num_rows()];
+ ) -> Result<(Vec<PartitionKey>, Vec<u32>, Vec<usize>)> {
+ let num_rows = batch.num_rows();
+ let mut batch_hashes = vec![0; num_rows];
create_hashes(columns, &self.random_state, &mut batch_hashes)?;
self.input_buffer_hashes.extend(&batch_hashes);
// reset row_map for new calculation
self.row_map_batch.clear();
- // res stores PartitionKey and row indices (indices where these
partition occurs in the `batch`) for each partition.
- let mut result: Vec<(PartitionKey, Vec<u32>)> = vec![];
+ let mut keys: Vec<PartitionKey> = vec![];
+ // Partition id of each row, in row order:
+ let mut row_partition_ids = Vec::with_capacity(num_rows);
+ // Number of rows in each partition:
+ let mut counts: Vec<usize> = vec![];
for (hash, row_idx) in batch_hashes.into_iter().zip(0u32..) {
let entry = self.row_map_batch.find_mut(hash, |(_, group_idx)| {
- // We can safely get the first index of the partition indices
- // since partition indices has one element during
initialization.
let row = get_row_at_idx(columns, row_idx as usize).unwrap();
- // Handle hash collusions with an equality check:
- row.eq(&result[*group_idx].0)
+ // Handle hash collisions with an equality check:
+ row.eq(&keys[*group_idx])
Review Comment:
Can't `==` be used here?
--
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]
