Skip to content

Increased memory usage due to round robin tie-breaker feature in SortPreservingMergeStream #23604

Description

@ariel-miculas

Describe the bug

prev_cursors in Sort Preserving Merge keeps an additional Cursor alive during the merge operation;
This is particularly problematic in the context of a single-column sort when merging streams read from spill files in IPC format, because in this case, due to apache/arrow-rs#6363, the single sort column kept alive by prev_cursors prevents an entire Record batch from being dropped. This is particularly problematic with large Record batches, leading to an excess in the observed peak memory.

Practical experiment:

with round robin tie-breaking feature disabled and this patch:

  /// 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 advance_cursors(&mut self, stream_idx: usize) -> bool {
        if let Some(cursor) = &mut self.cursors[stream_idx] {
            let _ = cursor.advance();
            if cursor.is_finished() {
                if self.enable_round_robin_tie_breaker {
                    self.prev_cursors[stream_idx] = self.cursors[stream_idx].take();
                } else {
                    self.cursors[stream_idx] = None;
                }
            }
            true
        } else {
            false
        }
    }

Observed peak memory (measured directly from the Rust allocator, not by jemalloc or other allocator stats) goes down from 425 MiBs to 347 MiBs, which is two Record batches worth of memory (each record batch is 39 MiBs).
In this experiment I'm always merging two streams read from the spill files via multi level merge.

To Reproduce

No response

Expected behavior

prev_cursors only needs the last row from the cursor, so why does it have to keep the entire cursor?

Additional context

Another similar issue, where extra Record Batches lead to increase in peak memory: #23559

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

Fields

No fields configured for Bug.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions