datafusion_physical_plan/repartition/

mod.rs

// 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.

//! This file implements the [`RepartitionExec`] operator, which maps N input
//! partitions to M output partitions based on a partitioning scheme, optionally
//! maintaining the order of the input rows in the output.

use std::cmp::Ordering;
use std::fmt::{Debug, Display, Formatter};
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use std::task::{Context, Poll};
use std::vec;

use super::common::SharedMemoryReservation;
use super::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
use super::{
    DisplayAs, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use crate::coalesce::LimitedBatchCoalescer;
use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
use crate::hash_utils::create_hashes;
use crate::metrics::{BaselineMetrics, SpillMetrics};
use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr};
use crate::sorts::streaming_merge::StreamingMergeBuilder;
use crate::spill::spill_manager::SpillManager;
use crate::spill::spill_pool::{self, SpillPoolSink, SpillPoolWriter};
use crate::statistics::{ChildStats, StatisticsArgs};
use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter};
use crate::{
    ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning,
    PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count,
};

use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array};
use arrow::compute::take_arrays;
use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type};
use arrow_schema::SortOptions;
use datafusion_common::config::ConfigOptions;
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose};
use datafusion_common::{
    ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint,
    assert_or_internal_err, internal_datafusion_err, internal_err,
    validate_range_split_points,
};
use datafusion_common::{Result, not_impl_err};
use datafusion_common_runtime::SpawnedTask;
use datafusion_execution::TaskContext;
use datafusion_execution::memory_pool::MemoryConsumer;
use datafusion_expr::ColumnarValue;
use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr, RangePartitioning};
use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
use datafusion_physical_expr_common::sort_expr::LexOrdering;
#[cfg(feature = "proto")]
use datafusion_physical_expr_common::sort_expr::{
    sort_exprs_try_from_proto, sort_exprs_try_to_proto,
};
#[cfg(feature = "proto")]
use datafusion_proto_models::protobuf;

use crate::filter_pushdown::{
    ChildPushdownResult, FilterDescription, FilterPushdownPhase,
    FilterPushdownPropagation,
};
use crate::joins::SeededRandomState;
use crate::sort_pushdown::SortOrderPushdownResult;
use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
use futures::stream::Stream;
use futures::{FutureExt, StreamExt, TryStreamExt};
use log::trace;
use parking_lot::Mutex;

mod distributor_channels;
use crate::repartition::distributor_channels::SendError;
use distributor_channels::{
    DistributionReceiver, DistributionSender, channels, partition_aware_channels,
};

/// A batch in the repartition queue - either in memory or spilled to disk.
///
/// This enum represents the two states a batch can be in during repartitioning.
/// The decision to spill is made based on memory availability when sending a batch
/// to an output partition.
///
/// # Batch Flow with Spilling
///
/// ```text
///                      Input Stream   ◀──────┐
///                           │                │
///                           ▼                │
///                    Partition Logic         │
///                           │           `batch_size` not
///                           ▼            reached yet
///                    Coalesce Batch          │
///           ┌───────────────┴────────────────┘
///           ▼
/// `batch_size` reached
///           │
///           └───────────────┐
///                           ▼
///                        try_grow()
///           ┌───────────────┴────────────────┐
///           ▼                                ▼
/// try_grow() succeeds              try_grow() fails
/// (Memory Available)               (Memory Pressure)
///           │                                │
///           ▼                                ▼
/// RepartitionBatch::Memory         spill_writer.push_batch()
/// (batch held in memory)           (batch written to disk)
///           │                                │
///           │                                ▼
///           │                      RepartitionBatch::Spilled
///           │                      (marker - no batch data)
///           └──────────────┬─────────────────┘
///                          │
///                          ▼
///                   Send to channel
///                          │
///                          ▼
///                 Output Stream (poll)
///                          │
///           ┌──────────────┴────────────────┐
///           ▼                               ▼
/// RepartitionBatch::Memory      RepartitionBatch::Spilled
/// Return batch immediately       Poll spill_stream (blocks)
///           └─────────────┬─────────────────┘
///                         │
///                         ▼
///                    Return batch
///               (FIFO order preserved)
/// ```
///
/// See [`RepartitionExec`] for overall architecture and [`StreamState`] for
/// the state machine that handles reading these batches.
#[derive(Debug)]
enum RepartitionBatch {
    /// Batch held in memory (counts against memory reservation)
    Memory(RecordBatch),
    /// Marker indicating a batch was spilled to the partition's SpillPool.
    /// The actual batch can be retrieved by reading from the SpillPoolStream.
    /// This variant contains no data itself - it's just a signal to the reader
    /// to fetch the next batch from the spill stream.
    Spilled,
}

type MaybeBatch = Option<Result<RepartitionBatch>>;
type InputPartitionsToCurrentPartitionSender = Vec<DistributionSender<MaybeBatch>>;
type InputPartitionsToCurrentPartitionReceiver = Vec<DistributionReceiver<MaybeBatch>>;

/// Output channel with its associated memory reservation and spill writer.
///
/// `coalescer` is `None` for preserve-order mode, where downstream
/// [`StreamingMergeBuilder`] performs the batching; otherwise it's a
/// [`SharedCoalescer`] cloned from the per-partition one held by
/// [`PartitionChannels`].
struct OutputChannel {
    sender: DistributionSender<MaybeBatch>,
    reservation: SharedMemoryReservation,
    spill_writer: SpillPoolSink,
    shared_coalescer: Option<SharedCoalescer>,
}

/// The set of spill-pool writers for a single output partition, before they are handed to the
/// per-input tasks. The variant encodes the repartition mode so the wrong writer topology cannot
/// be constructed for a given mode.
enum PartitionSpillWriters {
    /// `preserve_order`: one single-producer FIFO writer per input partition. Each is `take`n
    /// exactly once (moved into the matching input task), so the pool always has one writer.
    PerInput(Vec<Option<SpillPoolSink>>),
    /// Non-preserve-order: one shared writer, cloned into every input task.
    Shared(SpillPoolWriter),
}

impl PartitionSpillWriters {
    /// Hand out the writer for input partition `input`.
    ///
    /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per
    /// input); in `Shared` mode it clones the shared writer.
    fn take_for_input(&mut self, input: usize) -> Result<SpillPoolSink> {
        match self {
            PartitionSpillWriters::PerInput(writers) => {
                writers[input].take().ok_or_else(|| {
                    internal_datafusion_err!(
                        "spill writer for input partition requested more than once"
                    )
                })
            }
            PartitionSpillWriters::Shared(writer) => Ok(writer.new_sink()),
        }
    }
}

impl OutputChannel {
    fn coalesce(&mut self, batch: RecordBatch) -> Result<Vec<RecordBatch>> {
        match &self.shared_coalescer {
            Some(shared) => Ok(shared.push_and_drain(batch)?),
            None => Ok(vec![batch]),
        }
    }

    /// Send a single batch through the channel for `partition`, applying
    /// the memory reservation / spill-writer fallback. Removes the channel
    /// from `self.inner` if the receiver has hung up.
    ///
    /// Used after [`OutputChannel::coalesce`] for performance purposes.
    async fn send(&mut self, batch: RecordBatch) -> Result<(), SendError<MaybeBatch>> {
        let size = batch.get_array_memory_size();

        // Decide the payload outside of any await: never hold a MutexGuard
        // across an await point.
        let (payload, is_memory_batch) = {
            match self.reservation.try_grow(size) {
                Ok(_) => (Ok(RepartitionBatch::Memory(batch)), true),
                Err(_) => match self.spill_writer.push_batch(&batch) {
                    Ok(()) => (Ok(RepartitionBatch::Spilled), false),
                    Err(err) => (Err(err), false),
                },
            }
        };

        let result = self.sender.send(Some(payload)).await;
        if result.is_err() && is_memory_batch {
            self.reservation.shrink(size);
        }
        result
    }

    async fn finalize(mut self) -> Result<()> {
        let Some(shared) = self.shared_coalescer.take() else {
            return Ok(());
        };
        for batch in shared.finalize()? {
            // If this errored, it means that nobody is listening on the other side, which is fine
            // and can happen in certain cases, like when a LIMIT drops the stream that listens.
            let _ = self.send(batch).await;
        }
        Ok(())
    }
}

/// A producer-side coalescer shared across all input tasks targeting a
/// single output partition.
///
/// Bundles the [`LimitedBatchCoalescer`] (behind a [`Mutex`]) with the
/// active-sender counter that tracks how many input tasks may still push
/// into it. The last task to call [`Self::finalize`] is the one that
/// finalizes the coalescer and ships the residual batch.
///
/// Cheap to [`Clone`]: both fields are [`Arc`]s.
#[derive(Clone)]
struct SharedCoalescer {
    inner: Arc<Mutex<LimitedBatchCoalescer>>,
    active_senders: Arc<AtomicUsize>,
}

impl SharedCoalescer {
    fn new(schema: SchemaRef, target_batch_size: usize, num_senders: usize) -> Self {
        Self {
            inner: Arc::new(Mutex::new(LimitedBatchCoalescer::new(
                schema,
                target_batch_size,
                None,
            ))),
            active_senders: Arc::new(AtomicUsize::new(num_senders)),
        }
    }

    /// Push `batch` into the coalescer and drain any newly completed
    /// batches. The mutex is held only briefly.
    fn push_and_drain(&self, batch: RecordBatch) -> Result<Vec<RecordBatch>> {
        let mut acc = Vec::new();
        let mut c = self.inner.lock();
        c.push_batch(batch)?;
        while let Some(b) = c.next_completed_batch() {
            acc.push(b);
        }
        Ok(acc)
    }

    /// Decrement the active-senders counter. If this caller was the last
    /// sender, finalize the coalescer and return its residual batches; if
    /// other senders are still active, return `Ok(None)`.
    fn finalize(&self) -> Result<Vec<RecordBatch>> {
        let was_last = self.active_senders.fetch_sub(1, AtomicOrdering::AcqRel) == 1;
        if !was_last {
            return Ok(vec![]);
        }
        let mut acc = Vec::new();
        let mut c = self.inner.lock();
        c.finish()?;
        while let Some(b) = c.next_completed_batch() {
            acc.push(b);
        }
        Ok(acc)
    }
}

/// Channels and resources for a single output partition.
///
/// Each output partition has channels to receive data from all input partitions.
/// To handle memory pressure, each (input, output) pair gets its own
/// [`SpillPool`](crate::spill::spill_pool) channel via [`spill_pool::channel`].
///
/// # Structure
///
/// For an output partition receiving from N input partitions:
/// - `tx`: N senders (one per input) for sending batches to this output
/// - `rx`: N receivers (one per input) for receiving batches at this output
/// - `spill_writers`: N spill writers (one per input) for writing spilled data
/// - `spill_readers`: N spill readers (one per input) for reading spilled data
///
/// This 1:1 mapping between input partitions and spill channels ensures that
/// batches from each input are processed in FIFO order, even when some batches
/// are spilled to disk and others remain in memory.
///
/// See [`RepartitionExec`] for the overall N×M architecture.
///
/// [`spill_pool::channel`]: crate::spill::spill_pool::spsc_channel
struct PartitionChannels {
    /// Senders for each input partition to send data to this output partition
    tx: InputPartitionsToCurrentPartitionSender,
    /// Receivers for each input partition sending data to this output partition
    rx: InputPartitionsToCurrentPartitionReceiver,
    /// Memory reservation for this output partition
    reservation: SharedMemoryReservation,
    /// Shared coalescer used by all input tasks targeting this output
    /// partition. `None` in preserve-order mode (downstream
    /// `StreamingMergeBuilder` handles batching).
    shared_coalescer: Option<SharedCoalescer>,
    /// Spill writers for writing spilled data, before they are handed to the per-input tasks.
    /// The variant is chosen by the repartition mode (see [`PartitionSpillWriters`]): a dedicated
    /// single-producer FIFO writer per input in preserve-order mode, or one shared writer in
    /// non-preserve-order mode.
    spill_writers: PartitionSpillWriters,
    /// Spill readers for reading spilled data - one per input partition (FIFO semantics).
    /// Each (input, output) pair gets its own reader to maintain proper ordering.
    spill_readers: Vec<SendableRecordBatchStream>,
}

struct ConsumingInputStreamsState {
    /// Channels for sending batches from input partitions to output partitions.
    /// Key is the partition number.
    channels: HashMap<usize, PartitionChannels>,

    /// Helper that ensures that background jobs are killed once they are no longer needed.
    abort_helper: Arc<Vec<SpawnedTask<()>>>,
}

impl Debug for ConsumingInputStreamsState {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConsumingInputStreamsState")
            .field("num_channels", &self.channels.len())
            .field("abort_helper", &self.abort_helper)
            .finish()
    }
}

/// Inner state of [`RepartitionExec`].
#[derive(Default)]
enum RepartitionExecState {
    /// Not initialized yet. This is the default state stored in the RepartitionExec node
    /// upon instantiation.
    #[default]
    NotInitialized,
    /// Input streams are initialized, but they are still not being consumed. The node
    /// transitions to this state when the arrow's RecordBatch stream is created in
    /// RepartitionExec::execute(), but before any message is polled.
    InputStreamsInitialized(Vec<(SendableRecordBatchStream, RepartitionMetrics)>),
    /// The input streams are being consumed. The node transitions to this state when
    /// the first message in the arrow's RecordBatch stream is consumed.
    ConsumingInputStreams(ConsumingInputStreamsState),
}

impl Debug for RepartitionExecState {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            RepartitionExecState::NotInitialized => write!(f, "NotInitialized"),
            RepartitionExecState::InputStreamsInitialized(v) => {
                write!(f, "InputStreamsInitialized({:?})", v.len())
            }
            RepartitionExecState::ConsumingInputStreams(v) => {
                write!(f, "ConsumingInputStreams({v:?})")
            }
        }
    }
}

impl RepartitionExecState {
    fn ensure_input_streams_initialized(
        &mut self,
        input: &Arc<dyn ExecutionPlan>,
        metrics: &ExecutionPlanMetricsSet,
        output_partitions: usize,
        ctx: &Arc<TaskContext>,
    ) -> Result<()> {
        if !matches!(self, RepartitionExecState::NotInitialized) {
            return Ok(());
        }

        let num_input_partitions = input.output_partitioning().partition_count();
        let mut streams_and_metrics = Vec::with_capacity(num_input_partitions);

        for i in 0..num_input_partitions {
            let metrics = RepartitionMetrics::new(i, output_partitions, metrics);

            let timer = metrics.fetch_time.timer();
            let stream = input.execute(i, Arc::clone(ctx))?;
            timer.done();

            streams_and_metrics.push((stream, metrics));
        }
        *self = RepartitionExecState::InputStreamsInitialized(streams_and_metrics);
        Ok(())
    }

    #[expect(clippy::too_many_arguments)]
    fn consume_input_streams(
        &mut self,
        input: &Arc<dyn ExecutionPlan>,
        metrics: &ExecutionPlanMetricsSet,
        partitioning: &Partitioning,
        preserve_order: bool,
        name: &str,
        context: &Arc<TaskContext>,
        spill_manager: SpillManager,
    ) -> Result<&mut ConsumingInputStreamsState> {
        let streams_and_metrics = match self {
            RepartitionExecState::NotInitialized => {
                self.ensure_input_streams_initialized(
                    input,
                    metrics,
                    partitioning.partition_count(),
                    context,
                )?;
                let RepartitionExecState::InputStreamsInitialized(value) = self else {
                    // This cannot happen, as ensure_input_streams_initialized() was just called,
                    // but the compiler does not know.
                    return internal_err!(
                        "Programming error: RepartitionExecState must be in the InputStreamsInitialized state after calling RepartitionExecState::ensure_input_streams_initialized"
                    );
                };
                value
            }
            RepartitionExecState::ConsumingInputStreams(value) => return Ok(value),
            RepartitionExecState::InputStreamsInitialized(value) => value,
        };

        let num_input_partitions = streams_and_metrics.len();
        let num_output_partitions = partitioning.partition_count();
        let coalesce_batches = !preserve_order && !input.boundedness().is_unbounded();

        let spill_manager = Arc::new(spill_manager);

        let (txs, rxs) = if preserve_order {
            // Create partition-aware channels with one channel per (input, output) pair
            // This provides backpressure while maintaining proper ordering
            let (txs_all, rxs_all) =
                partition_aware_channels(num_input_partitions, num_output_partitions);
            // Take transpose of senders and receivers. `state.channels` keeps track of entries per output partition
            let txs = transpose(txs_all);
            let rxs = transpose(rxs_all);
            (txs, rxs)
        } else {
            // Create one channel per *output* partition with backpressure
            let (txs, rxs) = channels(num_output_partitions);
            // Clone sender for each input partitions
            let txs = txs
                .into_iter()
                .map(|item| vec![item; num_input_partitions])
                .collect::<Vec<_>>();
            let rxs = rxs.into_iter().map(|item| vec![item]).collect::<Vec<_>>();
            (txs, rxs)
        };

        let mut channels = HashMap::with_capacity(txs.len());
        for (partition, (tx, rx)) in txs.into_iter().zip(rxs).enumerate() {
            let reservation = Arc::new(
                MemoryConsumer::new(format!("{name}[{partition}]"))
                    .with_can_spill(true)
                    .register(context.memory_pool()),
            );

            // Create spill channels based on mode:
            // - preserve_order: one spill channel per (input, output) pair for proper FIFO ordering
            // - non-preserve-order: one shared spill channel per output partition since all inputs
            //   share the same receiver
            let max_file_size = context
                .session_config()
                .options()
                .execution
                .max_spill_file_size_bytes
                .get();

            let (spill_writers, spill_readers) = if preserve_order {
                // preserve_order: one dedicated single-producer FIFO pool per input partition.
                // Each writer is moved into exactly one input task (never cloned), so the ordering
                // the downstream merge relies on is preserved across the spill boundary.
                let mut writers = Vec::with_capacity(num_input_partitions);
                let mut readers = Vec::with_capacity(num_input_partitions);
                for _ in 0..num_input_partitions {
                    let (writer, reader) = spill_pool::spsc_channel(
                        max_file_size,
                        Arc::clone(&spill_manager),
                    );
                    writers.push(Some(writer));
                    readers.push(reader);
                }
                (PartitionSpillWriters::PerInput(writers), readers)
            } else {
                // non-preserve-order: one shared multi-producer pool per output partition, since
                // all inputs share the same receiver and the output is an unordered multiset.
                let (writer, reader) =
                    spill_pool::mpsc_channel(max_file_size, Arc::clone(&spill_manager));
                (PartitionSpillWriters::Shared(writer), vec![reader])
            };

            // Coalesce on the producer side, before the channel's gate, so
            // the consumer never sees the per-input-task small batches.
            // Skip in preserve-order mode, where `StreamingMergeBuilder`
            // handles batching, and for unbounded inputs, where a residual
            // batch could otherwise be withheld indefinitely.
            let shared_coalescer = coalesce_batches.then(|| {
                SharedCoalescer::new(
                    input.schema(),
                    context.session_config().batch_size(),
                    num_input_partitions,
                )
            });

            channels.insert(
                partition,
                PartitionChannels {
                    tx,
                    rx,
                    reservation,
                    spill_readers,
                    spill_writers,
                    shared_coalescer,
                },
            );
        }

        // launch one async task per *input* partition
        let mut spawned_tasks = Vec::with_capacity(num_input_partitions);
        for (i, (stream, metrics)) in
            std::mem::take(streams_and_metrics).into_iter().enumerate()
        {
            let txs: HashMap<_, _> = channels
                .iter_mut()
                .map(|(partition, channels)| {
                    // Hand this input task its spill writer: in preserve_order mode this moves
                    // the input's dedicated FIFO writer out; otherwise it clones the shared
                    // writer. See [`PartitionSpillWriters::take_for_input`].
                    Ok((
                        *partition,
                        OutputChannel {
                            sender: channels.tx[i].clone(),
                            reservation: Arc::clone(&channels.reservation),
                            spill_writer: channels.spill_writers.take_for_input(i)?,
                            shared_coalescer: channels.shared_coalescer.clone(),
                        },
                    ))
                })
                .collect::<Result<HashMap<_, _>>>()?;

            // Extract senders for wait_for_task before moving txs
            let senders: HashMap<_, _> = txs
                .iter()
                .map(|(partition, channel)| (*partition, channel.sender.clone()))
                .collect();

            let input_task = SpawnedTask::spawn(RepartitionExec::pull_from_input(
                stream,
                txs,
                partitioning.clone(),
                metrics,
                // preserve_order depends on partition index to start from 0
                if preserve_order { 0 } else { i },
                num_input_partitions,
            ));

            // In a separate task, wait for each input to be done
            // (and pass along any errors, including panic!s)
            let wait_for_task =
                SpawnedTask::spawn(RepartitionExec::wait_for_task(input_task, senders));
            spawned_tasks.push(wait_for_task);
        }
        *self = Self::ConsumingInputStreams(ConsumingInputStreamsState {
            channels,
            abort_helper: Arc::new(spawned_tasks),
        });
        match self {
            RepartitionExecState::ConsumingInputStreams(value) => Ok(value),
            _ => unreachable!(),
        }
    }
}

/// A utility that can be used to partition batches based on [`Partitioning`]
pub struct BatchPartitioner {
    state: BatchPartitionerState,
    timer: metrics::Time,
}

enum BatchPartitionerState {
    Hash {
        exprs: Vec<Arc<dyn PhysicalExpr>>,
        partition_reducer: StrengthReducedU64,
        hash_buffer: Vec<u64>,
        indices: Vec<Vec<u32>>,
    },
    RoundRobin {
        num_partitions: usize,
        next_idx: usize,
    },
    Range {
        /// Ordered partitioning key.
        ordering: LexOrdering,
        /// Sort options from the `LexOrdering`
        sort_options: Vec<SortOptions>,
        /// Boundaries between adjacent partitions.
        split_points: Vec<SplitPoint>,
        /// Row indices grouped by output partition
        indices: Vec<Vec<u32>>,
        /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points
        partition_buffer: Vec<ScalarValue>,
    },
}

/// Fixed RandomState used for hash repartitioning to ensure consistent behavior across
/// executions and runs.
pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0);

/// Physical expression that returns the Range partition for each input row.
///
/// This uses the same routing function as [`BatchPartitioner`], so dynamic
/// filtering and repartitioning agree for every [`ScalarValue`] comparison.
#[derive(Debug, Hash, PartialEq, Eq)]
pub struct RangeExpr {
    on_columns: Vec<PhysicalExprRef>,
    split_points: Vec<SplitPoint>,
    sort_options: Vec<SortOptions>,
}

impl RangeExpr {
    /// Creates a Range expression for `on_columns` using the supplied routing
    /// metadata.
    pub fn try_new(
        on_columns: Vec<PhysicalExprRef>,
        range_partitioning: &RangePartitioning,
    ) -> Result<Self> {
        let sort_options = range_partitioning
            .ordering()
            .iter()
            .map(|expr| expr.options)
            .collect();
        Self::try_new_parts(
            on_columns,
            range_partitioning.split_points().to_vec(),
            sort_options,
        )
    }

    fn try_new_parts(
        on_columns: Vec<PhysicalExprRef>,
        split_points: Vec<SplitPoint>,
        sort_options: Vec<SortOptions>,
    ) -> Result<Self> {
        assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key");
        assert_or_internal_err!(
            on_columns.len() == sort_options.len(),
            "RangeExpr key count must match sort options"
        );
        validate_range_split_points(&split_points, &sort_options)?;
        Ok(Self {
            on_columns,
            split_points,
            sort_options,
        })
    }

    /// Get the columns used to compute Range partition IDs.
    pub fn on_columns(&self) -> &[PhysicalExprRef] {
        &self.on_columns
    }

    /// Returns the Range split points used for routing.
    pub fn split_points(&self) -> &[SplitPoint] {
        &self.split_points
    }

    /// Returns the per-key sort options used for routing.
    pub fn sort_options(&self) -> &[SortOptions] {
        &self.sort_options
    }
}

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

impl PhysicalExpr for RangeExpr {
    fn children(&self) -> Vec<&PhysicalExprRef> {
        self.on_columns.iter().collect()
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<PhysicalExprRef>,
    ) -> Result<PhysicalExprRef> {
        assert_or_internal_err!(
            children.len() == self.on_columns.len(),
            "RangeExpr expected {} children, got {}",
            self.on_columns.len(),
            children.len()
        );
        Ok(Arc::new(Self::try_new_parts(
            children,
            self.split_points.clone(),
            self.sort_options.clone(),
        )?))
    }

    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
        Ok(DataType::UInt64)
    }

    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
        Ok(false)
    }

    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
        let arrays = evaluate_expressions_to_arrays(self.on_columns.iter(), batch)?;
        let mut row_key_buffer = Vec::with_capacity(arrays.len());
        let mut partition_ids = Vec::with_capacity(batch.num_rows());
        for row_idx in 0..batch.num_rows() {
            extract_row_at_idx_to_buf(&arrays, row_idx, &mut row_key_buffer)?;
            partition_ids.push(range_partition_id(
                &row_key_buffer,
                &self.split_points,
                &self.sort_options,
            )? as u64);
        }
        Ok(ColumnarValue::Array(Arc::new(UInt64Array::from(
            partition_ids,
        ))))
    }

    fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "range_partition")
    }

    #[cfg(feature = "proto")]
    fn try_to_proto(
        &self,
        ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
    ) -> Result<Option<protobuf::PhysicalExprNode>> {
        // Encode the raw ordered children: rebuilding a `LexOrdering` would
        // deduplicate equivalent children after dynamic-filter remapping.
        let sort_exprs = self
            .on_columns
            .iter()
            .zip(&self.sort_options)
            .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr), *options))
            .collect::<Vec<_>>();
        let sort_expr = sort_exprs_try_to_proto(&sort_exprs, ctx)?;
        let split_point = self
            .split_points
            .iter()
            .map(|split_point| {
                let value = split_point
                    .values()
                    .iter()
                    .map(|value| value.try_into().map_err(Into::into))
                    .collect::<Result<Vec<_>>>()?;
                Ok(protobuf::PhysicalRangeSplitPoint { value })
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Some(protobuf::PhysicalExprNode {
            expr_id: None,
            expr_type: Some(protobuf::physical_expr_node::ExprType::RangeExpr(
                protobuf::PhysicalRangeExprNode {
                    sort_expr,
                    split_point,
                },
            )),
        }))
    }
}

#[cfg(feature = "proto")]
impl RangeExpr {
    /// Reconstructs a [`RangeExpr`] from its protobuf representation.
    pub fn try_from_proto(
        node: &protobuf::PhysicalExprNode,
        ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
    ) -> Result<PhysicalExprRef> {
        // Decode the raw ordered children for the same reason as `try_to_proto`.
        let range_expr = match &node.expr_type {
            Some(protobuf::physical_expr_node::ExprType::RangeExpr(expr)) => expr,
            _ => return internal_err!("PhysicalExprNode is not a RangeExpr"),
        };
        let sort_exprs = sort_exprs_try_from_proto(&range_expr.sort_expr, ctx)?;
        let (on_columns, sort_options) = sort_exprs
            .into_iter()
            .map(|sort_expr| (sort_expr.expr, sort_expr.options))
            .unzip();
        let split_points = range_expr
            .split_point
            .iter()
            .map(|split_point| {
                let values = split_point
                    .value
                    .iter()
                    .map(|value| ScalarValue::try_from(value).map_err(Into::into))
                    .collect::<Result<Vec<_>>>()?;
                Ok(SplitPoint::new(values))
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Arc::new(Self::try_new_parts(
            on_columns,
            split_points,
            sort_options,
        )?))
    }
}

fn range_partition_id(
    row_key: &[ScalarValue],
    split_points: &[SplitPoint],
    sort_options: &[SortOptions],
) -> Result<usize> {
    let mut low = 0;
    let mut high = split_points.len();
    while low < high {
        let mid = low + (high - low) / 2;
        match compare_rows(row_key, split_points[mid].values(), sort_options)? {
            Ordering::Less => high = mid,
            Ordering::Equal | Ordering::Greater => low = mid + 1,
        }
    }
    Ok(low)
}

/// Computes `value % divisor` without division in the hot loop when `divisor`
/// is fixed for many values.
///
/// Hash repartitioning computes a remainder for every row. Integer division is
/// relatively expensive, so this precomputes the strength-reduced form of the
/// divisor: powers of two use a bit mask, and other divisors use a reciprocal
/// multiply to recover the quotient and therefore the remainder. This is the
/// same invariant-divisor optimization compilers use for `%` by a constant.
#[derive(Debug, Clone, Copy)]
enum StrengthReducedU64 {
    PowerOfTwo { mask: u64 },
    Reciprocal { divisor: u64, reciprocal: u128 },
}

impl StrengthReducedU64 {
    fn new(divisor: u64) -> Self {
        debug_assert!(divisor > 0);

        if divisor.is_power_of_two() {
            Self::PowerOfTwo { mask: divisor - 1 }
        } else {
            Self::Reciprocal {
                divisor,
                // ceil(2^128 / divisor), computed without representing 2^128
                reciprocal: u128::MAX / u128::from(divisor) + 1,
            }
        }
    }

    fn partition_indices(self, hash_buffer: &[u64], indices: &mut [Vec<u32>]) {
        match self {
            Self::PowerOfTwo { mask } => {
                for (index, hash) in hash_buffer.iter().enumerate() {
                    indices[(*hash & mask) as usize].push(index as u32);
                }
            }
            Self::Reciprocal {
                divisor,
                reciprocal,
            } => {
                for (index, hash) in hash_buffer.iter().enumerate() {
                    let quotient = Self::quotient(*hash, reciprocal);
                    let partition = *hash - quotient * divisor;
                    indices[partition as usize].push(index as u32);
                }
            }
        }
    }

    #[cfg(test)]
    fn remainder(self, value: u64) -> u64 {
        match self {
            Self::PowerOfTwo { mask } => value & mask,
            Self::Reciprocal {
                divisor,
                reciprocal,
            } => value - Self::quotient(value, reciprocal) * divisor,
        }
    }

    #[inline]
    fn quotient(value: u64, reciprocal: u128) -> u64 {
        let reciprocal_low = reciprocal as u64;
        let reciprocal_high = (reciprocal >> 64) as u64;
        let low_product = u128::from(value) * u128::from(reciprocal_low);
        let high_product = u128::from(value) * u128::from(reciprocal_high);
        let carry = ((high_product & u128::from(u64::MAX)) + (low_product >> 64)) >> 64;

        ((high_product >> 64) + carry) as u64
    }
}

impl BatchPartitioner {
    /// Create a new [`BatchPartitioner`] for hash-based repartitioning.
    ///
    /// # Parameters
    /// - `exprs`: Expressions used to compute the hash for each input row.
    /// - `num_partitions`: Total number of output partitions.
    /// - `timer`: Metric used to record time spent during repartitioning.
    ///
    /// The partition count is fixed for the lifetime of the partitioner, so this
    /// precomputes a strength-reduced reducer for `hash % num_partitions`.
    ///
    /// # Errors
    /// Returns an error if `num_partitions` is zero.
    pub fn new_hash_partitioner(
        exprs: Vec<Arc<dyn PhysicalExpr>>,
        num_partitions: usize,
        timer: metrics::Time,
    ) -> Result<Self> {
        if num_partitions == 0 {
            return internal_err!("Hash repartition requires at least one partition");
        }

        Ok(Self {
            state: BatchPartitionerState::Hash {
                exprs,
                partition_reducer: StrengthReducedU64::new(num_partitions as u64),
                hash_buffer: vec![],
                indices: vec![vec![]; num_partitions],
            },
            timer,
        })
    }

    /// Create a new [`BatchPartitioner`] for round-robin repartitioning.
    ///
    /// # Parameters
    /// - `num_partitions`: Total number of output partitions.
    /// - `timer`: Metric used to record time spent during repartitioning.
    /// - `input_partition`: Index of the current input partition.
    /// - `num_input_partitions`: Total number of input partitions.
    ///
    /// # Notes
    /// The starting output partition is derived from the input partition
    /// to avoid skew when multiple input partitions are used.
    pub fn new_round_robin_partitioner(
        num_partitions: usize,
        timer: metrics::Time,
        input_partition: usize,
        num_input_partitions: usize,
    ) -> Self {
        Self {
            state: BatchPartitionerState::RoundRobin {
                num_partitions,
                next_idx: (input_partition * num_partitions) / num_input_partitions,
            },
            timer,
        }
    }

    /// Create a new [`BatchPartitioner`] for range-based repartitioning.
    ///
    /// # Parameters
    /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions
    /// - `timer`: Metric used to record time spent during repartitioning.
    pub fn new_range_partitioner(
        range_partitioning: &RangePartitioning,
        timer: metrics::Time,
    ) -> Self {
        let ordering = range_partitioning.ordering().clone();
        let split_points = range_partitioning.split_points().to_vec();
        let num_partitions = range_partitioning.partition_count();
        let sort_options: Vec<SortOptions> = ordering.iter().map(|e| e.options).collect();

        Self {
            state: BatchPartitionerState::Range {
                partition_buffer: Vec::with_capacity(ordering.len()),
                ordering,
                sort_options,
                split_points,
                indices: vec![vec![]; num_partitions],
            },
            timer,
        }
    }

    /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme.
    ///
    /// This is a convenience constructor that delegates to the specialized
    /// hash, round-robin, or range constructors depending on the partitioning variant.
    ///
    /// # Parameters
    /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range).
    /// - `timer`: Metric used to record time spent during repartitioning.
    /// - `input_partition`: Index of the current input partition.
    /// - `num_input_partitions`: Total number of input partitions.
    ///
    /// # Errors
    /// Returns an error if the provided partitioning scheme is not supported,
    /// or if hash partitioning is requested with zero output partitions.
    pub fn try_new(
        partitioning: Partitioning,
        timer: metrics::Time,
        input_partition: usize,
        num_input_partitions: usize,
    ) -> Result<Self> {
        match partitioning {
            Partitioning::Hash(exprs, num_partitions) => {
                Self::new_hash_partitioner(exprs, num_partitions, timer)
            }
            Partitioning::RoundRobinBatch(num_partitions) => {
                Ok(Self::new_round_robin_partitioner(
                    num_partitions,
                    timer,
                    input_partition,
                    num_input_partitions,
                ))
            }
            Partitioning::Range(range_repartitioning) => {
                Ok(Self::new_range_partitioner(&range_repartitioning, timer))
            }
            other => {
                not_impl_err!("Unsupported repartitioning scheme {other:?}")
            }
        }
    }

    /// Partition the provided [`RecordBatch`] into one or more partitioned [`RecordBatch`]
    /// based on the [`Partitioning`] specified on construction
    ///
    /// `f` will be called for each partitioned [`RecordBatch`] with the corresponding
    /// partition index. Any error returned by `f` will be immediately returned by this
    /// function without attempting to publish further [`RecordBatch`]
    ///
    /// The time spent repartitioning, not including time spent in `f` will be recorded
    /// to the [`metrics::Time`] provided on construction
    pub fn partition<F>(&mut self, batch: RecordBatch, mut f: F) -> Result<()>
    where
        F: FnMut(usize, RecordBatch) -> Result<()>,
    {
        self.partition_iter(batch)?.try_for_each(|res| match res {
            Ok((partition, batch)) => f(partition, batch),
            Err(e) => Err(e),
        })
    }

    /// Returns an iterator of `(partition_index, RecordBatch)` pairs for the given batch.
    ///
    /// This is useful for async consumers that want to separate CPU-bound partitioning
    /// from I/O. For example, you can iterate results on the async side and send them
    /// through a channel, while performing file I/O on a blocking task:
    ///
    /// ```ignore
    /// for result in partitioner.partition_iter(batch)? {
    ///     let (partition, batch) = result?;
    ///     tx.send((partition, batch)).await?;
    /// }
    /// ```
    ///
    /// The sync [`partition`](Self::partition) method is implemented on top of this.
    pub fn partition_iter(
        &mut self,
        batch: RecordBatch,
    ) -> Result<impl Iterator<Item = Result<(usize, RecordBatch)>> + Send + '_> {
        let it: Box<dyn Iterator<Item = Result<(usize, RecordBatch)>> + Send> =
            match &mut self.state {
                BatchPartitionerState::RoundRobin {
                    num_partitions,
                    next_idx,
                } => {
                    let idx = *next_idx;
                    *next_idx = (*next_idx + 1) % *num_partitions;
                    Box::new(std::iter::once(Ok((idx, batch))))
                }
                BatchPartitionerState::Hash {
                    exprs,
                    partition_reducer,
                    hash_buffer,
                    indices,
                } => {
                    // Tracking time required for distributing indexes across output partitions
                    let timer = self.timer.timer();

                    let arrays =
                        evaluate_expressions_to_arrays(exprs.as_slice(), &batch)?;

                    hash_buffer.clear();
                    hash_buffer.resize(batch.num_rows(), 0);

                    create_hashes(
                        &arrays,
                        REPARTITION_RANDOM_STATE.random_state(),
                        hash_buffer,
                    )?;

                    indices.iter_mut().for_each(|v| v.clear());

                    partition_reducer.partition_indices(hash_buffer, indices);

                    // Finished building index-arrays for output partitions
                    timer.done();

                    let partitioned_batches =
                        Self::partition_grouped_take(&batch, indices, &self.timer)?;

                    Box::new(partitioned_batches.into_iter())
                }
                BatchPartitionerState::Range {
                    ordering,
                    sort_options,
                    split_points,
                    indices,
                    partition_buffer,
                } => {
                    // Tracking time required for distributing indexes across output partitions
                    let timer = self.timer.timer();
                    if split_points.is_empty() {
                        timer.done();
                        Box::new(std::iter::once(Ok((0, batch))))
                    } else {
                        let arrays = evaluate_expressions_to_arrays(
                            ordering.iter().map(|e| &e.expr),
                            &batch,
                        )?;

                        indices.iter_mut().for_each(|v| v.clear());

                        Self::partition_range_indices(
                            &arrays,
                            split_points,
                            sort_options,
                            partition_buffer,
                            indices,
                        )?;

                        // Finished building index-arrays for output partitions
                        timer.done();

                        let partitioned_batches =
                            Self::partition_grouped_take(&batch, indices, &self.timer)?;

                        Box::new(partitioned_batches.into_iter())
                    }
                }
            };

        Ok(it)
    }

    /// Groups input row indices by range partition. This populates `indices[p]` with the
    /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`.
    fn partition_range_indices(
        arrays: &[Arc<dyn Array>],
        split_points: &[SplitPoint],
        sort_options: &[SortOptions],
        row_key_buffer: &mut Vec<ScalarValue>,
        indices: &mut [Vec<u32>],
    ) -> Result<()> {
        let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0);
        for row_idx in 0..num_rows {
            // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row
            extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?;

            let partition =
                range_partition_id(row_key_buffer, split_points, sort_options)?;
            indices[partition].push(row_idx as u32)
        }

        Ok(())
    }

    // return the number of output partitions
    fn num_partitions(&self) -> usize {
        match &self.state {
            BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions,
            BatchPartitionerState::Hash { indices, .. }
            | BatchPartitionerState::Range { indices, .. } => indices.len(),
        }
    }

    /// Build repartitioned hash/range output batches using one `take` per input batch.
    ///
    /// The routers first fills one index vector per output partition. This method
    /// concatenates those index vectors, performs one grouped `take_arrays`, and
    /// then returns each output partition as a slice of the reordered batch.
    ///
    /// For example, given partition indices:
    ///
    /// ```text
    /// partition 0: [2, 5]
    /// partition 1: []
    /// partition 2: [0, 3, 4]
    /// ```
    ///
    /// this method takes rows in `[2, 5, 0, 3, 4]` order once, then returns
    /// `partition 0 = slice(0, 2)` and `partition 2 = slice(2, 3)`.
    fn partition_grouped_take(
        batch: &RecordBatch,
        indices: &mut [Vec<u32>],
        timer: &metrics::Time,
    ) -> Result<Vec<Result<(usize, RecordBatch)>>> {
        let mut partition_ranges = Vec::with_capacity(indices.len());
        let mut reordered_indices = Vec::with_capacity(batch.num_rows());

        for (partition, p_indices) in indices.iter_mut().enumerate() {
            if p_indices.is_empty() {
                continue;
            }

            let start = reordered_indices.len();
            reordered_indices.extend_from_slice(p_indices);
            partition_ranges.push((partition, start, p_indices.len()));
            p_indices.clear();
        }

        if reordered_indices.is_empty() {
            return Ok(vec![]);
        }

        let batches = {
            let _timer = timer.timer();
            let indices_array: PrimitiveArray<UInt32Type> = reordered_indices.into();
            let columns = take_arrays(batch.columns(), &indices_array, None)?;

            let mut options = RecordBatchOptions::new();
            options = options.with_row_count(Some(indices_array.len()));
            let reordered_batch =
                RecordBatch::try_new_with_options(batch.schema(), columns, &options)?;

            partition_ranges
                .into_iter()
                .map(|(partition, start, len)| {
                    Ok((partition, reordered_batch.slice(start, len)))
                })
                .collect()
        };

        Ok(batches)
    }
}

/// Maps `N` input partitions to `M` output partitions based on a
/// [`Partitioning`] scheme.
///
/// # Background
///
/// DataFusion, like most other commercial systems, with the
/// notable exception of DuckDB, uses the "Exchange Operator" based
/// approach to parallelism which works well in practice given
/// sufficient care in implementation.
///
/// DataFusion's planner picks the target number of partitions and
/// then [`RepartitionExec`] redistributes [`RecordBatch`]es to that number
/// of output partitions.
///
/// For example, given `target_partitions=3` (trying to use 3 cores)
/// but scanning an input with 2 partitions, `RepartitionExec` can be
/// used to get 3 even streams of `RecordBatch`es
///
///
/// ```text
///        ▲                  ▲                  ▲
///        │                  │                  │
///        │                  │                  │
///        │                  │                  │
/// ┌───────────────┐  ┌───────────────┐  ┌───────────────┐
/// │    GroupBy    │  │    GroupBy    │  │    GroupBy    │
/// │   (Partial)   │  │   (Partial)   │  │   (Partial)   │
/// └───────────────┘  └───────────────┘  └───────────────┘
///        ▲                  ▲                  ▲
///        └──────────────────┼──────────────────┘
///                           │
///              ┌─────────────────────────┐
///              │     RepartitionExec     │
///              │   (hash/round robin)    │
///              └─────────────────────────┘
///                         ▲   ▲
///             ┌───────────┘   └───────────┐
///             │                           │
///             │                           │
///        .─────────.                 .─────────.
///     ,─'           '─.           ,─'           '─.
///    ;      Input      :         ;      Input      :
///    :   Partition 0   ;         :   Partition 1   ;
///     ╲               ╱           ╲               ╱
///      '─.         ,─'             '─.         ,─'
///         `───────'                   `───────'
/// ```
///
/// # Error Handling
///
/// If any of the input partitions return an error, the error is propagated to
/// all output partitions and inputs are not polled again.
///
/// # Output Ordering
///
/// If more than one stream is being repartitioned, the output will be some
/// arbitrary interleaving (and thus unordered) unless
/// [`Self::with_preserve_order`] specifies otherwise.
///
/// # Batch coalescing
///
/// Repartitioning one [`RecordBatch`] implies creating multiple smaller batches, potentially
/// as many as the number of output partitions. [`RepartitionExec`] makes sure that the returned
/// batches adhere to the configured `datafusion.execution.batch_size` for efficient operations,
/// and for that, it will automatically coalesce batches right after repartitioning for bounded
/// inputs. Coalescing is skipped for unbounded inputs so partial batches are emitted promptly.
///
/// For this, one shared [`LimitedBatchCoalescer`] per output partition is used:
///
/// ```text
///                         ┌───┐                           ┌───┐
///                      ┌─▶│   │────────▶.───────────.     │   │     ┌──────────────────┐
///                      │  └───┘ ┌───┐  ( Coalescer 0 )──▶ ├───┤ ───▶│     Output 0     │
///                      │┌──────▶│   │──▶`───────────'     │   │     └──────────────────┘
///                      ││       └───┘                     └───┘
/// ┌──────────────────┐ ││                                           ┌──────────────────┐
/// │BatchPartitioner 0│─┘│                                           │     Output 1     │
/// └──────────────────┘  │                                           └──────────────────┘
///                       │
/// ┌──────────────────┐  │                ...                        ┌──────────────────┐
/// │BatchPartitioner 1│──┘                                           │     Output 2     │
/// └──────────────────┘                                              └──────────────────┘
///
///                                                                   ┌──────────────────┐
///                                                                   │     Output 3     │
///                                                                   └──────────────────┘
/// ```
///
/// # Spilling Architecture
///
/// RepartitionExec uses [`SpillPool`](crate::spill::spill_pool) channels to handle
/// memory pressure during repartitioning. Each (input partition, output partition)
/// pair gets its own SpillPool channel for FIFO ordering.
///
/// ```text
/// Input Partitions (N)          Output Partitions (M)
/// ────────────────────          ─────────────────────
///
///    Input 0 ──┐                      ┌──▶ Output 0
///              │  ┌──────────────┐    │
///              ├─▶│ SpillPool    │────┤
///              │  │ [In0→Out0]   │    │
///    Input 1 ──┤  └──────────────┘    ├──▶ Output 1
///              │                       │
///              │  ┌──────────────┐    │
///              ├─▶│ SpillPool    │────┤
///              │  │ [In1→Out0]   │    │
///    Input 2 ──┤  └──────────────┘    ├──▶ Output 2
///              │                      │
///              │       ... (N×M SpillPools total)
///              │                      │
///              │  ┌──────────────┐    │
///              └─▶│ SpillPool    │────┘
///                 │ [InN→OutM]   │
///                 └──────────────┘
///
/// Each SpillPool maintains FIFO order for its (input, output) pair.
/// See `RepartitionBatch` for details on the memory/spill decision logic.
/// ```
///
/// # Footnote
///
/// The "Exchange Operator" was first described in the 1989 paper
/// [Encapsulation of parallelism in the Volcano query processing
/// system Paper](https://dl.acm.org/doi/pdf/10.1145/93605.98720)
/// which uses the term "Exchange" for the concept of repartitioning
/// data across threads.
///
/// For more background, please also see the [Optimizing Repartitions in DataFusion] blog.
///
/// [Optimizing Repartitions in DataFusion]: https://datafusion.apache.org/blog/2025/12/15/avoid-consecutive-repartitions
#[derive(Debug, Clone)]
pub struct RepartitionExec {
    /// Input execution plan
    input: Arc<dyn ExecutionPlan>,
    /// Inner state that is initialized when the parent calls .execute() on this node
    /// and consumed as soon as the parent starts consuming this node.
    state: Arc<Mutex<RepartitionExecState>>,
    /// Execution metrics
    metrics: ExecutionPlanMetricsSet,
    /// Boolean flag to decide whether to preserve ordering. If true means
    /// `SortPreservingRepartitionExec`, false means `RepartitionExec`.
    preserve_order: bool,
    /// Cache holding plan properties like equivalences, output partitioning etc.
    cache: Arc<PlanProperties>,
}

#[derive(Debug, Clone)]
struct RepartitionMetrics {
    /// Time in nanos to execute child operator and fetch batches
    fetch_time: metrics::Time,
    /// Repartitioning elapsed time in nanos
    repartition_time: metrics::Time,
    /// Time in nanos for sending resulting batches to channels.
    ///
    /// One metric per output partition.
    send_time: Vec<metrics::Time>,
}

impl RepartitionMetrics {
    pub fn new(
        input_partition: usize,
        num_output_partitions: usize,
        metrics: &ExecutionPlanMetricsSet,
    ) -> Self {
        // Time in nanos to execute child operator and fetch batches
        let fetch_time =
            MetricBuilder::new(metrics).subset_time("fetch_time", input_partition);

        // Time in nanos to perform repartitioning
        let repartition_time =
            MetricBuilder::new(metrics).subset_time("repartition_time", input_partition);

        // Time in nanos for sending resulting batches to channels
        let send_time = (0..num_output_partitions)
            .map(|output_partition| {
                let label =
                    metrics::Label::new("outputPartition", output_partition.to_string());
                MetricBuilder::new(metrics)
                    .with_label(label)
                    .subset_time("send_time", input_partition)
            })
            .collect();

        Self {
            fetch_time,
            repartition_time,
            send_time,
        }
    }
}

impl RepartitionExec {
    /// Input execution plan
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
        &self.input
    }

    /// Partitioning scheme to use
    pub fn partitioning(&self) -> &Partitioning {
        &self.cache.partitioning
    }

    /// Get preserve_order flag of the RepartitionExec
    /// `true` means `SortPreservingRepartitionExec`, `false` means `RepartitionExec`
    pub fn preserve_order(&self) -> bool {
        self.preserve_order
    }

    /// Get name used to display this Exec
    pub fn name(&self) -> &str {
        "RepartitionExec"
    }
}

impl DisplayAs for RepartitionExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
        let input_partition_count = self.input.output_partitioning().partition_count();
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(
                    f,
                    "{}: partitioning={}, input_partitions={}",
                    self.name(),
                    self.partitioning(),
                    input_partition_count,
                )?;

                if self.preserve_order {
                    write!(f, ", preserve_order=true")?;
                } else if input_partition_count <= 1
                    && self.input.output_ordering().is_some()
                {
                    // Make it explicit that repartition maintains sortedness for a single input partition even
                    // when `preserve_sort order` is false
                    write!(f, ", maintains_sort_order=true")?;
                }

                if let Some(sort_exprs) = self.sort_exprs() {
                    write!(f, ", sort_exprs={}", sort_exprs.clone())?;
                }
                Ok(())
            }
            DisplayFormatType::TreeRender => {
                writeln!(f, "partitioning_scheme={}", self.partitioning(),)?;
                let output_partition_count = self.partitioning().partition_count();
                let input_to_output_partition_str =
                    format!("{input_partition_count} -> {output_partition_count}");
                writeln!(
                    f,
                    "partition_count(in->out)={input_to_output_partition_str}"
                )?;

                if self.preserve_order {
                    writeln!(f, "preserve_order={}", self.preserve_order)?;
                }
                Ok(())
            }
        }
    }
}

impl ExecutionPlan for RepartitionExec {
    fn name(&self) -> &'static str {
        "RepartitionExec"
    }

    /// Return a reference to Any that can be used for downcasting
    fn properties(&self) -> &Arc<PlanProperties> {
        &self.cache
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn apply_expressions(
        &self,
        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
    ) -> Result<TreeNodeRecursion> {
        match self.partitioning() {
            Partitioning::Hash(exprs, _) => crate::apply_expression_roots(exprs, f),
            Partitioning::Range(range) => crate::apply_expression_roots(
                range.ordering().iter().map(|sort_expr| &sort_expr.expr),
                f,
            ),
            _ => Ok(TreeNodeRecursion::Continue),
        }
    }

    fn replace_children(
        self: Arc<Self>,
        mut children: Vec<Arc<dyn ExecutionPlan>>,
        options: ReplaceChildrenOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        validate_child_count!(self, children);
        match options.children_properties {
            ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
                input: children.swap_remove(0),
                metrics: ExecutionPlanMetricsSet::new(),
                state: Default::default(),
                ..Self::clone(&*self)
            })),
            ChildrenPropertiesMode::Recompute => {
                let mut repartition = RepartitionExec::try_new(
                    children.swap_remove(0),
                    self.partitioning().clone(),
                )?;
                if self.preserve_order {
                    repartition = repartition.with_preserve_order();
                }
                Ok(Arc::new(repartition))
            }
        }
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        self.replace_children(
            children,
            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
        )
    }

    fn with_new_children_and_same_properties(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        self.replace_children(
            children,
            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
        )
    }

    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
        vec![matches!(self.partitioning(), Partitioning::Hash(_, _))]
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        Self::maintains_input_order_helper(self.input(), self.preserve_order)
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        trace!(
            "Start {}::execute for partition: {}",
            self.name(),
            partition
        );

        let spill_metrics = SpillMetrics::new(&self.metrics, partition);

        let input = Arc::clone(&self.input);
        let partitioning = self.partitioning().clone();
        let metrics = self.metrics.clone();
        let preserve_order = self.sort_exprs().is_some();
        let name = self.name().to_owned();
        let schema = self.schema();
        let schema_captured = Arc::clone(&schema);

        let spill_manager = SpillManager::new(
            Arc::clone(&context.runtime_env()),
            spill_metrics,
            input.schema(),
        );

        // Get existing ordering to use for merging
        let sort_exprs = self.sort_exprs().cloned();

        let state = Arc::clone(&self.state);
        if let Some(mut state) = state.try_lock() {
            state.ensure_input_streams_initialized(
                &input,
                &metrics,
                partitioning.partition_count(),
                &context,
            )?;
        }

        let num_input_partitions = input.output_partitioning().partition_count();

        let stream = futures::stream::once(async move {
            // lock scope
            let (rx, reservation, spill_readers, abort_helper) = {
                // lock mutexes
                let mut state = state.lock();
                let state = state.consume_input_streams(
                    &input,
                    &metrics,
                    &partitioning,
                    preserve_order,
                    &name,
                    &context,
                    spill_manager.clone(),
                )?;

                // now return stream for the specified *output* partition which will
                // read from the channel
                let PartitionChannels {
                    rx,
                    reservation,
                    spill_readers,
                    ..
                } = state
                    .channels
                    .remove(&partition)
                    .expect("partition not used yet");

                (
                    rx,
                    reservation,
                    spill_readers,
                    Arc::clone(&state.abort_helper),
                )
            };

            trace!(
                "Before returning stream in {name}::execute for partition: {partition}"
            );

            if preserve_order {
                // Store streams from all the input partitions:
                // Each input partition gets its own spill reader to maintain proper FIFO ordering
                //
                // Pass None for metrics here — these intermediate streams feed into
                // StreamingMerge which is the actual output. Only the merge's
                // BaselineMetrics should contribute to the operator's reported
                // output_rows. Without this, every row would be counted twice
                // (once by PerPartitionStream, once by StreamingMerge).
                let input_streams = rx
                    .into_iter()
                    .zip(spill_readers)
                    .map(|(receiver, spill_stream)| {
                        // In preserve_order mode, each receiver corresponds to exactly one input partition
                        Box::pin(PerPartitionStream::new(
                            Arc::clone(&schema_captured),
                            receiver,
                            Arc::clone(&abort_helper),
                            Arc::clone(&reservation),
                            spill_stream,
                            1, // Each receiver handles one input partition
                            None,
                        )) as SendableRecordBatchStream
                    })
                    .collect::<Vec<_>>();
                // Note that receiver size (`rx.len()`) and `num_input_partitions` are same.

                // Merge streams (while preserving ordering) coming from
                // input partitions to this partition:
                let fetch = None;
                let merge_reservation =
                    MemoryConsumer::new(format!("{name}[Merge {partition}]"))
                        .register(context.memory_pool());
                StreamingMergeBuilder::new()
                    .with_streams(input_streams)
                    .with_schema(schema_captured)
                    .with_expressions(&sort_exprs.unwrap())
                    .with_metrics(BaselineMetrics::new(&metrics, partition))
                    .with_batch_size(context.session_config().batch_size())
                    .with_fetch(fetch)
                    .with_reservation(merge_reservation)
                    .with_spill_manager(spill_manager)
                    .build()
            } else {
                // Non-preserve-order case: single input stream, so use the first spill reader
                let spill_stream = spill_readers
                    .into_iter()
                    .next()
                    .expect("at least one spill reader should exist");

                Ok(Box::pin(PerPartitionStream::new(
                    schema_captured,
                    rx.into_iter()
                        .next()
                        .expect("at least one receiver should exist"),
                    abort_helper,
                    reservation,
                    spill_stream,
                    num_input_partitions,
                    Some(BaselineMetrics::new(&metrics, partition)),
                )) as SendableRecordBatchStream)
            }
        })
        .try_flatten();
        let stream = RecordBatchStreamAdapter::new(schema, stream);
        Ok(Box::pin(stream))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())
    }

    fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> {
        vec![ChildStats::At(None)]
    }

    fn statistics_from_inputs(
        &self,
        input_stats: &[Arc<Statistics>],
        args: &StatisticsArgs,
    ) -> Result<Arc<Statistics>> {
        if args.partition().is_some() {
            let partition_count = self.partitioning().partition_count();
            // `StatisticsContext::compute` validates the partition index against
            // this same count before calling, so it is non-zero here; guard
            // defensively against a direct call so the division below cannot
            // divide by zero
            assert_or_internal_err!(
                partition_count > 0,
                "RepartitionExec statistics requested for a partition but the partition count is 0"
            );

            let mut stats = input_stats[0].as_ref().clone();

            // Distribute statistics across partitions
            stats.num_rows = stats
                .num_rows
                .get_value()
                .map(|rows| Precision::Inexact(rows / partition_count))
                .unwrap_or(Precision::Absent);
            stats.total_byte_size = stats
                .total_byte_size
                .get_value()
                .map(|bytes| Precision::Inexact(bytes / partition_count))
                .unwrap_or(Precision::Absent);

            // Make all column stats unknown
            stats.column_statistics = stats
                .column_statistics
                .iter()
                .map(|_| ColumnStatistics::new_unknown())
                .collect();

            Ok(Arc::new(stats))
        } else {
            Ok(Arc::clone(&input_stats[0]))
        }
    }

    fn cardinality_effect(&self) -> CardinalityEffect {
        CardinalityEffect::Equal
    }

    fn try_swapping_with_projection(
        &self,
        projection: &ProjectionExec,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
        // If the projection does not narrow the schema, we should not try to push it down.
        if projection.expr().len() >= projection.input().schema().fields().len() {
            return Ok(None);
        }

        // If pushdown is not beneficial or applicable, break it.
        if projection.benefits_from_input_partitioning()[0]
            || !all_columns(projection.expr())
        {
            return Ok(None);
        }

        let new_projection = make_with_child(projection, self.input())?;

        let new_partitioning = match self.partitioning() {
            Partitioning::Hash(partitions, size) => {
                let mut new_partitions = vec![];
                for partition in partitions {
                    let Some(new_partition) =
                        update_expr(partition, projection.expr(), false)?
                    else {
                        return Ok(None);
                    };
                    new_partitions.push(new_partition);
                }
                Partitioning::Hash(new_partitions, *size)
            }
            Partitioning::Range(range_partitioning) => {
                // Rewrite range key expressions through the projection.
                let mut sort_exprs =
                    Vec::with_capacity(range_partitioning.ordering().len());
                for sort_expr in range_partitioning.ordering() {
                    let Some(new_expr) =
                        update_expr(&sort_expr.expr, projection.expr(), false)?
                    else {
                        return Ok(None);
                    };
                    sort_exprs.push(PhysicalSortExpr::new(new_expr, sort_expr.options));
                }

                let Some(ordering) = LexOrdering::new(sort_exprs) else {
                    return internal_err!(
                        "failed to create LexOrdering for range partitioning"
                    );
                };

                Partitioning::Range(RangePartitioning::try_new(
                    ordering,
                    range_partitioning.split_points().to_vec(),
                )?)
            }
            others => others.clone(),
        };

        Ok(Some(Arc::new(RepartitionExec::try_new(
            new_projection,
            new_partitioning,
        )?)))
    }

    fn gather_filters_for_pushdown(
        &self,
        _phase: FilterPushdownPhase,
        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
        _config: &ConfigOptions,
    ) -> Result<FilterDescription> {
        FilterDescription::from_children(parent_filters, &self.children())
    }

    fn handle_child_pushdown_result(
        &self,
        _phase: FilterPushdownPhase,
        child_pushdown_result: ChildPushdownResult,
        _config: &ConfigOptions,
    ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
        Ok(FilterPushdownPropagation::if_all(child_pushdown_result))
    }

    fn try_pushdown_sort(
        &self,
        order: &[PhysicalSortExpr],
    ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
        // RepartitionExec only maintains input order if preserve_order is set
        // or if there's only one partition
        if !self.maintains_input_order()[0] {
            return Ok(SortOrderPushdownResult::Unsupported);
        }

        // Delegate to the child and wrap with a new RepartitionExec
        self.input.try_pushdown_sort(order)?.try_map(|new_input| {
            let mut new_repartition =
                RepartitionExec::try_new(new_input, self.partitioning().clone())?;
            if self.preserve_order {
                new_repartition = new_repartition.with_preserve_order();
            }
            Ok(Arc::new(new_repartition) as Arc<dyn ExecutionPlan>)
        })
    }

    fn repartitioned(
        &self,
        target_partitions: usize,
        _config: &ConfigOptions,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
        use Partitioning::*;
        let mut new_properties = PlanProperties::clone(&self.cache);
        new_properties.partitioning = match new_properties.partitioning {
            RoundRobinBatch(_) => RoundRobinBatch(target_partitions),
            Hash(hash, _) => Hash(hash, target_partitions),
            Range(_) => {
                // Number of partitions is constrained by the split points and cannot be changed
                return Ok(None);
            }
            UnknownPartitioning(_) => UnknownPartitioning(target_partitions),
        };
        Ok(Some(Arc::new(Self {
            input: Arc::clone(&self.input),
            state: Arc::clone(&self.state),
            metrics: self.metrics.clone(),
            preserve_order: self.preserve_order,
            cache: new_properties.into(),
        })))
    }

    #[cfg(feature = "proto")]
    fn try_to_proto(
        &self,
        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<protobuf::PhysicalPlanNode>> {
        let input = ctx.encode_child(self.input())?;

        let partitioning = self.partitioning().try_to_proto(&ctx.expr_ctx())?;

        Ok(Some(protobuf::PhysicalPlanNode {
            physical_plan_type: Some(
                protobuf::physical_plan_node::PhysicalPlanType::Repartition(Box::new(
                    protobuf::RepartitionExecNode {
                        input: Some(Box::new(input)),
                        partitioning: Some(partitioning),
                        preserve_order: self.preserve_order(),
                    },
                )),
            ),
        }))
    }
}

#[cfg(feature = "proto")]
impl RepartitionExec {
    /// Reconstruct a [`RepartitionExec`] from its protobuf representation.
    pub fn try_from_proto(
        node: &protobuf::PhysicalPlanNode,
        ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let repart = crate::expect_plan_variant!(
            node,
            protobuf::physical_plan_node::PhysicalPlanType::Repartition,
            "RepartitionExec",
        );
        let input = ctx.decode_required_child(
            repart.input.as_deref(),
            "RepartitionExec",
            "input",
        )?;
        let input_schema = input.schema();

        let partitioning = repart
            .partitioning
            .as_ref()
            .map(|partitioning| {
                Partitioning::try_from_proto(
                    partitioning,
                    &ctx.expr_ctx(input_schema.as_ref()),
                )
            })
            .transpose()?
            .flatten()
            .ok_or_else(|| {
                datafusion_common::internal_datafusion_err!(
                    "RepartitionExec is missing required field 'partitioning'"
                )
            })?;

        let mut repart_exec = RepartitionExec::try_new(input, partitioning)?;
        if repart.preserve_order {
            repart_exec = repart_exec.with_preserve_order();
        }
        Ok(Arc::new(repart_exec))
    }
}

impl RepartitionExec {
    /// Create a new RepartitionExec, that produces output `partitioning`, and
    /// does not preserve the order of the input (see [`Self::with_preserve_order`]
    /// for more details)
    pub fn try_new(
        input: Arc<dyn ExecutionPlan>,
        partitioning: Partitioning,
    ) -> Result<Self> {
        let preserve_order = false;
        let cache = Self::compute_properties(&input, partitioning, preserve_order);
        Ok(RepartitionExec {
            input,
            state: Default::default(),
            metrics: ExecutionPlanMetricsSet::new(),
            preserve_order,
            cache: Arc::new(cache),
        })
    }

    fn maintains_input_order_helper(
        input: &Arc<dyn ExecutionPlan>,
        preserve_order: bool,
    ) -> Vec<bool> {
        // We preserve ordering when repartition is order preserving variant or input partitioning is 1
        vec![preserve_order || input.output_partitioning().partition_count() <= 1]
    }

    fn eq_properties_helper(
        input: &Arc<dyn ExecutionPlan>,
        preserve_order: bool,
    ) -> EquivalenceProperties {
        // Equivalence Properties
        let mut eq_properties = input.equivalence_properties().clone();
        // If the ordering is lost, reset the ordering equivalence class:
        if !Self::maintains_input_order_helper(input, preserve_order)[0] {
            eq_properties.clear_orderings();
        }
        // When there are more than one input partitions, they will be fused at the output.
        // Therefore, remove per partition constants.
        if input.output_partitioning().partition_count() > 1 {
            eq_properties.clear_per_partition_constants();
        }
        eq_properties
    }

    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
    fn compute_properties(
        input: &Arc<dyn ExecutionPlan>,
        partitioning: Partitioning,
        preserve_order: bool,
    ) -> PlanProperties {
        PlanProperties::new(
            Self::eq_properties_helper(input, preserve_order),
            partitioning,
            input.pipeline_behavior(),
            input.boundedness(),
        )
        .with_scheduling_type(SchedulingType::Cooperative)
        .with_evaluation_type(EvaluationType::Eager)
    }

    /// Specify if this repartitioning operation should preserve the order of
    /// rows from its input when producing output. Preserving order is more
    /// expensive at runtime, so should only be set if the output of this
    /// operator can take advantage of it.
    ///
    /// If the input is not ordered, or has only one partition, this is a no op,
    /// and the node remains a `RepartitionExec`.
    pub fn with_preserve_order(mut self) -> Self {
        self.preserve_order =
                // If the input isn't ordered, there is no ordering to preserve
                self.input.output_ordering().is_some() &&
                // if there is only one input partition, merging is not required
                // to maintain order
                self.input.output_partitioning().partition_count() > 1;
        let eq_properties = Self::eq_properties_helper(&self.input, self.preserve_order);
        Arc::make_mut(&mut self.cache).set_eq_properties(eq_properties);
        self
    }

    /// Return the sort expressions that are used to merge
    fn sort_exprs(&self) -> Option<&LexOrdering> {
        if self.preserve_order {
            self.input.output_ordering()
        } else {
            None
        }
    }

    /// Pulls data from the specified input plan, feeding it to the
    /// output partitions based on the desired partitioning
    ///
    /// `output_channels` holds the output sending channels for each output partition
    async fn pull_from_input(
        mut stream: SendableRecordBatchStream,
        mut output_channels: HashMap<usize, OutputChannel>,
        partitioning: Partitioning,
        metrics: RepartitionMetrics,
        input_partition: usize,
        num_input_partitions: usize,
    ) -> Result<()> {
        let mut partitioner = BatchPartitioner::try_new(
            partitioning,
            metrics.repartition_time.clone(),
            input_partition,
            num_input_partitions,
        )?;

        // While there are still outputs to send to, keep pulling inputs
        let mut batches_until_yield = partitioner.num_partitions();
        while !output_channels.is_empty() {
            // fetch the next batch
            let timer = metrics.fetch_time.timer();
            let result = stream.next().await;
            timer.done();

            // Input is done
            let batch = match result {
                Some(result) => result?,
                None => break,
            };

            // Handle empty batch
            if batch.num_rows() == 0 {
                continue;
            }

            for res in partitioner.partition_iter(batch)? {
                let (partition, batch) = res?;

                let timer = metrics.send_time[partition].timer();
                // if there is still a receiver, send to it
                if let Some(output_channel) = output_channels.get_mut(&partition) {
                    for batch in output_channel.coalesce(batch)? {
                        if output_channel.send(batch).await.is_err() {
                            // If the other end has hung up, it was an early shutdown (e.g. LIMIT)
                            // so ignore this channel from now on.
                            output_channels.remove(&partition);
                            break;
                        }
                    }
                }
                timer.done();
            }

            // If the input stream is endless, we may spin forever and
            // never yield back to tokio.  See
            // https://github.com/apache/datafusion/issues/5278.
            //
            // However, yielding on every batch causes a bottleneck
            // when running with multiple cores. See
            // https://github.com/apache/datafusion/issues/6290
            //
            // Thus, heuristically yield after producing num_partition
            // batches
            //
            // In round robin this is ideal as each input will get a
            // new batch. In hash partitioning it may yield too often
            // on uneven distributions even if some partition can not
            // make progress, but parallelism is going to be limited
            // in that case anyways
            if batches_until_yield == 0 {
                tokio::task::yield_now().await;
                batches_until_yield = partitioner.num_partitions();
            } else {
                batches_until_yield -= 1;
            }
        }

        // End of input for this task. For each output partition we still
        // have a channel to, decrement the active-senders counter; whoever
        // sees the count drop to zero is the last input task and must
        // finalize the shared coalescer and ship its residual.
        for (_, output_channel) in output_channels.drain() {
            output_channel.finalize().await?;
        }

        // Spill writers will auto-finalize when dropped
        // No need for explicit flush
        Ok(())
    }

    /// Waits for `input_task` which is consuming one of the inputs to
    /// complete. Upon each successful completion, sends a `None` to
    /// each of the output tx channels to signal one of the inputs is
    /// complete. Upon error, propagates the errors to all output tx
    /// channels.
    async fn wait_for_task(
        input_task: SpawnedTask<Result<()>>,
        txs: HashMap<usize, DistributionSender<MaybeBatch>>,
    ) {
        // wait for completion, and propagate error
        // note we ignore errors on send (.ok) as that means the receiver has already shutdown.

        match input_task.join().await {
            // Error in joining task
            Err(e) => {
                let e = Arc::new(e);

                for (_, tx) in txs {
                    let err = Err(DataFusionError::Context(
                        "Join Error".to_string(),
                        Box::new(DataFusionError::External(Box::new(Arc::clone(&e)))),
                    ));
                    tx.send(Some(err)).await.ok();
                }
            }
            // Error from running input task
            Ok(Err(e)) => {
                // send the same Arc'd error to all output partitions
                let e = Arc::new(e);

                for (_, tx) in txs {
                    // wrap it because need to send error to all output partitions
                    let err = Err(DataFusionError::from(&e));
                    tx.send(Some(err)).await.ok();
                }
            }
            // Input task completed successfully
            Ok(Ok(())) => {
                // notify each output partition that this input partition has no more data
                for (_partition, tx) in txs {
                    tx.send(None).await.ok();
                }
            }
        }
    }
}

/// State for tracking whether we're reading from memory channel or spill stream.
///
/// This state machine ensures proper ordering when batches are mixed between memory
/// and spilled storage. When a [`RepartitionBatch::Spilled`] marker is received,
/// the stream must block on the spill stream until the corresponding batch arrives.
///
/// # State Machine
///
/// ```text
///                        ┌─────────────────┐
///                   ┌───▶│  ReadingMemory  │◀───┐
///                   │    └────────┬────────┘    │
///                   │             │             │
///                   │     Poll channel          │
///                   │             │             │
///                   │  ┌──────────┼─────────────┐
///                   │  │          │             │
///                   │  ▼          ▼             │
///                   │ Memory   Spilled          │
///       Got batch   │ batch    marker           │
///       from spill  │  │          │             │
///                   │  │          ▼             │
///                   │  │  ┌──────────────────┐  │
///                   │  │  │ ReadingSpilled   │  │
///                   │  │  └────────┬─────────┘  │
///                   │  │           │            │
///                   │  │   Poll spill_stream    │
///                   │  │           │            │
///                   │  │           ▼            │
///                   │  │      Get batch         │
///                   │  │           │            │
///                   └──┴───────────┴────────────┘
///                                  │
///                                  ▼
///                           Return batch
///                     (Order preserved within
///                      (input, output) pair)
/// ```
///
/// The transition to `ReadingSpilled` blocks further channel polling to maintain
/// FIFO ordering - we cannot read the next item from the channel until the spill
/// stream provides the current batch.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StreamState {
    /// Reading from the memory channel (normal operation)
    ReadingMemory,
    /// Waiting for a spilled batch from the spill stream.
    /// Must not poll channel until spilled batch is received to preserve ordering.
    ReadingSpilled,
}

/// This struct converts a receiver to a stream.
/// Receiver receives data on an SPSC channel.
struct PerPartitionStream {
    /// Schema wrapped by Arc
    schema: SchemaRef,

    /// channel containing the repartitioned batches
    receiver: DistributionReceiver<MaybeBatch>,

    /// Handle to ensure background tasks are killed when no longer needed.
    _drop_helper: Arc<Vec<SpawnedTask<()>>>,

    /// Memory reservation.
    reservation: SharedMemoryReservation,

    /// Infinite stream for reading from the spill pool
    spill_stream: SendableRecordBatchStream,

    /// Internal state indicating if we are reading from memory or spill stream
    state: StreamState,

    /// Number of input partitions that have not yet finished.
    /// In non-preserve-order mode, multiple input partitions send to the same channel,
    /// each sending None when complete. We must wait for all of them.
    remaining_partitions: usize,

    /// Execution metrics (None in preserve-order mode where StreamingMerge owns the metrics)
    baseline_metrics: Option<BaselineMetrics>,
}

impl PerPartitionStream {
    fn new(
        schema: SchemaRef,
        receiver: DistributionReceiver<MaybeBatch>,
        drop_helper: Arc<Vec<SpawnedTask<()>>>,
        reservation: SharedMemoryReservation,
        spill_stream: SendableRecordBatchStream,
        num_input_partitions: usize,
        baseline_metrics: Option<BaselineMetrics>,
    ) -> Self {
        Self {
            schema,
            receiver,
            _drop_helper: drop_helper,
            reservation,
            spill_stream,
            state: StreamState::ReadingMemory,
            remaining_partitions: num_input_partitions,
            baseline_metrics,
        }
    }

    fn poll_next_inner(
        self: &mut Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<RecordBatch>>> {
        use futures::StreamExt;
        let elapsed = self
            .baseline_metrics
            .as_ref()
            .map(|m| m.elapsed_compute().clone());
        let _timer = elapsed.as_ref().map(|t| t.timer());

        loop {
            match self.state {
                StreamState::ReadingMemory => {
                    // Poll the memory channel for next message
                    let value = match self.receiver.recv().poll_unpin(cx) {
                        Poll::Ready(v) => v,
                        Poll::Pending => {
                            // Nothing from channel, wait
                            return Poll::Pending;
                        }
                    };

                    match value {
                        Some(Some(v)) => match v {
                            Ok(RepartitionBatch::Memory(batch)) => {
                                // Release memory and return batch
                                self.reservation.shrink(batch.get_array_memory_size());
                                return Poll::Ready(Some(Ok(batch)));
                            }
                            Ok(RepartitionBatch::Spilled) => {
                                // Batch was spilled, transition to reading from spill stream
                                // We must block on spill stream until we get the batch
                                // to preserve ordering
                                self.state = StreamState::ReadingSpilled;
                                continue;
                            }
                            Err(e) => {
                                return Poll::Ready(Some(Err(e)));
                            }
                        },
                        Some(None) => {
                            // One input partition finished
                            self.remaining_partitions -= 1;
                            if self.remaining_partitions == 0 {
                                // All input partitions finished
                                return Poll::Ready(None);
                            }
                            // Continue to poll for more data from other partitions
                            continue;
                        }
                        None => {
                            // Channel closed unexpectedly
                            return Poll::Ready(None);
                        }
                    }
                }
                StreamState::ReadingSpilled => {
                    // Poll spill stream for the spilled batch
                    match self.spill_stream.poll_next_unpin(cx) {
                        Poll::Ready(Some(Ok(batch))) => {
                            self.state = StreamState::ReadingMemory;
                            return Poll::Ready(Some(Ok(batch)));
                        }
                        Poll::Ready(Some(Err(e))) => {
                            return Poll::Ready(Some(Err(e)));
                        }
                        Poll::Ready(None) => {
                            // Spill stream ended — release its resources before
                            // we go back to draining the memory channel.
                            let spill_schema = self.spill_stream.schema();
                            self.spill_stream =
                                Box::pin(EmptyRecordBatchStream::new(spill_schema));
                            self.state = StreamState::ReadingMemory;
                        }
                        Poll::Pending => {
                            // Spilled batch not ready yet, must wait
                            // This preserves ordering by blocking until spill data arrives
                            return Poll::Pending;
                        }
                    }
                }
            }
        }
    }
}

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

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let poll = self.poll_next_inner(cx);
        if let Some(metrics) = &self.baseline_metrics {
            metrics.record_poll(poll)
        } else {
            poll
        }
    }
}

impl RecordBatchStream for PerPartitionStream {
    /// Get the schema
    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;
    use crate::empty::EmptyExec;
    use crate::projection::ProjectionExpr;
    use crate::streaming::{PartitionStream, StreamingTableExec};
    use crate::test::TestMemoryExec;
    use crate::{
        test::{
            assert_is_pending,
            exec::{
                BarrierExec, BlockingExec, ErrorExec, MockExec,
                assert_strong_count_converges_to_zero,
            },
        },
        {collect, expressions::col},
    };

    use arrow::array::{ArrayRef, StringArray, UInt32Array};
    use arrow::datatypes::{DataType, Field, Schema};
    use datafusion_common::ScalarValue;
    use datafusion_common::cast::{as_string_array, as_uint32_array};
    use datafusion_common::exec_err;
    use datafusion_common::test_util::batches_to_sort_string;
    use datafusion_common_runtime::JoinSet;
    use datafusion_execution::config::SessionConfig;
    use datafusion_execution::runtime_env::RuntimeEnvBuilder;
    use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint};
    use insta::assert_snapshot;

    #[derive(Debug)]
    struct UnboundedTestPartition {
        schema: SchemaRef,
        batch: RecordBatch,
    }

    impl PartitionStream for UnboundedTestPartition {
        fn schema(&self) -> &SchemaRef {
            &self.schema
        }

        fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
            let stream = futures::stream::iter([Ok(self.batch.clone())])
                .chain(futures::stream::pending());
            Box::pin(RecordBatchStreamAdapter::new(
                Arc::clone(&self.schema),
                stream,
            ))
        }
    }

    #[test]
    fn range_expr_preserves_duplicate_remapped_children() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::UInt32, false),
            Field::new("b", DataType::UInt32, false),
        ]));
        let sort_options = [SortOptions::new(false, false), SortOptions::new(true, true)];
        let split_points = vec![SplitPoint::new(vec![
            ScalarValue::UInt32(Some(10)),
            ScalarValue::UInt32(Some(20)),
        ])];
        let range_partitioning = RangePartitioning::try_new(
            [
                PhysicalSortExpr::new(col("a", &schema)?, sort_options[0]),
                PhysicalSortExpr::new(col("b", &schema)?, sort_options[1]),
            ]
            .into(),
            split_points.clone(),
        )?;
        let expr = Arc::new(RangeExpr::try_new(
            vec![col("a", &schema)?, col("b", &schema)?],
            &range_partitioning,
        )?);
        let remapped = col("a", &schema)?;
        let rewritten =
            expr.with_new_children(vec![Arc::clone(&remapped), Arc::clone(&remapped)])?;

        let rewritten = rewritten
            .downcast_ref::<RangeExpr>()
            .expect("rewritten expression should remain a RangeExpr");
        assert_eq!(rewritten.on_columns().len(), 2);
        assert!(Arc::ptr_eq(
            &rewritten.on_columns()[0],
            &rewritten.on_columns()[1]
        ));
        assert_eq!(rewritten.sort_options(), sort_options);
        assert_eq!(rewritten.split_points(), split_points);

        Ok(())
    }

    #[test]
    fn strength_reduced_u64_remainder_matches_modulo() {
        let divisors = [
            1,
            2,
            3,
            4,
            5,
            7,
            8,
            10,
            16,
            31,
            32,
            63,
            64,
            65,
            97,
            u64::from(u32::MAX),
            u64::from(u32::MAX) + 1,
            1_u64 << 32,
            (1_u64 << 63) - 1,
            1_u64 << 63,
            u64::MAX - 1,
            u64::MAX,
        ];
        let values = [
            0,
            1,
            2,
            3,
            4,
            5,
            31,
            32,
            33,
            63,
            64,
            65,
            u64::from(u32::MAX) - 1,
            u64::from(u32::MAX),
            u64::from(u32::MAX) + 1,
            (1_u64 << 32) - 1,
            1_u64 << 32,
            (1_u64 << 32) + 1,
            (1_u64 << 63) - 1,
            1_u64 << 63,
            (1_u64 << 63) + 1,
            u64::MAX - 1,
            u64::MAX,
        ];

        for divisor in divisors {
            let reducer = StrengthReducedU64::new(divisor);
            for value in values {
                assert_eq!(
                    reducer.remainder(value),
                    value % divisor,
                    "value={value} divisor={divisor}"
                );
            }

            let mut value = 0x1234_5678_9abc_def0 ^ divisor;
            for _ in 0..10_000 {
                value = value
                    .wrapping_mul(6_364_136_223_846_793_005)
                    .wrapping_add(1_442_695_040_888_963_407);
                assert_eq!(
                    reducer.remainder(value),
                    value % divisor,
                    "value={value} divisor={divisor}"
                );
            }
        }
    }

    #[test]
    fn hash_partitioner_requires_nonzero_partitions() {
        let metrics = ExecutionPlanMetricsSet::new();
        let timer = MetricBuilder::new(&metrics).subset_time("test", 0);

        let err = BatchPartitioner::new_hash_partitioner(vec![], 0, timer)
            .err()
            .expect("zero hash partitions should fail")
            .to_string();

        assert!(
            err.contains("Hash repartition requires at least one partition"),
            "actual: {err}"
        );
    }

    #[tokio::test]
    async fn one_to_many_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let partitions = vec![partition];

        // repartition from 1 input to 4 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(4)).await?;

        assert_eq!(4, output_partitions.len());
        for partition in &output_partitions {
            assert_eq!(1, partition.len());
        }
        assert_eq!(13 * 8, output_partitions[0][0].num_rows());
        assert_eq!(13 * 8, output_partitions[1][0].num_rows());
        assert_eq!(12 * 8, output_partitions[2][0].num_rows());
        assert_eq!(12 * 8, output_partitions[3][0].num_rows());

        Ok(())
    }

    #[tokio::test]
    async fn many_to_one_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // repartition from 3 input to 1 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(1)).await?;

        assert_eq!(1, output_partitions.len());
        assert_eq!(150 * 8, output_partitions[0][0].num_rows());

        Ok(())
    }

    #[tokio::test]
    async fn many_to_many_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // repartition from 3 input to 5 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(5)).await?;

        let total_rows_per_partition = 8 * 50 * 3 / 5;
        assert_eq!(5, output_partitions.len());
        for partition in output_partitions {
            assert_eq!(1, partition.len());
            assert_eq!(total_rows_per_partition, partition[0].num_rows());
        }

        Ok(())
    }

    #[tokio::test]
    async fn many_to_many_hash_partition() -> Result<()> {
        // define input partitions
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        let output_partitions = repartition(
            &schema,
            partitions,
            Partitioning::Hash(vec![col("c0", &schema)?], 8),
        )
        .await?;

        let total_rows: usize = output_partitions
            .iter()
            .map(|x| x.iter().map(|x| x.num_rows()).sum::<usize>())
            .sum();

        assert_eq!(8, output_partitions.len());
        assert_eq!(total_rows, 8 * 50 * 3);

        Ok(())
    }

    #[tokio::test]
    async fn many_to_many_range_partition() -> Result<()> {
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // create_batch values are [1, 2, 3, 4, 5, 6, 7, 8]; split at 3 and 6 yields
        // 2, 3, and 3 rows per batch respectively
        let partitioning =
            u32_range_partitioning(&schema, SortOptions::default(), vec![3, 6])?;

        let output_partitions = repartition(&schema, partitions, partitioning).await?;

        assert_eq!(3, output_partitions.len());
        assert_eq!(300, partition_row_count(&output_partitions[0]));
        assert_eq!(450, partition_row_count(&output_partitions[1]));
        assert_eq!(450, partition_row_count(&output_partitions[2]));
        assert_eq!(
            collect_partition_u32_values(&output_partitions[0])
                .into_iter()
                .flatten()
                .collect::<HashSet<_>>(),
            HashSet::from([1, 2])
        );
        assert_eq!(
            collect_partition_u32_values(&output_partitions[1])
                .into_iter()
                .flatten()
                .collect::<HashSet<_>>(),
            HashSet::from([3, 4, 5])
        );
        assert_eq!(
            collect_partition_u32_values(&output_partitions[2])
                .into_iter()
                .flatten()
                .collect::<HashSet<_>>(),
            HashSet::from([6, 7, 8])
        );

        Ok(())
    }

    #[tokio::test]
    async fn range_repartition_routes_compound_keys() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::UInt32, false),
            Field::new("b", DataType::UInt32, false),
        ]));
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![
                Arc::new(UInt32Array::from(vec![5, 10, 10, 10, 10, 15])),
                Arc::new(UInt32Array::from(vec![1, 1, 3, 5, 7, 0])),
            ],
        )?;
        let partitioning = Partitioning::Range(RangePartitioning::try_new(
            [
                PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()),
                PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()),
            ]
            .into(),
            vec![
                SplitPoint::new(vec![
                    ScalarValue::UInt32(Some(10)),
                    ScalarValue::UInt32(Some(1)),
                ]),
                SplitPoint::new(vec![
                    ScalarValue::UInt32(Some(10)),
                    ScalarValue::UInt32(Some(5)),
                ]),
            ],
        )?);

        let output_partitions =
            repartition(&schema, vec![vec![batch]], partitioning).await?;

        assert_eq!(3, output_partitions.len());
        assert_eq!(
            vec![(5, 1)],
            collect_partition_u32_pairs(&output_partitions[0])
        );
        assert_eq!(
            vec![(10, 1), (10, 3)],
            collect_partition_u32_pairs(&output_partitions[1])
        );
        assert_eq!(
            vec![(10, 5), (10, 7), (15, 0)],
            collect_partition_u32_pairs(&output_partitions[2])
        );

        Ok(())
    }

    #[tokio::test]
    async fn range_repartition_routes_nulls_asc_nulls_last() -> Result<()> {
        let schema = test_schema(true);
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(UInt32Array::from(vec![
                None,
                Some(5),
                Some(10),
                Some(15),
            ]))],
        )?;
        let partitioning =
            u32_range_partitioning(&schema, SortOptions::new(false, false), vec![10])?;

        let output_partitions =
            repartition(&schema, vec![vec![batch]], partitioning).await?;

        assert_eq!(2, output_partitions.len());
        assert_eq!(
            vec![Some(5)],
            collect_partition_u32_values(&output_partitions[0])
        );
        assert_eq!(
            vec![None, Some(10), Some(15)],
            collect_partition_u32_values(&output_partitions[1])
        );

        Ok(())
    }

    #[tokio::test]
    async fn range_repartition_routes_nulls_asc_nulls_first() -> Result<()> {
        let schema = test_schema(true);
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(UInt32Array::from(vec![
                None,
                Some(5),
                Some(10),
                Some(15),
            ]))],
        )?;
        let partitioning =
            u32_range_partitioning(&schema, SortOptions::new(false, true), vec![10])?;

        let output_partitions =
            repartition(&schema, vec![vec![batch]], partitioning).await?;

        assert_eq!(2, output_partitions.len());
        assert_eq!(
            vec![None, Some(5)],
            collect_partition_u32_values(&output_partitions[0])
        );
        assert_eq!(
            vec![Some(10), Some(15)],
            collect_partition_u32_values(&output_partitions[1])
        );

        Ok(())
    }

    #[tokio::test]
    async fn range_repartition_routes_rows_asc() -> Result<()> {
        let schema = test_schema(false);
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 25]))],
        )?;
        let partitioning =
            u32_range_partitioning(&schema, SortOptions::default(), vec![10, 20])?;

        let output_partitions =
            repartition(&schema, vec![vec![batch]], partitioning).await?;

        assert_eq!(3, output_partitions.len());
        assert_eq!(
            vec![Some(5)],
            collect_partition_u32_values(&output_partitions[0])
        );
        assert_eq!(
            vec![Some(10), Some(15)],
            collect_partition_u32_values(&output_partitions[1])
        );
        assert_eq!(
            vec![Some(25)],
            collect_partition_u32_values(&output_partitions[2])
        );

        Ok(())
    }

    #[tokio::test]
    async fn range_repartition_routes_rows_desc() -> Result<()> {
        let schema = test_schema(false);
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 20, 25]))],
        )?;
        let partitioning =
            u32_range_partitioning(&schema, SortOptions::new(true, false), vec![20, 10])?;

        let output_partitions =
            repartition(&schema, vec![vec![batch]], partitioning).await?;

        assert_eq!(3, output_partitions.len());
        assert_eq!(
            vec![Some(25)],
            collect_partition_u32_values(&output_partitions[0])
        );
        assert_eq!(
            vec![Some(15), Some(20)],
            collect_partition_u32_values(&output_partitions[1])
        );
        assert_eq!(
            vec![Some(5), Some(10)],
            collect_partition_u32_values(&output_partitions[2])
        );

        Ok(())
    }

    #[tokio::test]
    async fn range_repartition_routes_string_rows() -> Result<()> {
        let task_ctx = Arc::new(TaskContext::default());
        let batch = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["bar", "baz", "foo", "qux"])) as ArrayRef,
        )])?;

        let schema = batch.schema();
        let expr = col("my_awesome_field", &schema)?;
        let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema));
        let partitioning = Partitioning::Range(RangePartitioning::try_new(
            [PhysicalSortExpr::new_default(expr)].into(),
            vec![SplitPoint::new(vec![ScalarValue::Utf8(Some(
                "foo".to_string(),
            ))])],
        )?);
        let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?;

        let mut partition_0 = Vec::new();
        let mut stream = exec.execute(0, Arc::clone(&task_ctx))?;
        while let Some(result) = stream.next().await {
            partition_0.push(result?);
        }

        let mut partition_1 = Vec::new();
        let mut stream = exec.execute(1, task_ctx)?;
        while let Some(result) = stream.next().await {
            partition_1.push(result?);
        }

        assert_eq!(
            vec!["bar", "baz"],
            collect_partition_string_values(&partition_0)
        );
        assert_eq!(
            vec!["foo", "qux"],
            collect_partition_string_values(&partition_1)
        );

        Ok(())
    }

    #[test]
    fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> {
        // Three columns so the projection both narrows the schema (required for
        // swap) and moves the range key from @0 to @1.
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("region", DataType::Utf8, false),
            Field::new("payload", DataType::UInt32, false),
        ]));
        let repartition = Arc::new(RepartitionExec::try_new(
            Arc::new(EmptyExec::new(Arc::clone(&schema))),
            range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?,
        )?);

        let projection =
            projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?;

        let swapped = repartition
            .try_swapping_with_projection(&projection)?
            .expect("swap should succeed when projection keeps the range key");
        let swapped_repartition = swapped
            .downcast_ref::<RepartitionExec>()
            .expect("top node should be RepartitionExec");

        assert!(swapped_repartition.input().is::<ProjectionExec>());
        let range = expect_range_partitioning(swapped_repartition.partitioning());
        assert_eq!(range.ordering()[0].to_string(), "id@1 ASC");
        assert_eq!(
            range.split_points(),
            &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])]
        );

        Ok(())
    }

    #[test]
    fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> {
        // Drop a simple range key.
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("payload", DataType::UInt32, false),
        ]));
        let repartition = Arc::new(RepartitionExec::try_new(
            Arc::new(EmptyExec::new(Arc::clone(&schema))),
            range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?,
        )?);
        let projection =
            projection_on_columns(&(Arc::clone(&repartition) as _), &["payload"])?;
        assert!(
            repartition
                .try_swapping_with_projection(&projection)?
                .is_none()
        );

        // Drop part of a compound range key.
        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::UInt32, false),
            Field::new("b", DataType::UInt32, false),
            Field::new("c", DataType::UInt32, false),
        ]));
        let repartition = Arc::new(RepartitionExec::try_new(
            Arc::new(EmptyExec::new(Arc::clone(&schema))),
            range_partitioning_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?,
        )?);
        let projection =
            projection_on_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?;
        assert!(
            repartition
                .try_swapping_with_projection(&projection)?
                .is_none()
        );

        Ok(())
    }

    #[test]
    fn range_repartition_try_pushdown_sort_when_maintains_order() -> Result<()> {
        let schema =
            Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)]));
        let ordering = LexOrdering::new([PhysicalSortExpr::new(
            col("id", &schema)?,
            SortOptions::default(),
        )])
        .expect("ordering must not be empty");

        // Multi-partition source with preserve_order: Range maintains input order.
        let source = Arc::new(ExactSortPushdownExec::new(
            Arc::clone(&schema),
            2,
            ordering.clone(),
        ));
        let repartition = Arc::new(
            RepartitionExec::try_new(
                source,
                range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?,
            )?
            .with_preserve_order(),
        );
        assert!(repartition.maintains_input_order()[0]);

        match repartition.try_pushdown_sort(ordering.as_ref())? {
            SortOrderPushdownResult::Exact { inner } => {
                let pushed = inner
                    .downcast_ref::<RepartitionExec>()
                    .expect("pushdown should keep RepartitionExec");

                assert!(pushed.preserve_order());
                assert!(pushed.maintains_input_order()[0]);

                let range = expect_range_partitioning(pushed.partitioning());
                assert_eq!(range.ordering()[0].to_string(), "id@0 ASC");
                assert_eq!(
                    inner.properties().output_ordering().map(|o| o.to_string()),
                    Some(ordering.to_string()),
                    "pushed repartition output ordering should match the requested sort"
                );
            }
            other => panic!("expected Exact sort pushdown, got {other:?}"),
        }

        Ok(())
    }

    #[test]
    fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance()
    -> Result<()> {
        let schema =
            Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)]));
        let ordering = LexOrdering::new([PhysicalSortExpr::new(
            col("id", &schema)?,
            SortOptions::default(),
        )])
        .expect("ordering must not be empty");

        // Multi-partition source without preserve_order: Range does not maintain order.
        let source = Arc::new(ExactSortPushdownExec::new(
            Arc::clone(&schema),
            2,
            ordering.clone(),
        ));
        let repartition = Arc::new(RepartitionExec::try_new(
            source,
            range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?,
        )?);
        assert!(!repartition.maintains_input_order()[0]);

        assert!(matches!(
            repartition.try_pushdown_sort(ordering.as_ref())?,
            SortOrderPushdownResult::Unsupported
        ));

        Ok(())
    }

    fn range_partitioning_on_columns(
        schema: &SchemaRef,
        key_columns: &[&str],
        split_points: Vec<Vec<u32>>,
    ) -> Result<Partitioning> {
        let Some(ordering) = LexOrdering::new(
            key_columns
                .iter()
                .map(|name| {
                    Ok(PhysicalSortExpr::new(
                        col(name, schema)?,
                        SortOptions::default(),
                    ))
                })
                .collect::<Result<Vec<_>>>()?,
        ) else {
            return exec_err!("range ordering must not be empty");
        };
        Ok(Partitioning::Range(RangePartitioning::try_new(
            ordering,
            split_points
                .into_iter()
                .map(|values| {
                    SplitPoint::new(
                        values
                            .into_iter()
                            .map(|value| ScalarValue::UInt32(Some(value)))
                            .collect(),
                    )
                })
                .collect(),
        )?))
    }

    fn projection_on_columns(
        input: &Arc<dyn ExecutionPlan>,
        names: &[&str],
    ) -> Result<ProjectionExec> {
        let exprs = names
            .iter()
            .map(|name| {
                Ok(ProjectionExpr {
                    expr: col(name, &input.schema())?,
                    alias: (*name).to_string(),
                })
            })
            .collect::<Result<Vec<_>>>()?;
        ProjectionExec::try_new(exprs, Arc::clone(input))
    }

    fn expect_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning {
        match partitioning {
            Partitioning::Range(range) => range,
            other => panic!("expected Range partitioning, got {other:?}"),
        }
    }

    /// Test source that claims Exact support for any sort pushdown request.
    #[derive(Debug, Clone)]
    struct ExactSortPushdownExec {
        cache: Arc<PlanProperties>,
    }

    impl ExactSortPushdownExec {
        fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self {
            use crate::execution_plan::{Boundedness, EmissionType};
            Self {
                cache: Arc::new(PlanProperties::new(
                    EquivalenceProperties::new_with_orderings(schema, [ordering]),
                    Partitioning::UnknownPartitioning(num_partitions),
                    EmissionType::Incremental,
                    Boundedness::Bounded,
                )),
            }
        }
    }

    impl DisplayAs for ExactSortPushdownExec {
        fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
            write!(f, "ExactSortPushdownExec")
        }
    }

    impl ExecutionPlan for ExactSortPushdownExec {
        fn name(&self) -> &str {
            "ExactSortPushdownExec"
        }

        fn properties(&self) -> &Arc<PlanProperties> {
            &self.cache
        }

        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
            vec![]
        }

        fn apply_expressions(
            &self,
            _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
        ) -> Result<TreeNodeRecursion> {
            Ok(TreeNodeRecursion::Continue)
        }

        fn replace_children(
            self: Arc<Self>,
            _: Vec<Arc<dyn ExecutionPlan>>,
            _: ReplaceChildrenOptions,
        ) -> Result<Arc<dyn ExecutionPlan>> {
            Ok(self)
        }

        fn with_new_children(
            self: Arc<Self>,
            children: Vec<Arc<dyn ExecutionPlan>>,
        ) -> Result<Arc<dyn ExecutionPlan>> {
            self.replace_children(
                children,
                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
            )
        }

        fn execute(
            &self,
            _partition: usize,
            _context: Arc<TaskContext>,
        ) -> Result<SendableRecordBatchStream> {
            Ok(Box::pin(EmptyRecordBatchStream::new(self.schema())))
        }

        fn try_pushdown_sort(
            &self,
            _order: &[PhysicalSortExpr],
        ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
            Ok(SortOrderPushdownResult::Exact {
                inner: Arc::new(self.clone()),
            })
        }
    }

    #[tokio::test]
    async fn test_repartition_with_coalescing() -> Result<()> {
        let schema = test_schema(false);
        // create 50 batches, each having 8 rows
        let partition = create_vec_batches(50);
        let partitions = vec![partition.clone(), partition.clone()];
        let partitioning = Partitioning::RoundRobinBatch(1);

        let session_config = SessionConfig::new().with_batch_size(200);
        let task_ctx = TaskContext::default().with_session_config(session_config);
        let task_ctx = Arc::new(task_ctx);

        // create physical plan
        let exec = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                let batch = result?;
                assert_eq!(200, batch.num_rows());
            }
        }
        Ok(())
    }

    #[tokio::test]
    async fn unbounded_input_emits_before_batch_size() -> Result<()> {
        let schema = test_schema(false);
        let batch = create_batch();
        let source = Arc::new(StreamingTableExec::try_new(
            Arc::clone(&schema),
            vec![Arc::new(UnboundedTestPartition {
                schema: Arc::clone(&schema),
                batch: batch.clone(),
            })],
            None,
            vec![],
            true,
            None,
        )?);
        let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(1))?;
        let session_config = SessionConfig::new().with_batch_size(batch.num_rows() * 2);
        let task_ctx =
            Arc::new(TaskContext::default().with_session_config(session_config));

        let mut stream = exec.execute(0, task_ctx)?;
        let output =
            tokio::time::timeout(std::time::Duration::from_secs(5), stream.next())
                .await
                .expect("unbounded repartition withheld a partial batch")
                .expect("unbounded input ended unexpectedly")?;

        assert_eq!(batch, output);
        Ok(())
    }

    fn test_schema(nullable: bool) -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new(
            "c0",
            DataType::UInt32,
            nullable,
        )]))
    }

    fn u32_range_partitioning(
        schema: &SchemaRef,
        sort_options: SortOptions,
        split_values: Vec<u32>,
    ) -> Result<Partitioning> {
        let expr = col("c0", schema)?;
        Ok(Partitioning::Range(RangePartitioning::try_new(
            [PhysicalSortExpr::new(expr, sort_options)].into(),
            split_values
                .into_iter()
                .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))]))
                .collect(),
        )?))
    }

    fn partition_row_count(batches: &[RecordBatch]) -> usize {
        batches.iter().map(|batch| batch.num_rows()).sum()
    }

    fn collect_partition_u32_values(batches: &[RecordBatch]) -> Vec<Option<u32>> {
        batches
            .iter()
            .flat_map(|batch| {
                let array =
                    as_uint32_array(batch.column(0)).expect("expected UInt32 column");
                (0..array.len())
                    .map(|idx| {
                        if array.is_null(idx) {
                            None
                        } else {
                            Some(array.value(idx))
                        }
                    })
                    .collect::<Vec<_>>()
            })
            .collect()
    }

    fn collect_partition_u32_pairs(batches: &[RecordBatch]) -> Vec<(u32, u32)> {
        batches
            .iter()
            .flat_map(|batch| {
                let a = as_uint32_array(batch.column(0)).expect("expected UInt32 column");
                let b = as_uint32_array(batch.column(1)).expect("expected UInt32 column");
                (0..a.len())
                    .map(|idx| (a.value(idx), b.value(idx)))
                    .collect::<Vec<_>>()
            })
            .collect()
    }

    fn collect_partition_string_values(batches: &[RecordBatch]) -> Vec<&str> {
        batches
            .iter()
            .flat_map(|batch| {
                let array =
                    as_string_array(batch.column(0)).expect("expected Utf8 column");
                (0..array.len())
                    .map(|idx| array.value(idx))
                    .collect::<Vec<_>>()
            })
            .collect()
    }

    async fn repartition(
        schema: &SchemaRef,
        input_partitions: Vec<Vec<RecordBatch>>,
        partitioning: Partitioning,
    ) -> Result<Vec<Vec<RecordBatch>>> {
        let task_ctx = Arc::new(TaskContext::default());
        // create physical plan
        let exec =
            TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        // execute and collect results
        let mut output_partitions = vec![];
        for i in 0..exec.partitioning().partition_count() {
            // execute this *output* partition and collect all batches
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            let mut batches = vec![];
            while let Some(result) = stream.next().await {
                batches.push(result?);
            }
            output_partitions.push(batches);
        }
        Ok(output_partitions)
    }

    #[tokio::test]
    async fn many_to_many_round_robin_within_tokio_task() -> Result<()> {
        let handle: SpawnedTask<Result<Vec<Vec<RecordBatch>>>> =
            SpawnedTask::spawn(async move {
                // define input partitions
                let schema = test_schema(false);
                let partition = create_vec_batches(50);
                let partitions =
                    vec![partition.clone(), partition.clone(), partition.clone()];

                // repartition from 3 input to 5 output
                repartition(&schema, partitions, Partitioning::RoundRobinBatch(5)).await
            });

        let output_partitions = handle.join().await.unwrap().unwrap();

        let total_rows_per_partition = 8 * 50 * 3 / 5;
        assert_eq!(5, output_partitions.len());
        for partition in output_partitions {
            assert_eq!(1, partition.len());
            assert_eq!(total_rows_per_partition, partition[0].num_rows());
        }

        Ok(())
    }

    #[tokio::test]
    async fn unsupported_partitioning() {
        let task_ctx = Arc::new(TaskContext::default());
        // have to send at least one batch through to provoke error
        let batch = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef,
        )])
        .unwrap();

        let schema = batch.schema();
        let input = MockExec::new(vec![Ok(batch)], schema);
        // This generates an error (partitioning type not supported)
        // but only after the plan is executed. The error should be
        // returned and no results produced
        let partitioning = Partitioning::UnknownPartitioning(1);
        let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap();
        let output_stream = exec.execute(0, task_ctx).unwrap();

        // Expect that an error is returned
        let result_string = crate::common::collect(output_stream)
            .await
            .unwrap_err()
            .to_string();
        assert!(
            result_string
                .contains("Unsupported repartitioning scheme UnknownPartitioning(1)"),
            "actual: {result_string}"
        );
    }

    #[tokio::test]
    async fn error_for_input_exec() {
        // This generates an error on a call to execute. The error
        // should be returned and no results produced.

        let task_ctx = Arc::new(TaskContext::default());
        let input = ErrorExec::new();
        let partitioning = Partitioning::RoundRobinBatch(1);
        let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap();

        // Expect that an error is returned
        let result_string = exec.execute(0, task_ctx).err().unwrap().to_string();

        assert!(
            result_string.contains("ErrorExec, unsurprisingly, errored in partition 0"),
            "actual: {result_string}"
        );
    }

    #[tokio::test]
    async fn repartition_with_error_in_stream() {
        let task_ctx = Arc::new(TaskContext::default());
        let batch = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef,
        )])
        .unwrap();

        // input stream returns one good batch and then one error. The
        // error should be returned.
        let err = exec_err!("bad data error");

        let schema = batch.schema();
        let input = MockExec::new(vec![Ok(batch), err], schema);
        let partitioning = Partitioning::RoundRobinBatch(1);
        let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap();

        // Note: this should pass (the stream can be created) but the
        // error when the input is executed should get passed back
        let output_stream = exec.execute(0, task_ctx).unwrap();

        // Expect that an error is returned
        let result_string = crate::common::collect(output_stream)
            .await
            .unwrap_err()
            .to_string();
        assert!(
            result_string.contains("bad data error"),
            "actual: {result_string}"
        );
    }

    #[tokio::test]
    async fn repartition_with_delayed_stream() {
        let task_ctx = Arc::new(TaskContext::default());
        let batch1 = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef,
        )])
        .unwrap();

        let batch2 = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["frob", "baz"])) as ArrayRef,
        )])
        .unwrap();

        // The mock exec doesn't return immediately (instead it
        // requires the input to wait at least once)
        let schema = batch1.schema();
        let expected_batches = vec![batch1.clone(), batch2.clone()];
        let input = MockExec::new(vec![Ok(batch1), Ok(batch2)], schema);
        let partitioning = Partitioning::RoundRobinBatch(1);

        let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap();

        assert_snapshot!(batches_to_sort_string(&expected_batches), @r"
        +------------------+
        | my_awesome_field |
        +------------------+
        | bar              |
        | baz              |
        | foo              |
        | frob             |
        +------------------+
        ");

        let output_stream = exec.execute(0, task_ctx).unwrap();
        let batches = crate::common::collect(output_stream).await.unwrap();

        assert_snapshot!(batches_to_sort_string(&batches), @r"
        +------------------+
        | my_awesome_field |
        +------------------+
        | bar              |
        | baz              |
        | foo              |
        | frob             |
        +------------------+
        ");
    }

    #[tokio::test]
    async fn robin_repartition_with_dropping_output_stream() {
        let task_ctx = Arc::new(TaskContext::default());
        let partitioning = Partitioning::RoundRobinBatch(2);
        // The barrier exec waits to be pinged
        // requires the input to wait at least once)
        let input = Arc::new(make_barrier_exec());

        // partition into two output streams
        let exec = RepartitionExec::try_new(
            Arc::clone(&input) as Arc<dyn ExecutionPlan>,
            partitioning,
        )
        .unwrap();

        let output_stream0 = exec.execute(0, Arc::clone(&task_ctx)).unwrap();
        let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap();

        // now, purposely drop output stream 0
        // *before* any outputs are produced
        drop(output_stream0);

        // Now, start sending input
        let mut background_task = JoinSet::new();
        background_task.spawn(async move {
            input.wait().await;
        });

        // output stream 1 should *not* error and have one of the input batches
        let batches = crate::common::collect(output_stream1).await.unwrap();

        assert_snapshot!(batches_to_sort_string(&batches), @r"
        +------------------+
        | my_awesome_field |
        +------------------+
        | baz              |
        | frob             |
        | gar              |
        | goo              |
        +------------------+
        ");
    }

    #[tokio::test]
    // As the hash results might be different on different platforms or
    // with different compilers, we will compare the same execution with
    // and without dropping the output stream.
    async fn hash_repartition_with_dropping_output_stream() {
        let task_ctx = Arc::new(TaskContext::default());
        let partitioning = Partitioning::Hash(
            vec![Arc::new(crate::expressions::Column::new(
                "my_awesome_field",
                0,
            ))],
            2,
        );

        // We first collect the results without dropping the output stream.
        let input = Arc::new(make_barrier_exec());
        let exec = RepartitionExec::try_new(
            Arc::clone(&input) as Arc<dyn ExecutionPlan>,
            partitioning.clone(),
        )
        .unwrap();
        let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap();
        let mut background_task = JoinSet::new();
        background_task.spawn(async move {
            input.wait().await;
        });
        let batches_without_drop = crate::common::collect(output_stream1).await.unwrap();

        // run some checks on the result
        let items_vec = str_batches_to_vec(&batches_without_drop);
        let items_set: HashSet<&str> = items_vec.iter().copied().collect();
        assert_eq!(items_vec.len(), items_set.len());
        let source_str_set: HashSet<&str> =
            ["foo", "bar", "frob", "baz", "goo", "gar", "grob", "gaz"]
                .iter()
                .copied()
                .collect();
        assert_eq!(items_set.difference(&source_str_set).count(), 0);

        // Now do the same but dropping the stream before waiting for the barrier
        let input = Arc::new(make_barrier_exec());
        let exec = RepartitionExec::try_new(
            Arc::clone(&input) as Arc<dyn ExecutionPlan>,
            partitioning,
        )
        .unwrap();
        let output_stream0 = exec.execute(0, Arc::clone(&task_ctx)).unwrap();
        let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap();
        // now, purposely drop output stream 0
        // *before* any outputs are produced
        drop(output_stream0);
        let mut background_task = JoinSet::new();
        background_task.spawn(async move {
            input.wait().await;
        });
        let batches_with_drop = crate::common::collect(output_stream1).await.unwrap();

        let items_vec_with_drop = str_batches_to_vec(&batches_with_drop);
        let items_set_with_drop: HashSet<&str> =
            items_vec_with_drop.iter().copied().collect();
        assert_eq!(
            items_set_with_drop.symmetric_difference(&items_set).count(),
            0
        );
    }

    fn str_batches_to_vec(batches: &[RecordBatch]) -> Vec<&str> {
        batches
            .iter()
            .flat_map(|batch| {
                assert_eq!(batch.columns().len(), 1);
                let string_array = as_string_array(batch.column(0))
                    .expect("Unexpected type for repartitioned batch");

                string_array
                    .iter()
                    .map(|v| v.expect("Unexpected null"))
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }

    /// Create a BarrierExec that returns two partitions of two batches each
    fn make_barrier_exec() -> BarrierExec {
        let batch1 = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef,
        )])
        .unwrap();

        let batch2 = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["frob", "baz"])) as ArrayRef,
        )])
        .unwrap();

        let batch3 = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["goo", "gar"])) as ArrayRef,
        )])
        .unwrap();

        let batch4 = RecordBatch::try_from_iter(vec![(
            "my_awesome_field",
            Arc::new(StringArray::from(vec!["grob", "gaz"])) as ArrayRef,
        )])
        .unwrap();

        // The barrier exec waits to be pinged
        // requires the input to wait at least once)
        let schema = batch1.schema();
        BarrierExec::new(vec![vec![batch1, batch2], vec![batch3, batch4]], schema)
    }

    #[tokio::test]
    async fn test_drop_cancel() -> Result<()> {
        let task_ctx = Arc::new(TaskContext::default());
        let schema =
            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)]));

        let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 2));
        let refs = blocking_exec.refs();
        let repartition_exec = Arc::new(RepartitionExec::try_new(
            blocking_exec,
            Partitioning::UnknownPartitioning(1),
        )?);

        let fut = collect(repartition_exec, task_ctx);
        let mut fut = fut.boxed();

        assert_is_pending(&mut fut);
        drop(fut);
        assert_strong_count_converges_to_zero(refs).await;

        Ok(())
    }

    #[tokio::test]
    async fn hash_repartition_avoid_empty_batch() -> Result<()> {
        let task_ctx = Arc::new(TaskContext::default());
        let batch = RecordBatch::try_from_iter(vec![(
            "a",
            Arc::new(StringArray::from(vec!["foo"])) as ArrayRef,
        )])
        .unwrap();
        let partitioning = Partitioning::Hash(
            vec![Arc::new(crate::expressions::Column::new("a", 0))],
            2,
        );
        let schema = batch.schema();
        let input = MockExec::new(vec![Ok(batch)], schema);
        let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap();
        let output_stream0 = exec.execute(0, Arc::clone(&task_ctx)).unwrap();
        let batch0 = crate::common::collect(output_stream0).await.unwrap();
        let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap();
        let batch1 = crate::common::collect(output_stream1).await.unwrap();
        assert!(batch0.is_empty() || batch1.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn repartition_with_spilling() -> Result<()> {
        // Test that repartition successfully spills to disk when memory is constrained
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let input_partitions = vec![partition];
        let partitioning = Partitioning::RoundRobinBatch(4);

        // Set up context with very tight memory limit to force spilling
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(1, 1.0)
            .build_arc()?;

        let task_ctx = TaskContext::default().with_runtime(runtime);
        let task_ctx = Arc::new(task_ctx);

        // create physical plan
        let exec =
            TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        // Collect all partitions - should succeed by spilling to disk
        let mut total_rows = 0;
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                let batch = result?;
                total_rows += batch.num_rows();
            }
        }

        // Verify we got all the data (50 batches * 8 rows each)
        assert_eq!(total_rows, 50 * 8);

        // Verify spilling metrics to confirm spilling actually happened
        let metrics = exec.metrics().unwrap();
        assert!(
            metrics.spill_count().unwrap() > 0,
            "Expected spill_count > 0, but got {:?}",
            metrics.spill_count()
        );
        println!("Spilled {} times", metrics.spill_count().unwrap());
        assert!(
            metrics.spilled_bytes().unwrap() > 0,
            "Expected spilled_bytes > 0, but got {:?}",
            metrics.spilled_bytes()
        );
        println!(
            "Spilled {} bytes in {} spills",
            metrics.spilled_bytes().unwrap(),
            metrics.spill_count().unwrap()
        );
        assert!(
            metrics.spilled_rows().unwrap() > 0,
            "Expected spilled_rows > 0, but got {:?}",
            metrics.spilled_rows()
        );
        println!("Spilled {} rows", metrics.spilled_rows().unwrap());

        Ok(())
    }

    #[tokio::test]
    async fn repartition_with_partial_spilling() -> Result<()> {
        // Test that repartition can handle partial spilling (some batches in memory, some spilled)
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let input_partitions = vec![partition];
        let partitioning = Partitioning::RoundRobinBatch(4);

        // With `batch_size = 1024` and a single UInt32 column, each
        // coalesced residual is ~4 KiB. An 8 KiB pool fits one and forces
        // the rest to spill.
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(8 * 1024, 1.0)
            .build_arc()?;

        let session_config = SessionConfig::new().with_batch_size(1024);
        let task_ctx = TaskContext::default()
            .with_runtime(runtime)
            .with_session_config(session_config);
        let task_ctx = Arc::new(task_ctx);

        // create physical plan
        let exec =
            TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        // Collect all partitions - should succeed with partial spilling
        let mut total_rows = 0;
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                let batch = result?;
                total_rows += batch.num_rows();
            }
        }

        // Verify we got all the data (50 batches * 8 rows each)
        assert_eq!(total_rows, 50 * 8);

        // Verify partial spilling metrics
        let metrics = exec.metrics().unwrap();
        let spill_count = metrics.spill_count().unwrap();
        let spilled_rows = metrics.spilled_rows().unwrap();
        let spilled_bytes = metrics.spilled_bytes().unwrap();

        assert!(
            spill_count > 0,
            "Expected some spilling to occur, but got spill_count={spill_count}"
        );
        assert!(
            spilled_rows > 0 && spilled_rows < total_rows,
            "Expected partial spilling (0 < spilled_rows < {total_rows}), but got spilled_rows={spilled_rows}"
        );
        assert!(
            spilled_bytes > 0,
            "Expected some bytes to be spilled, but got spilled_bytes={spilled_bytes}"
        );

        println!(
            "Partial spilling: spilled {} out of {} rows ({:.1}%) in {} spills, {} bytes",
            spilled_rows,
            total_rows,
            (spilled_rows as f64 / total_rows as f64) * 100.0,
            spill_count,
            spilled_bytes
        );

        Ok(())
    }

    #[tokio::test]
    async fn repartition_without_spilling() -> Result<()> {
        // Test that repartition does not spill when there's ample memory
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let input_partitions = vec![partition];
        let partitioning = Partitioning::RoundRobinBatch(4);

        // Set up context with generous memory limit - no spilling should occur
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(10 * 1024 * 1024, 1.0) // 10MB
            .build_arc()?;

        let task_ctx = TaskContext::default().with_runtime(runtime);
        let task_ctx = Arc::new(task_ctx);

        // create physical plan
        let exec =
            TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        // Collect all partitions - should succeed without spilling
        let mut total_rows = 0;
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                let batch = result?;
                total_rows += batch.num_rows();
            }
        }

        // Verify we got all the data (50 batches * 8 rows each)
        assert_eq!(total_rows, 50 * 8);

        // Verify no spilling occurred
        let metrics = exec.metrics().unwrap();
        assert_eq!(
            metrics.spill_count(),
            Some(0),
            "Expected no spilling, but got spill_count={:?}",
            metrics.spill_count()
        );
        assert_eq!(
            metrics.spilled_bytes(),
            Some(0),
            "Expected no bytes spilled, but got spilled_bytes={:?}",
            metrics.spilled_bytes()
        );
        assert_eq!(
            metrics.spilled_rows(),
            Some(0),
            "Expected no rows spilled, but got spilled_rows={:?}",
            metrics.spilled_rows()
        );

        println!("No spilling occurred - all data processed in memory");

        Ok(())
    }

    #[tokio::test]
    async fn oom() -> Result<()> {
        use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};

        // Test that repartition fails with OOM when disk manager is disabled
        let schema = test_schema(false);
        let partition = create_vec_batches(50);
        let input_partitions = vec![partition];
        let partitioning = Partitioning::RoundRobinBatch(4);

        // Setup context with memory limit but NO disk manager (explicitly disabled)
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(1, 1.0)
            .with_disk_manager_builder(
                DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled),
            )
            .build_arc()?;

        let task_ctx = TaskContext::default().with_runtime(runtime);
        let task_ctx = Arc::new(task_ctx);

        // create physical plan
        let exec =
            TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        // Attempt to execute - should fail with ResourcesExhausted error
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            let err = stream.next().await.unwrap().unwrap_err();
            let err = err.find_root();
            assert!(
                matches!(err, DataFusionError::ResourcesExhausted(_)),
                "Wrong error type: {err}",
            );
        }

        Ok(())
    }

    /// Create vector batches
    fn create_vec_batches(n: usize) -> Vec<RecordBatch> {
        let batch = create_batch();
        std::iter::repeat_n(batch, n).collect()
    }

    /// Create batch
    fn create_batch() -> RecordBatch {
        let schema = test_schema(false);
        RecordBatch::try_new(
            schema,
            vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))],
        )
        .unwrap()
    }

    /// Create batches with sequential values for ordering tests
    fn create_ordered_batches(num_batches: usize) -> Vec<RecordBatch> {
        let schema = test_schema(false);
        (0..num_batches)
            .map(|i| {
                let start = (i * 8) as u32;
                RecordBatch::try_new(
                    Arc::clone(&schema),
                    vec![Arc::new(UInt32Array::from(
                        (start..start + 8).collect::<Vec<_>>(),
                    ))],
                )
                .unwrap()
            })
            .collect()
    }

    #[tokio::test]
    async fn test_repartition_ordering_with_spilling() -> Result<()> {
        // Test that repartition preserves ordering when spilling occurs
        // This tests the state machine fix where we must block on spill_stream
        // when a Spilled marker is received, rather than continuing to poll the channel

        let schema = test_schema(false);
        // Create batches with sequential values: batch 0 has [0,1,2,3,4,5,6,7],
        // batch 1 has [8,9,10,11,12,13,14,15], etc.
        let partition = create_ordered_batches(20);
        let input_partitions = vec![partition];

        // Use RoundRobinBatch to ensure predictable ordering
        let partitioning = Partitioning::RoundRobinBatch(2);

        // Set up context with very tight memory limit to force spilling
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(1, 1.0)
            .build_arc()?;

        let task_ctx = TaskContext::default().with_runtime(runtime);
        let task_ctx = Arc::new(task_ctx);

        // create physical plan
        let exec =
            TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?;
        let exec = RepartitionExec::try_new(exec, partitioning)?;

        // Collect all output partitions
        let mut all_batches = Vec::new();
        for i in 0..exec.partitioning().partition_count() {
            let mut partition_batches = Vec::new();
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                let batch = result?;
                partition_batches.push(batch);
            }
            all_batches.push(partition_batches);
        }

        // Verify spilling occurred
        let metrics = exec.metrics().unwrap();
        assert!(
            metrics.spill_count().unwrap() > 0,
            "Expected spilling to occur, but spill_count = 0"
        );

        // Verify ordering is preserved within each partition
        // With RoundRobinBatch, even batches go to partition 0, odd batches to partition 1
        for (partition_idx, batches) in all_batches.iter().enumerate() {
            let mut last_value = None;
            for batch in batches {
                let array = batch
                    .column(0)
                    .as_any()
                    .downcast_ref::<UInt32Array>()
                    .unwrap();

                for i in 0..array.len() {
                    let value = array.value(i);
                    if let Some(last) = last_value {
                        assert!(
                            value > last,
                            "Ordering violated in partition {partition_idx}: {value} is not greater than {last}"
                        );
                    }
                    last_value = Some(value);
                }
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::test::TestMemoryExec;
    use crate::union::UnionExec;
    use arrow::array::{UInt32Array, record_batch};
    use arrow::compute::SortOptions;
    use arrow::datatypes::{DataType, Field, Schema};
    use datafusion_common::assert_batches_eq;
    use datafusion_common::config::ConfigNonZeroUsize;

    use datafusion_physical_expr::expressions::col;

    /// Asserts that the plan is as expected
    ///
    /// `$EXPECTED_PLAN_LINES`: input plan
    /// `$PLAN`: the plan to optimized
    macro_rules! assert_plan {
        ($PLAN: expr,  @ $EXPECTED: expr) => {
            let formatted = crate::displayable($PLAN).indent(true).to_string();

            insta::assert_snapshot!(
                formatted,
                @$EXPECTED
            );
        };
    }

    #[tokio::test]
    async fn test_preserve_order() -> Result<()> {
        let schema = test_schema();
        let sort_exprs = sort_exprs(&schema);
        let source1 = sorted_memory_exec(&schema, sort_exprs.clone());
        let source2 = sorted_memory_exec(&schema, sort_exprs);
        // output has multiple partitions, and is sorted
        let union = UnionExec::try_new(vec![source1, source2])?;
        let exec = RepartitionExec::try_new(union, Partitioning::RoundRobinBatch(10))?
            .with_preserve_order();

        // Repartition should preserve order
        assert_plan!(&exec, @r"
        RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c0@0 ASC
          UnionExec
            DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC
            DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC
        ");
        Ok(())
    }

    #[tokio::test]
    async fn test_preserve_order_one_partition() -> Result<()> {
        let schema = test_schema();
        let sort_exprs = sort_exprs(&schema);
        let source = sorted_memory_exec(&schema, sort_exprs);
        // output is sorted, but has only a single partition, so no need to sort
        let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(10))?
            .with_preserve_order();

        // Repartition should not preserve order
        assert_plan!(&exec, @r"
        RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
          DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC
        ");

        Ok(())
    }

    #[tokio::test]
    async fn test_preserve_order_input_not_sorted() -> Result<()> {
        let schema = test_schema();
        let source1 = memory_exec(&schema);
        let source2 = memory_exec(&schema);
        // output has multiple partitions, but is not sorted
        let union = UnionExec::try_new(vec![source1, source2])?;
        let exec = RepartitionExec::try_new(union, Partitioning::RoundRobinBatch(10))?
            .with_preserve_order();

        // Repartition should not preserve order, as there is no order to preserve
        assert_plan!(&exec, @r"
        RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2
          UnionExec
            DataSourceExec: partitions=1, partition_sizes=[0]
            DataSourceExec: partitions=1, partition_sizes=[0]
        ");
        Ok(())
    }

    #[tokio::test]
    async fn test_preserve_order_with_spilling() -> Result<()> {
        use datafusion_execution::runtime_env::RuntimeEnvBuilder;

        // Create sorted input data across multiple partitions
        // Partition1: [1,3], [5,7], [9,11]
        // Partition2: [2,4], [6,8], [10,12]
        let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap();
        let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap();
        let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap();
        let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap();
        let batch5 = record_batch!(("c0", UInt32, [9, 11])).unwrap();
        let batch6 = record_batch!(("c0", UInt32, [10, 12])).unwrap();
        let schema = batch1.schema();
        let sort_exprs = LexOrdering::new([PhysicalSortExpr {
            expr: col("c0", &schema).unwrap(),
            options: SortOptions::default().asc(),
        }])
        .unwrap();
        let partition1 = vec![batch1.clone(), batch3.clone(), batch5.clone()];
        let partition2 = vec![batch2.clone(), batch4.clone(), batch6.clone()];
        let input_partitions = vec![partition1, partition2];

        // Set up context with tight memory limit to force spilling
        // Sorting needs some non-spillable memory, so 608 bytes should force spilling while still allowing the query to complete
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(608, 1.0)
            .build_arc()?;

        let task_ctx = TaskContext::default().with_runtime(runtime);
        let task_ctx = Arc::new(task_ctx);

        // Create physical plan with order preservation
        let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?
            .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?;
        let exec = Arc::new(exec);
        let exec = Arc::new(TestMemoryExec::update_cache(&exec));
        // Repartition into 3 partitions with order preservation
        // We expect 1 batch per output partition after repartitioning
        let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))?
            .with_preserve_order();

        let mut batches = vec![];

        // Collect all partitions - should succeed by spilling to disk
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                let batch = result?;
                batches.push(batch);
            }
        }

        #[rustfmt::skip]
        let expected = [
            [
                "+----+",
                "| c0 |",
                "+----+",
                "| 1  |",
                "| 2  |",
                "| 3  |",
                "| 4  |",
                "+----+",
            ],
            [
                "+----+",
                "| c0 |",
                "+----+",
                "| 5  |",
                "| 6  |",
                "| 7  |",
                "| 8  |",
                "+----+",
            ],
            [
                "+----+",
                "| c0 |",
                "+----+",
                "| 9  |",
                "| 10 |",
                "| 11 |",
                "| 12 |",
                "+----+",
            ],
        ];

        for (batch, expected) in batches.iter().zip(expected.iter()) {
            assert_batches_eq!(expected, std::slice::from_ref(batch));
        }

        // We should have spilled
        let metrics = exec.metrics().unwrap();
        assert!(
            metrics.spill_count().unwrap() > 0,
            "Expected spilling to occur for order-preserving repartition at this \
             memory limit. If this fails, the memory limit may need adjustment."
        );
        Ok(())
    }

    /// Regression test for order preservation across spill *file rotation*.
    ///
    /// A `preserve_order` repartition relies on each per-(input, output) spill pool delivering
    /// batches in strict FIFO order (see [`spill_pool::spsc_channel`] / [`SpillPoolSink`]). This uses
    /// the same memory profile as [`Self::test_preserve_order_with_spilling`] — which is tuned to
    /// force spilling while still completing — but additionally sets `max_spill_file_size_bytes`
    /// to 1 so every spilled batch lands in its own file. That exercises the FIFO-across-rotation
    /// path: if ordering were lost across rotated files (e.g. by feeding an ordered pool with a
    /// shared multi-producer writer), the downstream `StreamingMerge` would emit out-of-order rows
    /// and the sortedness assertion below would fail.
    #[tokio::test]
    async fn test_preserve_order_with_spill_file_rotation() -> Result<()> {
        use datafusion_execution::config::SessionConfig;
        use datafusion_execution::runtime_env::RuntimeEnvBuilder;

        // Same sorted input as `test_preserve_order_with_spilling`:
        // Partition1: [1,3], [5,7], [9,11]; Partition2: [2,4], [6,8], [10,12]
        let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap();
        let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap();
        let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap();
        let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap();
        let batch5 = record_batch!(("c0", UInt32, [9, 11])).unwrap();
        let batch6 = record_batch!(("c0", UInt32, [10, 12])).unwrap();
        let schema = batch1.schema();
        let sort_exprs = LexOrdering::new([PhysicalSortExpr {
            expr: col("c0", &schema).unwrap(),
            options: SortOptions::default().asc(),
        }])
        .unwrap();
        let partition1 = vec![batch1, batch3, batch5];
        let partition2 = vec![batch2, batch4, batch6];
        let input_partitions = vec![partition1, partition2];

        // Force a new spill file per spilled batch to exercise FIFO across rotation.
        let mut session_config = SessionConfig::new();
        session_config
            .options_mut()
            .execution
            .max_spill_file_size_bytes = ConfigNonZeroUsize::try_new(1).unwrap();
        // Same tight limit as `test_preserve_order_with_spilling`: forces spilling while leaving
        // the merge enough non-spillable headroom to complete.
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(608, 1.0)
            .build_arc()?;
        let task_ctx = Arc::new(
            TaskContext::default()
                .with_session_config(session_config)
                .with_runtime(runtime),
        );

        let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?
            .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?;
        let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec)));
        let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))?
            .with_preserve_order();

        // Each output partition merges sorted substreams, so its rows must be non-decreasing.
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            let mut last: Option<u32> = None;
            while let Some(result) = stream.next().await {
                let batch = result?;
                let col = batch
                    .column(0)
                    .as_any()
                    .downcast_ref::<UInt32Array>()
                    .unwrap();
                for r in 0..col.len() {
                    let v = col.value(r);
                    if let Some(prev) = last {
                        assert!(
                            prev <= v,
                            "output partition {i} not sorted: {prev} came before {v}"
                        );
                    }
                    last = Some(v);
                }
            }
        }

        let metrics = exec.metrics().unwrap();
        assert!(
            metrics.spill_count().unwrap() > 0,
            "Expected spilling to occur for order-preserving repartition at this \
             memory limit. If this fails, the memory limit may need adjustment."
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_hash_partitioning_with_spilling() -> Result<()> {
        use datafusion_execution::runtime_env::RuntimeEnvBuilder;

        // Create input data similar to the round-robin test
        let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap();
        let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap();
        let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap();
        let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap();
        let schema = batch1.schema();

        let partition1 = vec![batch1.clone(), batch3.clone()];
        let partition2 = vec![batch2.clone(), batch4.clone()];
        let input_partitions = vec![partition1, partition2];

        // Set up context with memory limit to test hash partitioning with spilling infrastructure
        let runtime = RuntimeEnvBuilder::default()
            .with_memory_limit(1, 1.0)
            .build_arc()?;

        let task_ctx = TaskContext::default().with_runtime(runtime);
        let task_ctx = Arc::new(task_ctx);

        // Create physical plan with hash partitioning
        let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?;
        let exec = Arc::new(exec);
        let exec = Arc::new(TestMemoryExec::update_cache(&exec));
        // Hash partition into 2 partitions by column c0
        let hash_expr = col("c0", &schema)?;
        let exec =
            RepartitionExec::try_new(exec, Partitioning::Hash(vec![hash_expr], 2))?;

        // Collect all partitions concurrently using JoinSet - this prevents deadlock
        // where the distribution channel gate closes when all output channels are full
        let mut join_set = tokio::task::JoinSet::new();
        for i in 0..exec.partitioning().partition_count() {
            let stream = exec.execute(i, Arc::clone(&task_ctx))?;
            join_set.spawn(async move {
                let mut count = 0;
                futures::pin_mut!(stream);
                while let Some(result) = stream.next().await {
                    let batch = result?;
                    count += batch.num_rows();
                }
                Ok::<usize, DataFusionError>(count)
            });
        }

        // Wait for all partitions and sum the rows
        let mut total_rows = 0;
        while let Some(result) = join_set.join_next().await {
            total_rows += result.unwrap()?;
        }

        // Verify we got all rows back
        let all_batches = [batch1, batch2, batch3, batch4];
        let expected_rows: usize = all_batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total_rows, expected_rows);

        // Verify metrics are available
        let metrics = exec.metrics().unwrap();
        // Just verify the metrics can be retrieved (spilling may or may not occur)
        let spill_count = metrics.spill_count().unwrap_or(0);
        assert!(spill_count > 0);
        let spilled_bytes = metrics.spilled_bytes().unwrap_or(0);
        assert!(spilled_bytes > 0);
        let spilled_rows = metrics.spilled_rows().unwrap_or(0);
        assert!(spilled_rows > 0);

        Ok(())
    }

    #[tokio::test]
    async fn test_repartition() -> Result<()> {
        let schema = test_schema();
        let sort_exprs = sort_exprs(&schema);
        let source = sorted_memory_exec(&schema, sort_exprs);
        // output is sorted, but has only a single partition, so no need to sort
        let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(10))?
            .repartitioned(20, &Default::default())?
            .unwrap();

        // Repartition should not preserve order
        assert_plan!(exec.as_ref(), @r"
        RepartitionExec: partitioning=RoundRobinBatch(20), input_partitions=1, maintains_sort_order=true
          DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC
        ");
        Ok(())
    }

    #[test]
    fn test_range_repartitioned_returns_none() -> Result<()> {
        let schema = test_schema();
        let source = memory_exec(&schema);
        let partitioning = Partitioning::Range(RangePartitioning::try_new(
            [PhysicalSortExpr::new(
                col("c0", &schema)?,
                SortOptions::default(),
            )]
            .into(),
            vec![
                SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]),
                SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]),
            ],
        )?);
        let exec = RepartitionExec::try_new(source, partitioning)?;

        let mut expressions = vec![];
        exec.apply_expressions(&mut |expr| {
            expressions.push(expr.to_string());
            Ok(TreeNodeRecursion::Continue)
        })?;
        assert_eq!(expressions, ["c0@0"]);

        // Range partition count is fixed by split points, so repartitioned()
        // cannot change it to an arbitrary target.
        let result = exec.repartitioned(10, &Default::default())?;
        assert!(
            result.is_none(),
            "range repartitioning should not support changing partition count"
        );
        Ok(())
    }

    fn test_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
    }

    fn sort_exprs(schema: &Schema) -> LexOrdering {
        [PhysicalSortExpr {
            expr: col("c0", schema).unwrap(),
            options: SortOptions::default(),
        }]
        .into()
    }

    fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
        TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(schema), None).unwrap()
    }

    fn sorted_memory_exec(
        schema: &SchemaRef,
        sort_exprs: LexOrdering,
    ) -> Arc<dyn ExecutionPlan> {
        let exec = TestMemoryExec::try_new(&[vec![]], Arc::clone(schema), None)
            .unwrap()
            .try_with_sort_information(vec![sort_exprs])
            .unwrap();
        let exec = Arc::new(exec);
        Arc::new(TestMemoryExec::update_cache(&exec))
    }

    /// preserve_order repartition should not double-count
    /// output rows.
    #[tokio::test]
    async fn test_preserve_order_output_rows_not_double_counted() -> Result<()> {
        use datafusion_execution::TaskContext;

        // Two sorted input partitions, 2 rows each (4 total)
        let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap();
        let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap();
        let schema = batch1.schema();
        let sort_exprs = sort_exprs(&schema);

        let input_partitions = vec![vec![batch1], vec![batch2]];
        let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?
            .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?;
        let exec = Arc::new(exec);
        let exec = Arc::new(TestMemoryExec::update_cache(&exec));

        let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))?
            .with_preserve_order();

        let task_ctx = Arc::new(TaskContext::default());
        let mut total_rows = 0;
        for i in 0..exec.partitioning().partition_count() {
            let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
            while let Some(result) = stream.next().await {
                total_rows += result?.num_rows();
            }
        }

        assert_eq!(total_rows, 4, "actual rows collected should be 4");

        let metrics = exec.metrics().unwrap();
        let reported_output_rows = metrics.output_rows().unwrap();
        assert_eq!(
            reported_output_rows, total_rows,
            "metrics output_rows ({reported_output_rows}) should match \
             actual rows collected ({total_rows}), not double-count"
        );

        Ok(())
    }
}