Dandandan commented on code in PR #13133:
URL: https://github.com/apache/datafusion/pull/13133#discussion_r1818551635
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -327,16 +409,79 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
self.loser_tree_adjusted = true;
}
- /// Attempts to update the loser tree, following winner replacement, if
possible
+ /// Resets the poll count by incrementing the reset epoch.
+ fn reset_poll_counts(&mut self) {
+ self.current_reset_epoch += 1;
+ }
+
+ /// Handles tie-breaking logic during the adjustment of the loser tree.
+ ///
+ /// When comparing elements from multiple partitions in the
`update_loser_tree` process, a tie can occur
+ /// between the current winner and a challenger. This function is invoked
when such a tie needs to be
+ /// resolved according to the round-robin tie-breaker mode.
+ ///
+ /// If round-robin tie-breaking is not active, it is enabled, and the poll
counts for all elements are reset.
+ /// The function then compares the poll counts of the current winner and
the challenger:
+ /// - If the winner remains at the top after the final comparison, it
increments the winner's poll count.
+ /// - If the challenger has a lower poll count than the current winner,
the challenger becomes the new winner.
+ /// - If the poll counts are equal but the challenger's index is smaller,
the challenger is preferred.
+ ///
+ /// # Parameters
+ /// - `cmp_node`: The index of the comparison node in the loser tree where
the tie-breaking is happening.
+ /// - `winner`: A mutable reference to the current winner, which may be
updated based on the tie-breaking result.
+ /// - `challenger`: The index of the challenger being compared against the
winner.
+ ///
+ /// This function ensures fair selection among elements with equal values
when tie-breaking mode is enabled,
+ /// aiming to balance the polling across different partitions.
+ #[inline]
+ fn handle_tie(&mut self, cmp_node: usize, winner: &mut usize, challenger:
usize) {
+ if !self.round_robin_tie_breaker_mode {
+ self.round_robin_tie_breaker_mode = true;
+ // Reset poll count for tie-breaker
+ self.reset_poll_counts();
+ }
+ // Update poll count if the winner survives in the final match
+ if *winner == self.loser_tree[0] {
+ self.update_poll_count_on_the_same_value(*winner);
+ if self.is_poll_count_gt(*winner, challenger) {
+ self.update_winner(cmp_node, winner, challenger);
+ }
+ } else if challenger < *winner {
+ // The winner loses, assign the new winner
+ self.update_winner(cmp_node, winner, challenger);
+ }
+ }
+
+ /// Updates the loser tree to reflect the new winner after the previous
winner is consumed.
+ /// This function adjusts the tree by comparing the current winner with
challengers from
+ /// other partitions.
+ ///
+ /// If `enable_round_robin_tie_breaker` is true and a tie occurs at the
final level, the
+ /// tie-breaker logic will be applied to ensure fair selection among equal
elements.
fn update_loser_tree(&mut self) {
+ // Start with the current winner
let mut winner = self.loser_tree[0];
- // Replace overall winner by walking tree of losers
+
+ // Find the leaf node index of the winner in the loser tree.
let mut cmp_node = self.lt_leaf_node_index(winner);
+
+ // Traverse up the tree to adjust comparisons until reaching the root.
while cmp_node != 0 {
let challenger = self.loser_tree[cmp_node];
- if self.is_gt(winner, challenger) {
- self.loser_tree[cmp_node] = winner;
- winner = challenger;
+ // If round-robin tie-breaker is enabled and we're at the final
comparison (cmp_node == 1)
+ if self.enable_round_robin_tie_breaker && cmp_node == 1 {
+ let (is_eq, is_gt) = self.is_eq_and_gt(winner, challenger);
+ if is_eq {
+ self.handle_tie(cmp_node, &mut winner, challenger);
Review Comment:
I wonder whether it will be possible somehow to combine `is_eq_and_gt` and
`handle_tie`? It seems it could be more efficient?
--
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]
