datafusion-physical-plan 55.0.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Create a stream that do a multi level merge stream

use crate::metrics::BaselineMetrics;
use crate::{EmptyRecordBatchStream, SpillManager};
use arrow::array::RecordBatch;
use std::fmt::{Debug, Formatter};
use std::mem;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use arrow::datatypes::SchemaRef;
use datafusion_common::{Result, internal_err, resources_err};
use datafusion_execution::memory_pool::MemoryReservation;

use crate::sorts::builder::try_grow_reservation_to_at_least;
use crate::sorts::sort::get_reserved_bytes_for_record_batch_size;
use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
use crate::stream::{ObservedStream, RecordBatchStreamAdapter};
use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream};
use datafusion_physical_expr_common::sort_expr::LexOrdering;
use futures::TryStreamExt;
use futures::{Stream, StreamExt};

/// Merges a stream of sorted cursors and record batches into a single sorted stream
///
/// This is a wrapper around [`SortPreservingMergeStream`](crate::sorts::merge::SortPreservingMergeStream)
/// that provide it the sorted streams/files to merge while making sure we can merge them in memory.
/// In case we can't merge all of them in a single pass we will spill the intermediate results to disk
/// and repeat the process.
///
/// ## High level Algorithm
/// 1. Get the maximum amount of sorted in-memory streams and spill files we can merge with the available memory
/// 2. Sort them to a sorted stream
/// 3. Do we have more spill files to merge?
///  - Yes: write that sorted stream to a spill file,
///    add that spill file back to the spill files to merge and
///    repeat the process
///
///  - No: return that sorted stream as the final output stream
///
/// ```text
/// Initial State: Multiple sorted streams + spill files
///      ┌───────────┐
///      │  Phase 1  │
///      └───────────┘
/// ┌──Can hold in memory─┐
/// │   ┌──────────────┐  │
/// │   │  In-memory   │
/// │   │sorted stream │──┼────────┐
/// │   │      1       │  │        │
///     └──────────────┘  │        │
/// │   ┌──────────────┐  │        │
/// │   │  In-memory   │           │
/// │   │sorted stream │──┼────────┤
/// │   │      2       │  │        │
///     └──────────────┘  │        │
/// │   ┌──────────────┐  │        │
/// │   │  In-memory   │           │
/// │   │sorted stream │──┼────────┤
/// │   │      3       │  │        │
///     └──────────────┘  │        │
/// │   ┌──────────────┐  │        │            ┌───────────┐
/// │   │ Sorted Spill │           │            │  Phase 2  │
/// │   │    file 1    │──┼────────┤            └───────────┘
/// │   └──────────────┘  │        │
///  ──── ──── ──── ──── ─┘        │       ┌──Can hold in memory─┐
///                                │       │                     │
///     ┌──────────────┐           │       │   ┌──────────────┐
///     │ Sorted Spill │           │       │   │ Sorted Spill │  │
///     │    file 2    │──────────────────────▶│    file 2    │──┼─────┐
///     └──────────────┘           │           └──────────────┘  │     │
///     ┌──────────────┐           │       │   ┌──────────────┐  │     │
///     │ Sorted Spill │           │       │   │ Sorted Spill │        │
///     │    file 3    │──────────────────────▶│    file 3    │──┼─────┤
///     └──────────────┘           │       │   └──────────────┘  │     │
///     ┌──────────────┐           │           ┌──────────────┐  │     │
///     │ Sorted Spill │           │       │   │ Sorted Spill │  │     │
///     │    file 4    │──────────────────────▶│    file 4    │────────┤          ┌───────────┐
///     └──────────────┘           │       │   └──────────────┘  │     │          │  Phase 3  │
///                                │       │                     │     │          └───────────┘
///                                │        ──── ──── ──── ──── ─┘     │     ┌──Can hold in memory─┐
///                                │                                   │     │                     │
///     ┌──────────────┐           │           ┌──────────────┐        │     │  ┌──────────────┐
///     │ Sorted Spill │           │           │ Sorted Spill │        │     │  │ Sorted Spill │   │
///     │    file 5    │──────────────────────▶│    file 5    │────────────────▶│    file 5    │───┼───┐
///     └──────────────┘           │           └──────────────┘        │     │  └──────────────┘   │   │
///                                │                                   │     │                     │   │
///                                │           ┌──────────────┐        │     │  ┌──────────────┐       │
///                                │           │ Sorted Spill │        │     │  │ Sorted Spill │   │   │       ┌── ─── ─── ─── ─── ─── ─── ──┐
///                                └──────────▶│    file 6    │────────────────▶│    file 6    │───┼───┼──────▶         Output Stream
///                                            └──────────────┘        │     │  └──────────────┘   │   │       └── ─── ─── ─── ─── ─── ─── ──┘
///                                                                    │     │                     │   │
///                                                                    │     │  ┌──────────────┐       │
///                                                                    │     │  │ Sorted Spill │   │   │
///                                                                    └───────▶│    file 7    │───┼───┘
///                                                                          │  └──────────────┘   │
///                                                                          │                     │
///                                                                          └─ ──── ──── ──── ────
/// ```
///
/// ## Memory Management Strategy
///
/// This multi-level merge make sure that we can handle any amount of data to sort as long as
/// we have enough memory to merge at least 2 streams at a time, even when individual record
/// batches are skewed (very wide).
///
/// 1. **Worst-Case Memory Reservation**: Reserves memory based on the largest
///    batch size encountered in each spill file to merge, ensuring sufficient memory is always
///    available during merge operations.
/// 2. **Adaptive Buffer Sizing**: Reduces buffer sizes when memory is constrained
/// 3. **Spill-to-Disk**: Spill to disk when we cannot merge all files in memory
/// 4. **Re-spilling Skewed Runs**: If even at the smallest read-buffer size we still cannot
///    reserve memory for the minimum of 2 streams - because a single run's largest batch is so
///    wide that two streams' worth of reservation exceeds the budget - the larger of the two
///    runs is re-spilled with each batch sliced in half. This shrinks its largest batch,
///    lowering the per-stream reservation, and the merge pass is retried. The re-spilled run
///    is tracked alongside a per-run batch-size limit equal to half the batch size it was
///    written with, so any later merge that includes it caps its output batch size to match -
///    otherwise the merged run could rebuild a full-size batch and reintroduce the skew.
///    Crucially the global merge batch size is *not* lowered, so re-spilling more than one run
///    does not compound the reduction. If a batch cannot be split any further (a single row
///    wider than the budget), the merge surfaces `ResourcesExhausted` instead of looping
///    forever.
pub(crate) struct MultiLevelMergeBuilder {
    spill_manager: SpillManager,
    schema: SchemaRef,
    /// Sorted runs still to be merged. Each run is paired with the batch-size limit a
    /// merge consuming it must cap its output at. Runs written at the full batch size
    /// carry `batch_size`. A run re-spilled smaller to resolve skew carries its halved
    /// limit (see [`Self::split_spill_file_in_half`]). Tracking it here keeps this limit
    /// out of the public [`SortedSpillFile`], so no external caller has to set it.
    sorted_spill_files: Vec<(SortedSpillFile, usize)>,
    sorted_streams: Vec<SendableRecordBatchStream>,
    expr: LexOrdering,
    metrics: BaselineMetrics,
    batch_size: usize,
    reservation: MemoryReservation,
    fetch: Option<usize>,
    enable_round_robin_tie_breaker: bool,
}

impl Debug for MultiLevelMergeBuilder {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "MultiLevelMergeBuilder")
    }
}

impl MultiLevelMergeBuilder {
    #[expect(clippy::too_many_arguments)]
    pub(crate) fn new(
        spill_manager: SpillManager,
        schema: SchemaRef,
        sorted_spill_files: Vec<SortedSpillFile>,
        sorted_streams: Vec<SendableRecordBatchStream>,
        expr: LexOrdering,
        metrics: BaselineMetrics,
        batch_size: usize,
        reservation: MemoryReservation,
        fetch: Option<usize>,
        enable_round_robin_tie_breaker: bool,
    ) -> Self {
        Self {
            spill_manager,
            schema,
            // Initial runs are written at the full batch size, so they impose no cap
            // on later merges - record `batch_size` as their (unconstrained) limit.
            sorted_spill_files: sorted_spill_files
                .into_iter()
                .map(|file| (file, batch_size))
                .collect(),
            sorted_streams,
            expr,
            metrics,
            batch_size,
            reservation,
            enable_round_robin_tie_breaker,
            fetch,
        }
    }

    pub(crate) fn create_spillable_merge_stream(self) -> SendableRecordBatchStream {
        Box::pin(RecordBatchStreamAdapter::new(
            Arc::clone(&self.schema),
            futures::stream::once(self.create_stream()).try_flatten(),
        ))
    }

    async fn create_stream(mut self) -> Result<SendableRecordBatchStream> {
        loop {
            let (mut stream, batch_size_limit) =
                match self.merge_sorted_runs_within_mem_limit()? {
                    MergeStep::Stream {
                        stream,
                        batch_size_limit,
                    } => (stream, batch_size_limit),
                    MergeStep::SplitThenRetry(index) => {
                        // Couldn't reserve memory for the minimum of 2 streams. Re-spill
                        // the larger of the two we're trying to merge with half its batch
                        // size so its largest batch shrinks, lowering the per-stream
                        // reservation, then retry. Makes the merge resilient to skewed
                        // (very wide) rows.
                        self.split_spill_file_in_half(index).await?;
                        continue;
                    }
                };

            // TODO - add a threshold for number of files to disk even if empty and reading from disk so
            //        we can avoid the memory reservation

            // If no spill files are left, we can return the stream as this is the last sorted run
            // TODO - We can write to disk before reading it back to avoid having multiple streams in memory
            if self.sorted_spill_files.is_empty() {
                assert!(
                    self.sorted_streams.is_empty(),
                    "We should not have any sorted streams left"
                );

                return Ok(stream);
            }

            // Need to sort to a spill file
            let Some((spill_file, max_record_batch_memory)) = self
                .spill_manager
                .spill_record_batch_stream_and_return_max_batch_memory(
                    &mut stream,
                    "MultiLevelMergeBuilder intermediate spill",
                )
                .await?
            else {
                continue;
            };

            // Add the spill file paired with the batch-size limit of the merge that
            // produced it: if that merge consumed a shrunk (skew-resolved) run, its
            // output was capped and this intermediate run is likewise capped, so a
            // later pass that re-merges it won't rebuild an oversized batch.
            self.sorted_spill_files.push((
                SortedSpillFile {
                    file: spill_file,
                    max_record_batch_memory,
                },
                batch_size_limit,
            ));
        }
    }

    /// This tries to create a stream that merges the most sorted streams and sorted spill files
    /// as possible within the memory limit.
    fn merge_sorted_runs_within_mem_limit(&mut self) -> Result<MergeStep> {
        match (self.sorted_spill_files.len(), self.sorted_streams.len()) {
            // No data so empty batch
            (0, 0) => {
                let empty_stream =
                    Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema)));
                Ok(MergeStep::Stream {
                    stream: self.observe_output(empty_stream),
                    batch_size_limit: self.batch_size,
                })
            }

            // Only in-memory stream, return that
            (0, 1) => {
                let output_stream = self.sorted_streams.remove(0);
                Ok(MergeStep::Stream {
                    stream: self.observe_output(output_stream),
                    batch_size_limit: self.batch_size,
                })
            }

            // Only single sorted spill file so return it
            (1, 0) => {
                let (spill_file, batch_size) = self.sorted_spill_files.remove(0);

                // Not reserving any memory for this disk as we are not holding it in memory
                let output_stream = self
                    .spill_manager
                    .read_spill_as_stream(spill_file.file, None)?;

                Ok(MergeStep::Stream {
                    stream: self.observe_output(output_stream),
                    batch_size_limit: batch_size,
                })
            }

            // Only in memory streams, so merge them all in a single pass. In-memory
            // runs are never shrunk for skew, so this merge runs at the full batch
            // size and its output carries no limit.
            (0, _) => {
                let sorted_stream = mem::take(&mut self.sorted_streams);
                // No need to wrap with observed stream since merge sort will update the observed metrics
                Ok(MergeStep::Stream {
                    stream: self.create_new_merge_sort(
                        sorted_stream,
                        // If we have no sorted spill files left, this is the last run
                        true,
                        true,
                        self.batch_size,
                    )?,
                    batch_size_limit: self.batch_size,
                })
            }

            // Need to merge multiple streams
            (_, _) => {
                // Transfer any pre-reserved bytes (from sort_spill_reservation_bytes)
                // to the merge memory reservation. This prevents starvation when
                // concurrent sort partitions compete for pool memory: the pre-reserved
                // bytes cover spill file buffer reservations without additional pool
                // allocation.
                let mut memory_reservation = self.reservation.take();

                // Compute the minimum before taking the in-memory streams so that, if we
                // need to re-spill and retry, `self.sorted_streams` is left untouched.
                let minimum_number_of_required_streams =
                    2_usize.saturating_sub(self.sorted_streams.len());

                let (sorted_spill_files, buffer_size) = match self
                    .get_sorted_spill_files_to_merge(
                        2,
                        // we must have at least 2 streams to merge
                        minimum_number_of_required_streams,
                        &mut memory_reservation,
                    )? {
                    SpillFilesToMerge::Ready(sorted_spill_files, buffer_size) => {
                        (sorted_spill_files, buffer_size)
                    }
                    // Not enough memory to seat 2 streams. Re-spill the blocking file
                    // smaller and retry. `get_sorted_spill_files_to_merge` already freed
                    // the reservation and `self.sorted_streams` is untouched, so the
                    // retry starts clean.
                    SpillFilesToMerge::SplitThenRetry(index) => {
                        return Ok(MergeStep::SplitThenRetry(index));
                    }
                };

                // Don't account for existing streams memory
                // as we are not holding the memory for them
                let mut sorted_streams = mem::take(&mut self.sorted_streams);

                let is_only_merging_memory_streams = sorted_spill_files.is_empty();

                // If no spill files were selected (e.g. all too large for
                // available memory but enough in-memory streams exist),
                // return the pre-reserved bytes to self.reservation so
                // create_new_merge_sort can transfer them to the merge
                // stream's BatchBuilder.
                if is_only_merging_memory_streams {
                    mem::swap(&mut self.reservation, &mut memory_reservation);
                }

                // Cap the merge output at the smallest limit among the runs we're
                // about to merge. Runs that were shrunk for skew carry a smaller limit,
                // if none do, every run carries `self.batch_size` and the merge runs at
                // the full batch size. The output stream is tagged with the same limit
                // (see the `MergeStep::Stream` returns below) so a re-spilled
                // intermediate run stays shrunk and won't rebuild an oversized batch on
                // a later pass.
                let mut output_batch_size = self.batch_size;
                for (spill, batch_size_limit) in sorted_spill_files {
                    let stream = self
                        .spill_manager
                        .clone()
                        .with_batch_read_buffer_capacity(buffer_size)
                        .read_spill_as_stream(
                            spill.file,
                            Some(spill.max_record_batch_memory),
                        )?;
                    output_batch_size = output_batch_size.min(batch_size_limit);
                    sorted_streams.push(stream);
                }
                let merge_sort_stream = self.create_new_merge_sort(
                    sorted_streams,
                    // If we have no sorted spill files left, this is the last run
                    self.sorted_spill_files.is_empty(),
                    is_only_merging_memory_streams,
                    output_batch_size,
                )?;

                // If we're only merging memory streams, we don't need to attach the memory reservation
                // as it's empty
                if is_only_merging_memory_streams {
                    assert_eq!(
                        memory_reservation.size(),
                        0,
                        "when only merging memory streams, we should not have any memory reservation and let the merge sort handle the memory"
                    );

                    Ok(MergeStep::Stream {
                        stream: merge_sort_stream,
                        batch_size_limit: output_batch_size,
                    })
                } else {
                    // Attach the memory reservation to the stream to make sure we have enough memory
                    // throughout the merge process as we bypassed the memory pool for the merge sort stream
                    Ok(MergeStep::Stream {
                        stream: Box::pin(StreamAttachedReservation::new(
                            merge_sort_stream,
                            memory_reservation,
                        )),
                        batch_size_limit: output_batch_size,
                    })
                }
            }
        }
    }

    fn create_new_merge_sort(
        &mut self,
        streams: Vec<SendableRecordBatchStream>,
        is_output: bool,
        all_in_memory: bool,
        output_batch_size: usize,
    ) -> Result<SendableRecordBatchStream> {
        let mut builder = StreamingMergeBuilder::new()
            .with_schema(Arc::clone(&self.schema))
            .with_expressions(&self.expr)
            .with_batch_size(output_batch_size)
            .with_fetch(self.fetch)
            .with_metrics(if is_output {
                // Only add the metrics to the last run
                self.metrics.clone()
            } else {
                self.metrics.intermediate()
            })
            .with_round_robin_tie_breaker(self.enable_round_robin_tie_breaker)
            .with_streams(streams);

        if !all_in_memory {
            // Don't track memory used by this stream as we reserve that memory by worst case sceneries
            // (reserving memory for the biggest batch in each stream)
            // TODO - avoid this hack as this can be broken easily when `SortPreservingMergeStream`
            //        changes the implementation to use more/less memory
            builder = builder.with_bypass_mempool();
        } else {
            // If we are only merging in-memory streams, we need to use the memory reservation
            // because we don't know the maximum size of the batches in the streams.
            // Use take() to transfer any pre-reserved bytes so the merge can use them
            // as its initial budget without additional pool allocation.
            builder = builder.with_reservation(self.reservation.take());
        }

        builder.build()
    }

    /// Return the sorted spill files to use for the next phase, and the buffer size
    /// This will try to get as many spill files as possible to merge, and if we don't have enough streams
    /// it will try to reduce the buffer size until we have enough streams to merge
    /// otherwise it will return an error
    fn get_sorted_spill_files_to_merge(
        &mut self,
        buffer_len: usize,
        minimum_number_of_required_streams: usize,
        reservation: &mut MemoryReservation,
    ) -> Result<SpillFilesToMerge> {
        assert_ne!(buffer_len, 0, "Buffer length must be greater than 0");
        let mut number_of_spills_to_read_for_current_phase = 0;
        let configured_fan_in = self
            .spill_manager
            .env()
            .disk_manager
            .max_spill_merge_fan_in();
        let max_spill_files = effective_spill_merge_fan_in(configured_fan_in);
        // Track total memory needed for spill file buffers. When the
        // reservation has pre-reserved bytes (from sort_spill_reservation_bytes),
        // those bytes cover the first N spill files without additional pool
        // allocation, preventing starvation under memory pressure.
        let mut total_needed: usize = 0;

        for (spill, _) in &self.sorted_spill_files {
            if number_of_spills_to_read_for_current_phase >= max_spill_files {
                break;
            }

            let per_spill = get_reserved_bytes_for_record_batch_size(
                spill.max_record_batch_memory,
                // Size will be the same as the sliced size, bc it is a spilled batch.
                spill.max_record_batch_memory,
            ) * buffer_len;
            total_needed += per_spill;

            // For memory pools that are not shared this is good, for other
            // this is not and there should be some upper limit to memory
            // reservation so we won't starve the system.
            match try_grow_reservation_to_at_least(reservation, total_needed) {
                Ok(_) => {
                    number_of_spills_to_read_for_current_phase += 1;
                }
                // If we can't grow the reservation, we need to stop
                Err(err) => {
                    // We must have at least 2 streams to merge, so if we don't have enough memory
                    // fail
                    if minimum_number_of_required_streams
                        > number_of_spills_to_read_for_current_phase
                    {
                        // Free the memory we reserved for this merge as we either try again or fail
                        reservation.free();
                        if buffer_len > 1 {
                            // Try again with smaller buffer size, it will be slower but at least we can merge
                            return self.get_sorted_spill_files_to_merge(
                                buffer_len - 1,
                                minimum_number_of_required_streams,
                                reservation,
                            );
                        }

                        // buffer_len == 1 and we still can't seat the minimum of 2 streams.
                        if number_of_spills_to_read_for_current_phase == 0 {
                            // We couldn't even reserve a single stream - one record batch
                            // is larger than the whole merge budget. That's the lone-batch
                            // case, not the 2-stream merge skew we rescue here - surface it.
                            return Err(err);
                        }

                        // We seated one stream (index 0) but not the second (index 1, the
                        // batch that just failed to reserve). Those are by definition the
                        // only two streams we are trying to merge, so re-spill the larger
                        // of them with a smaller batch size and retry, the smaller max
                        // batch lowers the per-stream reservation enough to seat both.
                        let split_index = usize::from(
                            self.sorted_spill_files[1].0.max_record_batch_memory
                                > self.sorted_spill_files[0].0.max_record_batch_memory,
                        );
                        return Ok(SpillFilesToMerge::SplitThenRetry(split_index));
                    }

                    // We reached the maximum amount of memory we can use
                    // for this merge
                    break;
                }
            }
        }

        let spills = self
            .sorted_spill_files
            .drain(..number_of_spills_to_read_for_current_phase)
            .collect::<Vec<_>>();

        Ok(SpillFilesToMerge::Ready(spills, buffer_len))
    }

    /// Re-spill the spill file at `index` with half its batch size, putting it back
    /// at the same position. We read the file back and re-spill it through the normal
    /// spill API (which owns batch layout), slicing every batch in two, which halves
    /// the largest written batch and so lowers the per-stream merge reservation enough
    /// for the next attempt to seat both streams. One stream's worth of memory is
    /// reserved for the duration and freed afterwards. Makes the merge resilient to skew.
    ///
    /// Instead of halving the *global* merge batch size (which would compound when more
    /// than one run is re-spilled), the shrunk run records its own smaller batch-size
    /// limit (tracked alongside the run in `sorted_spill_files`), so only merges that
    /// actually consume it pay the reduced batch size.
    async fn split_spill_file_in_half(&mut self, index: usize) -> Result<()> {
        log::debug!(
            "2 spilled streams could not be loaded into memory for merge \
        (requires 2x of the largest batch from both), re-spilling the larger of the two with half \
        the batch size to reduce memory needs for the next merge attempt. the shrunk run carries \
        a halved batch-size limit so only merges consuming it use the smaller batch size"
        );

        // Extract the target in O(1) instead of `remove(index)`, which would shift
        // every following spill file. Swap it to the back and pop it; the matching
        // swap after re-spilling restores the original order, so the vec ends up
        // exactly as it started, just with the target file shrunk.
        // `old_batch_size` is the batch size this run was written with (the full merge
        // batch size unless it was already shrunk once). Halving it caps the next merge
        // that reads this run so the merged output can't rebuild a full-size batch.
        let last = self.sorted_spill_files.len() - 1;
        self.sorted_spill_files.swap(index, last);
        let (target, old_batch_size) = self
            .sorted_spill_files
            .pop()
            .expect("index is in bounds, so the vec is non-empty");
        let old_max = target.max_record_batch_memory;

        // Reserve enough to hold a single stream of this file while we re-spill it.
        let reservation = self.reservation.new_empty();
        reservation
            .try_grow(get_reserved_bytes_for_record_batch_size(old_max, old_max))?;

        let source = self
            .spill_manager
            .read_spill_as_stream(target.file, Some(old_max))?;
        // Re-spill with half the batch size: slice every batch in two. The spill
        // writer owns the batch layout, we only change how many rows per batch.
        let mut halved: SendableRecordBatchStream =
            Box::pin(RecordBatchStreamAdapter::new(
                Arc::clone(&self.schema),
                source.flat_map(|batch| {
                    futures::stream::iter(match batch {
                        Ok(batch) => split_batch_in_half(batch)
                            .into_iter()
                            .map(Ok)
                            .collect::<Vec<_>>(),
                        Err(e) => vec![Err(e)],
                    })
                }),
            ));

        let result = self
            .spill_manager
            .spill_record_batch_stream_and_return_max_batch_memory(
                &mut halved,
                "MultiLevelMergeBuilder split skewed spill",
            )
            .await?;

        reservation.free();

        let Some((file, new_max)) = result else {
            return internal_err!("re-spilling a skewed spill file produced no data");
        };

        // If halving could not reduce the largest batch (e.g. a single row that is
        // itself wider than the budget), there is nothing more we can do - surface
        // the out-of-memory condition instead of looping forever.
        if new_max >= old_max {
            return resources_err!(
                "Cannot merge sorted runs: a single record batch of {old_max} bytes \
                 exceeds the available merge memory and cannot be split further"
            );
        }

        // Record the halved batch size as a *per-run* limit rather than lowering the
        // global batch size. Merges that don't touch this run keep the full batch
        // size. a merge that reads it caps its output at this limit so the merged run
        // can't rebuild a full-size batch and reintroduce the skew.
        let new_batch_size_limit = (old_batch_size / 2).max(1);

        // Push the re-spilled (smaller) file and swap it back into `index`, undoing
        // the swap-to-back above so the order is preserved.
        self.sorted_spill_files.push((
            SortedSpillFile {
                file,
                max_record_batch_memory: new_max,
            },
            new_batch_size_limit,
        ));
        let last = self.sorted_spill_files.len() - 1;
        self.sorted_spill_files.swap(index, last);

        Ok(())
    }

    fn observe_output(
        &self,
        stream: SendableRecordBatchStream,
    ) -> SendableRecordBatchStream {
        Box::pin(ObservedStream::new(stream, self.metrics.clone(), None))
    }
}

/// Outcome of trying to reserve memory for one multi-level merge pass.
enum SpillFilesToMerge {
    /// Enough memory: the spill files to read this pass (each paired with its
    /// batch-size limit) and the read-ahead buffer size.
    Ready(Vec<(SortedSpillFile, usize)>, usize),
    /// Could not seat the minimum of 2 streams. Re-spill the spill file at this index
    /// with a smaller (halved) batch size, then retry the pass.
    SplitThenRetry(usize),
}

/// What one iteration of the multi-level merge loop should do next.
enum MergeStep {
    /// A merged stream is ready to be consumed (and possibly spilled back).
    Stream {
        stream: SendableRecordBatchStream,
        /// The batch-size limit to stamp on the run if this stream is re-spilled as an
        /// intermediate result: the batch size its merge ran at. It equals the full
        /// merge batch size unless the merge consumed a skew-resolved run, in which
        /// case it is that run's smaller limit so the re-spilled result stays capped
        /// and can't rebuild an oversized batch.
        batch_size_limit: usize,
    },
    /// Re-spill the spill file at this index smaller, then retry the merge step.
    SplitThenRetry(usize),
}

/// Slice `batch` into two row-halves so a re-spill writes batches half the size.
fn split_batch_in_half(batch: RecordBatch) -> Vec<RecordBatch> {
    let num_rows = batch.num_rows();
    if num_rows <= 1 {
        return vec![batch];
    }
    let mid = num_rows / 2;
    vec![batch.slice(0, mid), batch.slice(mid, num_rows - mid)]
}

fn effective_spill_merge_fan_in(configured_fan_in: usize) -> usize {
    if configured_fan_in == 0 {
        usize::MAX
    } else {
        configured_fan_in.max(2)
    }
}

struct StreamAttachedReservation {
    stream: SendableRecordBatchStream,
    reservation: MemoryReservation,
}

impl StreamAttachedReservation {
    fn new(stream: SendableRecordBatchStream, reservation: MemoryReservation) -> Self {
        Self {
            stream,
            reservation,
        }
    }
}

impl Stream for StreamAttachedReservation {
    type Item = Result<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let res = self.stream.poll_next_unpin(cx);

        match res {
            Poll::Ready(res) => {
                match res {
                    Some(Ok(batch)) => Poll::Ready(Some(Ok(batch))),
                    Some(Err(err)) => {
                        // Had an error so drop the data
                        self.reservation.free();
                        Poll::Ready(Some(Err(err)))
                    }
                    None => {
                        // Stream is done so free the memory
                        self.reservation.free();

                        Poll::Ready(None)
                    }
                }
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl RecordBatchStream for StreamAttachedReservation {
    fn schema(&self) -> SchemaRef {
        self.stream.schema()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::expressions::PhysicalSortExpr;
    use arrow::array::{AsArray, Int64Array};
    use arrow::compute::concat_batches;
    use arrow::datatypes::{DataType, Field, Int64Type, Schema};
    use datafusion_execution::memory_pool::{
        GreedyMemoryPool, MemoryConsumer, MemoryPool,
    };
    use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
    use datafusion_physical_expr::expressions::{Column, col};
    use datafusion_physical_expr_common::metrics::{
        ExecutionPlanMetricsSet, SpillMetrics,
    };

    fn test_schema() -> SchemaRef {
        Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]))
    }

    fn build_spill_manager(env: &Arc<RuntimeEnv>, schema: &SchemaRef) -> SpillManager {
        SpillManager::new(
            Arc::clone(env),
            SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            Arc::clone(schema),
        )
    }

    /// Spill `values` (which must already be sorted) as a single sorted run and
    /// return it as a `SortedSpillFile` carrying its recorded largest-batch memory.
    fn make_sorted_spill_file(
        spill_manager: &SpillManager,
        schema: &SchemaRef,
        values: Vec<i64>,
    ) -> SortedSpillFile {
        let batch = RecordBatch::try_new(
            Arc::clone(schema),
            vec![Arc::new(Int64Array::from(values))],
        )
        .unwrap();
        let batches: Vec<Result<RecordBatch>> = vec![Ok(batch)];
        let (file, max_record_batch_memory) = spill_manager
            .spill_record_batch_iter_and_return_max_batch_memory(
                batches.into_iter(),
                "test input run",
            )
            .unwrap()
            .expect("spill should produce a file");
        SortedSpillFile {
            file,
            max_record_batch_memory,
        }
    }

    fn build_merge_builder(
        spill_manager: SpillManager,
        schema: SchemaRef,
        sorted_spill_files: Vec<SortedSpillFile>,
        pool: &Arc<dyn MemoryPool>,
        batch_size: usize,
    ) -> MultiLevelMergeBuilder {
        let reservation = MemoryConsumer::new("test merge").register(pool);
        let expr: LexOrdering =
            [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into();
        MultiLevelMergeBuilder::new(
            spill_manager,
            schema,
            sorted_spill_files,
            vec![],
            expr,
            BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            batch_size,
            reservation,
            None,
            false,
        )
    }

    /// Two sorted runs whose largest batches are too big to both
    /// be seated in the merge budget at once are re-spilled (halved) until they
    /// fit, and the merge then completes with fully sorted, complete output.
    #[tokio::test]
    async fn skewed_runs_are_respilled_so_the_merge_fits() -> Result<()> {
        let env = Arc::new(RuntimeEnv::default());
        let schema = test_schema();
        let spill_manager = build_spill_manager(&env, &schema);

        let n: i64 = 16384;
        let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect());
        let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect());
        let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory);

        // Seating two streams needs ~4*m (2*m each), which does NOT fit, but the
        // budget is large enough once a run is halved. The rescue keeps halving
        // the blocking run until two streams fit (here, after one halving).
        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 7 / 2));

        let builder = build_merge_builder(
            spill_manager,
            Arc::clone(&schema),
            vec![f0, f1],
            &pool,
            8192,
        );
        let stream = builder.create_spillable_merge_stream();
        let batches: Vec<RecordBatch> = stream.try_collect().await?;

        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(
            total_rows,
            (2 * n) as usize,
            "the merge must emit every input row"
        );

        let merged = concat_batches(&schema, &batches)?;
        let col = merged.column(0).as_primitive::<Int64Type>();
        for i in 1..col.len() {
            assert!(
                col.value(i - 1) <= col.value(i),
                "merge output must be sorted: {} > {} at {i}",
                col.value(i - 1),
                col.value(i),
            );
        }

        Ok(())
    }

    /// Tests the `new_max >= old_max` guard: a single-row run cannot be split
    /// any smaller, so re-spilling it does not shrink the largest batch and the
    /// rescue surfaces `ResourcesExhausted` rather than looping forever.
    #[tokio::test]
    async fn respilling_an_unsplittable_run_surfaces_resources_exhausted() -> Result<()> {
        let env = Arc::new(RuntimeEnv::default());
        let schema = test_schema();
        let spill_manager = build_spill_manager(&env, &schema);

        // A one-row run: `split_batch_in_half` returns it unchanged, so the
        // re-spilled file's largest batch cannot drop below the original.
        let f0 = make_sorted_spill_file(&spill_manager, &schema, vec![42]);

        // Ample budget so the only possible failure is the un-splittable guard,
        // not the single-stream reservation itself.
        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024 * 1024));
        let mut builder =
            build_merge_builder(spill_manager, schema, vec![f0], &pool, 1024);

        let err = builder
            .split_spill_file_in_half(0)
            .await
            .expect_err("re-spilling a one-row run cannot shrink it");
        assert!(
            err.to_string().contains("cannot be split further"),
            "expected the un-splittable guard error, got: {err}"
        );

        Ok(())
    }

    /// Proves the re-spill also halves the merge output batch size: after one
    /// re-spill the merged run is emitted in 4096-row batches (not the original
    /// 8192), so it cannot rebuild a full-size batch and reintroduce the skew.
    #[tokio::test]
    async fn respill_halves_the_merge_output_batch_size() -> Result<()> {
        let env = Arc::new(RuntimeEnv::default());
        let schema = test_schema();
        let spill_manager = build_spill_manager(&env, &schema);

        let n: i64 = 16384;
        let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect());
        let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect());
        let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory);

        // 3.5*m forces exactly one re-spill (split one run, then both fit), which
        // halves the merge output batch size.
        let initial_batch_size = 8192;
        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 7 / 2));

        let builder = build_merge_builder(
            spill_manager,
            Arc::clone(&schema),
            vec![f0, f1],
            &pool,
            initial_batch_size,
        );
        let stream = builder.create_spillable_merge_stream();
        let batches: Vec<RecordBatch> = stream.try_collect().await?;

        // All rows are still present.
        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total_rows, (2 * n) as usize);

        // The largest emitted batch is the halved size, not the original 8192: the
        // shrunk run carries a halved batch-size limit, and the final pass consumes
        // it, so the merge output is capped there. Without the per-run limit the merge
        // would rebuild 8192-row batches.
        let expected_batch_size = initial_batch_size / 2;
        let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0);
        assert_eq!(
            max_batch_rows, expected_batch_size,
            "after one re-spill the merge must emit {expected_batch_size}-row \
             batches, got a largest batch of {max_batch_rows} rows"
        );

        Ok(())
    }

    /// Same as [`respill_halves_the_merge_output_batch_size`], but under a budget tight
    /// enough that *both* runs must be re-spilled before the merge fits - the scenario
    /// where the batch-size reduction could compound. Because the reduction is tracked
    /// per-run (each run capped at half) rather than by halving the global batch size on
    /// every split, the merged output is emitted in 4096-row batches - half, not a
    /// quarter. A global-halving implementation would have halved once per re-spill and
    /// emitted 2048-row batches.
    #[tokio::test]
    async fn respilling_two_skewed_runs_halves_the_output_without_compounding()
    -> Result<()> {
        let env = Arc::new(RuntimeEnv::default());
        let schema = test_schema();
        let spill_manager = build_spill_manager(&env, &schema);

        let n: i64 = 16384;
        let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect());
        let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect());
        let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory);

        // 2.5*m is tight enough that even after halving one run the two still don't
        // fit, so *both* runs are re-spilled once before the merge succeeds. (3.5*m,
        // as in the single-split test, would let the pair fit after one split.) This
        // is exactly the scenario where a compounding, global-halving implementation
        // would drive the output batch size down to a quarter.
        let initial_batch_size = 8192;
        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 5 / 2));

        let builder = build_merge_builder(
            spill_manager,
            Arc::clone(&schema),
            vec![f0, f1],
            &pool,
            initial_batch_size,
        );
        let stream = builder.create_spillable_merge_stream();
        let batches: Vec<RecordBatch> = stream.try_collect().await?;

        // All rows are still present.
        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total_rows, (2 * n) as usize);

        // Each run was re-spilled once, so each is capped at half the original batch
        // size and the merge caps its output at that half - NOT a quarter. A global
        // halving-per-split implementation would have emitted 2048-row batches here.
        let expected_batch_size = initial_batch_size / 2;
        let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0);
        assert_eq!(
            max_batch_rows, expected_batch_size,
            "two re-spills must halve (not quarter) the output: expected \
             {expected_batch_size}-row batches, got a largest batch of \
             {max_batch_rows} rows"
        );

        Ok(())
    }

    #[test]
    fn spill_merge_fan_in_is_unlimited_by_default() {
        assert_eq!(effective_spill_merge_fan_in(0), usize::MAX);
    }

    #[test]
    fn spill_merge_fan_in_preserves_merge_progress() {
        assert_eq!(effective_spill_merge_fan_in(1), 2);
        assert_eq!(effective_spill_merge_fan_in(2), 2);
        assert_eq!(effective_spill_merge_fan_in(8), 8);
    }

    #[test]
    fn spill_merge_phase_respects_configured_fan_in() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
        let runtime = RuntimeEnvBuilder::new()
            .with_max_spill_merge_fan_in(2)
            .build_arc()?;
        let spill_manager = SpillManager::new(
            Arc::clone(&runtime),
            SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
            Arc::clone(&schema),
        );
        let sorted_spill_files = (0..4)
            .map(|idx| {
                Ok(SortedSpillFile {
                    file: runtime
                        .disk_manager
                        .create_tmp_file(&format!("spill fan-in test {idx}"))?,
                    max_record_batch_memory: 1,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        let expr = LexOrdering::new([PhysicalSortExpr::new_default(col("a", &schema)?)])
            .unwrap();
        let reservation =
            MemoryConsumer::new("spill_merge_phase_respects_configured_fan_in")
                .register(&runtime.memory_pool);
        let metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
        let mut builder = MultiLevelMergeBuilder::new(
            spill_manager,
            schema,
            sorted_spill_files,
            vec![],
            expr,
            metrics,
            1024,
            reservation,
            None,
            false,
        );
        let mut merge_reservation = MemoryConsumer::new("spill_merge_fan_in_phase")
            .register(&runtime.memory_pool);

        let (spills, buffer_len) = match builder.get_sorted_spill_files_to_merge(
            1,
            2,
            &mut merge_reservation,
        )? {
            SpillFilesToMerge::Ready(spills, buffer_len) => (spills, buffer_len),
            SpillFilesToMerge::SplitThenRetry(index) => {
                panic!("expected ready spill files, got retry for index {index}")
            }
        };

        assert_eq!(spills.len(), 2);
        assert_eq!(buffer_len, 1);
        assert_eq!(builder.sorted_spill_files.len(), 2);

        Ok(())
    }
}