Struct SortPreservingMergeStream 

pub(crate) struct SortPreservingMergeStream<C: CursorValues> {

    in_progress: BatchBuilder,
    streams: Box<dyn PartitionedStream<Output = Result<(C, RecordBatch)>>>,
    metrics: BaselineMetrics,
    done: bool,
    loser_tree: Vec<usize>,
    loser_tree_adjusted: bool,
    batch_size: usize,
    cursors: Vec<Option<Cursor<C>>>,
    enable_round_robin_tie_breaker: bool,
    round_robin_tie_breaker_mode: bool,
    num_of_polled_with_same_value: Vec<usize>,
    poll_reset_epochs: Vec<usize>,
    current_reset_epoch: usize,
    prev_cursors: Vec<Option<Cursor<C>>>,
    fetch: Option<usize>,
    produced: usize,
    uninitiated_partitions: Vec<usize>,
}

Merges a stream of sorted cursors and record batches into a single sorted stream

Fields

in_progress: BatchBuilder

streams: Box<dyn PartitionedStream<Output = Result<(C, RecordBatch)>>>

The sorted input streams to merge together

metrics: BaselineMetrics

used to record execution metrics

done: bool

If the stream has encountered an error or reaches the fetch limit.

loser_tree: Vec<usize>

A loser tree that always produces the minimum cursor

Node 0 stores the top winner, Nodes 1..num_streams store the loser nodes

This implements a “Tournament Tree” (aka Loser Tree) to keep track of the current smallest element at the top. When the top record is taken, the tree structure is not modified, and only the path from bottom to top is visited, keeping the number of comparisons close to the theoretical limit of log(S).

The current implementation uses a vector to store the tree. Conceptually, it looks like this (assuming 8 streams):

    0 (winner)

    1
   / \
  2   3
 / \ / \
4  5 6  7

Where element at index 0 in the vector is the current winner. Element at index 1 is the root of the loser tree, element at index 2 is the left child of the root, and element at index 3 is the right child of the root and so on.

reference: https://en.wikipedia.org/wiki/K-way_merge_algorithm#Tournament_Tree

loser_tree_adjusted: bool

If the most recently yielded overall winner has been replaced within the loser tree. A value of false indicates that the overall winner has been yielded but the loser tree has not been updated

batch_size: usize

Target batch size

cursors: Vec<Option<Cursor<C>>>

Cursors for each input partition. None means the input is exhausted

enable_round_robin_tie_breaker: bool

Configuration parameter to enable round-robin selection of tied winners of loser tree.

This option controls the tie-breaker strategy and attempts to avoid the issue of unbalanced polling between partitions

If true, when multiple partitions have the same value, the partition that has the fewest poll counts is selected. This strategy ensures that multiple partitions with the same value are chosen equally, distributing the polling load in a round-robin fashion. This approach balances the workload more effectively across partitions and avoids excessive buffer growth.

if false, partitions with smaller indices are consistently chosen as the winners, which can lead to an uneven distribution of polling and potentially causing upstream operator buffers for the other partitions to grow excessively, as they continued receiving data without consuming it.

For example, an upstream operator like RepartitionExec execution would keep sending data to certain partitions, but those partitions wouldn’t consume the data if they weren’t selected as winners. This resulted in inefficient buffer usage.

round_robin_tie_breaker_mode: bool

Flag indicating whether we are in the mode of round-robin tie breaker for the loser tree winners.

num_of_polled_with_same_value: Vec<usize>

Total number of polls returning the same value, as per partition. We select the one that has less poll counts for tie-breaker in loser tree.

poll_reset_epochs: Vec<usize>

To keep track of reset counts

current_reset_epoch: usize

Current reset count

prev_cursors: Vec<Option<Cursor<C>>>

Stores the previous value of each partitions for tracking the poll counts on the same value.

fetch: Option<usize>

Optional number of rows to fetch

produced: usize

number of rows produced

uninitiated_partitions: Vec<usize>

This vector contains the indices of the partitions that have not started emitting yet.

Implementations

impl<C: CursorValues> SortPreservingMergeStream<C>

pub(crate) fn new( streams: Box<dyn PartitionedStream<Output = Result<(C, RecordBatch)>>>, schema: SchemaRef, metrics: BaselineMetrics, batch_size: usize, fetch: Option<usize>, reservation: MemoryReservation, enable_round_robin_tie_breaker: bool, ) -> Self

fn maybe_poll_stream( &mut self, cx: &mut Context<'_>, idx: usize, ) -> Poll<Result<()>>

If the stream at the given index is not exhausted, and the last cursor for the stream is finished, poll the stream for the next RecordBatch and create a new cursor for the stream from the returned result

fn poll_next_inner( &mut self, cx: &mut Context<'_>, ) -> Poll<Option<Result<RecordBatch>>>

fn update_poll_count_on_the_same_value(&mut self, partition_idx: usize)

For the given partition, updates the poll count. If the current value is the same of the previous value, it increases the count by 1; otherwise, it is reset as 0.

fn fetch_reached(&mut self) -> bool

fn advance_cursors(&mut self, stream_idx: usize) -> bool

Advances the actual cursor. If it reaches its end, update the previous cursor with it.

If the given partition is not exhausted, the function returns true.

fn is_gt(&self, a: usize, b: usize) -> bool

Returns true if the cursor at index a is greater than at index b. In an equality case, it compares the partition indices given.

fn is_poll_count_gt(&self, a: usize, b: usize) -> bool

fn update_winner( &mut self, cmp_node: usize, winner: &mut usize, challenger: usize, )

fn lt_leaf_node_index(&self, cursor_index: usize) -> usize

Find the leaf node index in the loser tree for the given cursor index

Note that this is not necessarily a leaf node in the tree, but it can also be a half-node (a node with only one child). This happens when the number of cursors/streams is not a power of two. Thus, the loser tree will be unbalanced, but it will still work correctly.

For example, with 5 streams, the loser tree will look like this:

          0 (winner)

          1
       /     \
      2       3
    /  \     / \
   4    |   |   |
  / \   |   |   |
-+---+--+---+---+---- Below is not a part of loser tree
 S3 S4 S0   S1  S2

S0, S1, … S4 are the streams (read: stream at index 0, stream at index 1, etc.)

Zooming in at node 2 in the loser tree as an example, we can see that it takes as input the next item at (S0) and the loser of (S3, S4).

fn lt_parent_node_index(&self, node_idx: usize) -> usize

Find the parent node index for the given node index

fn init_loser_tree(&mut self)

Attempts to initialize the loser tree with one value from each non exhausted input, if possible

fn reset_poll_counts(&mut self)

Resets the poll count by incrementing the reset epoch.

fn handle_tie(&mut self, cmp_node: usize, winner: &mut usize, challenger: usize)

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.

fn update_loser_tree(&mut self)

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.

Trait Implementations

impl<C: Debug + CursorValues> Debug for SortPreservingMergeStream<C>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<C: CursorValues + Unpin> RecordBatchStream for SortPreservingMergeStream<C>

fn schema(&self) -> SchemaRef

Returns the schema of this RecordBatchStream.

impl<C: CursorValues + Unpin> Stream for SortPreservingMergeStream<C>

type Item = Result<RecordBatch, DataFusionError>

Values yielded by the stream.

fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Option<Self::Item>>

Attempt to pull out the next value of this stream, registering the current task for wakeup if the value is not yet available, and returning None if the stream is exhausted.

fn size_hint(&self) -> (usize, Option<usize>)

Returns the bounds on the remaining length of the stream.