use crate::queryplanner::check_memory::CheckMemoryExec;
use crate::queryplanner::group_by_limit_aggregate::GroupByLimitAggregateExec;
use crate::queryplanner::inline_aggregate::{InlineAggregateExec, InlineAggregateMode};
use crate::queryplanner::planning::WorkerExec;
use crate::queryplanner::query_executor::ClusterSendExec;
use crate::queryplanner::tail_limit::TailLimitExec;
use crate::queryplanner::topk::AggregateTopKExec;
use datafusion::arrow::compute::SortOptions;
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::common::{internal_datafusion_err, HashMap};
use datafusion::config::ConfigOptions;
use datafusion::error::DataFusionError;
use datafusion::physical_expr::{
    LexOrdering, LexRequirement, PhysicalSortExpr, PhysicalSortRequirement,
};
use datafusion::physical_optimizer::limit_pushdown::LimitPushdown;
use datafusion::physical_optimizer::PhysicalOptimizerRule as _;
use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::expressions::Column;
use datafusion::physical_plan::filter::FilterExec;
use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::union::UnionExec;
use datafusion::physical_plan::{
    ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PhysicalExpr,
};
use itertools::Itertools as _;
use std::collections::HashSet;
use std::sync::Arc;

/// Transforms from:
///     AggregateFinal
///     `- AggregatePartial
///        `- ClusterSend
/// to:
///     AggregateFinal
///     `- ClusterSend
///        `- AggregatePartial
///
/// The latter gives results in more parallelism and less network.
pub fn push_aggregate_to_workers(
    p_final: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let p_final_agg: &AggregateExec;
    let p_partial: &Arc<dyn ExecutionPlan>;
    if let Some(a) = p_final.as_any().downcast_ref::<AggregateExec>() {
        if matches!(
            a.mode(),
            AggregateMode::Final | AggregateMode::FinalPartitioned
        ) {
            p_final_agg = a;
            p_partial = a.input();
        } else {
            return Ok(p_final);
        }
    } else {
        return Ok(p_final);
    }

    let agg;
    if let Some(a) = p_partial.as_any().downcast_ref::<AggregateExec>() {
        agg = a;
    } else {
        return Ok(p_final);
    }
    if *agg.mode() != AggregateMode::Partial {
        return Ok(p_final);
    }

    let p_final_input: Arc<dyn ExecutionPlan> =
        if let Some(cs) = agg.input().as_any().downcast_ref::<ClusterSendExec>() {
            let clustersend_input = p_partial
                .clone()
                .with_new_children(vec![cs.input_for_optimizations.clone()])?;

            // Note that required_input_ordering is applicable when p_final_agg has a Sorted input mode.

            // Router plan, replace partial aggregate with cluster send.
            Arc::new(
                cs.with_changed_schema(
                    clustersend_input,
                    p_final_agg
                        .required_input_ordering()
                        .into_iter()
                        .next()
                        .unwrap(),
                ),
            )
        } else if let Some(w) = agg.input().as_any().downcast_ref::<WorkerExec>() {
            let worker_input = p_partial.clone().with_new_children(vec![w.input.clone()])?;

            // Worker plan, execute partial aggregate inside the worker.
            Arc::new(WorkerExec::new(
                worker_input,
                w.max_batch_rows,
                // TODO upgrade DF: WorkerExec limit_and_reverse must be wrong here.  Should be
                // None.  Same applies to cs.with_changed_schema.
                w.limit_and_reverse.clone(),
                p_final_agg
                    .required_input_ordering()
                    .into_iter()
                    .next()
                    .unwrap(),
                w.worker_sort_and_limit.clone(),
                w.properties().output_partitioning().partition_count(),
            ))
        } else {
            return Ok(p_final);
        };

    // We change AggregateMode::FinalPartitioned to AggregateMode::Final, because the ClusterSend
    // node ends up creating an incompatible partitioning for FinalPartitioned.  Some other ideas,
    // like adding a RepartitionExec node, would just be redundant with the behavior of
    // AggregateExec::Final, and also, tricky to set up with the ideal number of partitions in the
    // middle of optimization passes.  Having ClusterSend be able to pass through hash partitions in
    // some form is another option.
    let p_final_input_schema = p_final_input.schema();
    Ok(Arc::new(AggregateExec::try_new(
        AggregateMode::Final,
        p_final_agg.group_expr().clone(),
        p_final_agg.aggr_expr().to_vec(),
        p_final_agg.filter_expr().to_vec(),
        p_final_input,
        p_final_input_schema,
    )?))
}

/// A Linear aggregate doesn't use input ordering, but the scan still merges its partitions with a
/// SortPreservingMergeExec when the index has a sort key (e.g. picked from the filters for partition
/// pruning). Replace such merges under the aggregate with plain partition coalescing: the per-row
/// key comparisons are pure waste there, and particularly bad when the leading sort columns are
/// constant (low-cardinality / NULL-heavy), turning every comparison into a full-length tie.
///
/// Always applied to global (no GROUP BY) aggregates. A grouped hash aggregate is equally safe --
/// CoalescePartitions still yields a single output partition, hence one hash table and no
/// per-partition memory blowup -- but it is gated behind `coalesce_grouped_hash` while the win is
/// validated.
pub fn drop_sort_merge_under_global_aggregate(
    p: Arc<dyn ExecutionPlan>,
    coalesce_grouped_hash: bool,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let Some(agg) = p.as_any().downcast_ref::<AggregateExec>() else {
        return Ok(p);
    };
    if !matches!(agg.input_order_mode(), InputOrderMode::Linear) {
        return Ok(p);
    }
    if !agg.group_expr().expr().is_empty() && !coalesce_grouped_hash {
        return Ok(p);
    }
    // Order-sensitive aggregates (first_value, array_agg(ORDER BY), ...) stay Linear but still
    // require their input ordered -- leave their merge in place.
    if agg.required_input_ordering()[0].is_some() {
        return Ok(p);
    }
    let new_input = replace_merges_with_coalesce(agg.input())?;
    if Arc::ptr_eq(&new_input, agg.input()) {
        return Ok(p);
    }
    p.with_new_children(vec![new_input])
}

/// Replaces sort preserving merges with plain partition coalescing, looking through the nodes
/// that don't require an input ordering of their own.
fn replace_merges_with_coalesce(
    p: &Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let p_any = p.as_any();
    if let Some(merge) = p_any.downcast_ref::<SortPreservingMergeExec>() {
        if merge.fetch().is_some() {
            return Ok(p.clone());
        }
        let child = replace_merges_with_coalesce(merge.input())?;
        return Ok(Arc::new(CoalescePartitionsExec::new(child)));
    }
    if p_any.is::<FilterExec>() || p_any.is::<ProjectionExec>() || p_any.is::<UnionExec>() {
        let new_children = p
            .children()
            .into_iter()
            .map(replace_merges_with_coalesce)
            .collect::<Result<Vec<_>, _>>()?;
        if p.children()
            .into_iter()
            .zip(new_children.iter())
            .all(|(old, new)| Arc::ptr_eq(old, new))
        {
            return Ok(p.clone());
        }
        return p.clone().with_new_children(new_children);
    }
    Ok(p.clone())
}

/// Transforms from:
///     AggregatePartial, Sorted
///     `- SortPreservingMerge
///        `- source(N partitions)
/// to:
///     SortPreservingMerge
///     `- AggregatePartial, Sorted (executed per partition)
///        `- source(N partitions)
///
/// The merge then carries one row per group per partition instead of all raw rows. Duplicate
/// group keys from different partitions are adjacent in the merged stream and get combined by
/// the Final aggregate, the same way partial states from different workers are.
///
/// Only sorted (streaming) partial aggregates are pushed: they hold O(1) accumulators per
/// partition, while a hash aggregate holds O(num_groups) and would multiply its memory usage by
/// the partition count.
pub fn push_sorted_partial_aggregate_below_merge(
    p: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let Some(agg) = p.as_any().downcast_ref::<AggregateExec>() else {
        return Ok(p);
    };
    if *agg.mode() != AggregateMode::Partial
        || !matches!(agg.input_order_mode(), InputOrderMode::Sorted)
        // Restrict to aggregates convertible to InlineAggregateExec (no grouping sets)
        || !agg.group_expr().is_single()
    {
        return Ok(p);
    }
    let Some(merge) = agg
        .input()
        .as_any()
        .downcast_ref::<SortPreservingMergeExec>()
    else {
        return Ok(p);
    };
    if merge.fetch().is_some() {
        return Ok(p);
    }
    let merge_input = merge.input();
    if merge_input.output_partitioning().partition_count() <= 1 {
        return Ok(p);
    }

    let new_agg = AggregateExec::try_new(
        AggregateMode::Partial,
        agg.group_expr().clone(),
        agg.aggr_expr().to_vec(),
        agg.filter_expr().to_vec(),
        merge_input.clone(),
        agg.input_schema(),
    )?;
    // Per-partition input must still be sorted on the group keys, otherwise the aggregate
    // becomes hash-based and must stay above the merge.
    if !matches!(new_agg.input_order_mode(), InputOrderMode::Sorted) {
        return Ok(p);
    }
    let Some(ordering) = new_agg.properties().output_ordering().cloned() else {
        return Ok(p);
    };
    Ok(Arc::new(SortPreservingMergeExec::new(
        ordering,
        Arc::new(new_agg),
    )))
}

/// Bounds worker memory for `GROUP BY ... ORDER BY <group cols> LIMIT n` when the ORDER BY is a
/// subset of the group keys but not a prefix of the index sort key (so the cheaper
/// `add_limit_to_workers` early-termination doesn't apply and the worker would otherwise emit every
/// group). Driven by the `worker_sort_and_limit` descriptor computed in `ChooseIndex` and carried
/// through `ClusterSendNode` serialization, since the worker's plan has no ORDER BY of its own.
///
/// On the worker (a `WorkerExec` carrying the descriptor), the worker subtree
///     SortPreservingMerge(index) <- PartialAggregate
/// becomes
///     SortPreservingMerge(worker_order) <- Sort(worker_order, fetch=n, per partition) <- PartialAggregate
///
/// On the router (a final aggregate over a `ClusterSend` carrying the descriptor), the final
/// aggregate is rebuilt over a `SortPreservingMerge(worker_order)`: the bounded worker streams are
/// merged in `worker_order`, so equal group keys are adjacent and the (sorted) final aggregate
/// combines them. Its output is `worker_order`-sorted, whose prefix is the query's ORDER BY, so the
/// limit above it stays correct.
///
/// `worker_order` is the full group key reordered with the ORDER BY columns as its prefix. The
/// per-partition `fetch=n` is sound only because the key is the full group key: a group in the
/// global first `n` has fewer than `n` smaller-keyed groups, hence stays within every partition's
/// first `n`, so all of its partial states survive the fetch and the final aggregate sums it
/// correctly -- even when the same group spans partitions. Sorting by `order_by` alone (a partial
/// key) would drop partial states of tied groups.
pub fn push_worker_sort_and_limit(
    p: Arc<dyn ExecutionPlan>,
    group_by_limit_factor: usize,
    group_by_limit_per_partition: bool,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    // Worker side: bound the partial aggregate's output. A hash aggregate uses the trimming top-k
    // when factor > 0, otherwise (and for the sorted/inline aggregate, whose group count is unknown)
    // it is bounded with a per-partition Sort(fetch). `resort_worker_subtree` returns None only when
    // there is no locatable partial aggregate, leaving the plan as planned.
    if let Some(w) = p.as_any().downcast_ref::<WorkerExec>() {
        let Some((cols, fetch)) = w.worker_sort_and_limit.clone() else {
            return Ok(p);
        };
        let Some((new_input, _is_hash)) = resort_worker_subtree(
            &w.input,
            &cols,
            fetch,
            group_by_limit_factor,
            group_by_limit_per_partition,
        ) else {
            return Ok(p);
        };
        return Ok(Arc::new(WorkerExec::new(
            new_input,
            w.max_batch_rows,
            w.limit_and_reverse.clone(),
            w.required_input_ordering.clone(),
            w.worker_sort_and_limit.clone(),
            w.properties().output_partitioning().partition_count(),
        )));
    }

    // Router side: combine the workers' top-k with a hash final aggregate, then re-apply the top-k
    // sort by `worker_order` (the total order T).
    let Some(final_agg) = FinalAggregateInfo::extract(&p) else {
        return Ok(p);
    };
    let Some(cs_plan) = find_cluster_send(&final_agg.input) else {
        return Ok(p);
    };
    let cs = cs_plan
        .as_any()
        .downcast_ref::<ClusterSendExec>()
        .expect("find_cluster_send returns a ClusterSendExec");
    let Some((cols, fetch)) = cs.worker_sort_and_limit.clone() else {
        return Ok(p);
    };
    // `resort_worker_subtree` reports `is_hash` from the same descent it already does, so we don't
    // walk to the partial aggregate twice. The hash aggregate emits unordered trimmed groups
    // (combined by a hash final + re-sort); the sorted/inline aggregate emits bounded sorted streams
    // (combined by a sort-preserving merge + sorted final).
    let Some((new_worker_subtree, is_hash)) = resort_worker_subtree(
        &cs.input_for_optimizations,
        &cols,
        fetch,
        group_by_limit_factor,
        group_by_limit_per_partition,
    ) else {
        return Ok(p);
    };
    let new_cs: Arc<dyn ExecutionPlan> =
        Arc::new(cs.with_changed_schema(new_worker_subtree, cs.required_input_ordering.clone()));
    let worker_order = worker_ordering(&final_agg.group_expr, &cols)?;

    if is_hash {
        // Hash final over coalesced (unordered) streams, then re-apply the top-k sort by the total
        // order T. The Sort(fetch) is required even for a bare LIMIT: it keeps the k smallest by T --
        // exactly the groups every worker kept and fully combined here -- where a plain limit could
        // take a group only one worker kept (undercounted). The coalesce drains the workers in
        // parallel.
        let coalesced: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(new_cs));
        let final_hash: Arc<dyn ExecutionPlan> = Arc::new(AggregateExec::try_new(
            AggregateMode::Final,
            final_agg.group_expr,
            final_agg.aggr_expr,
            final_agg.filter_expr,
            coalesced,
            final_agg.input_schema,
        )?);
        Ok(Arc::new(
            SortExec::new(worker_order, final_hash).with_fetch(Some(fetch)),
        ))
    } else {
        // Sorted final over a sort-preserving merge in worker_order: equal keys are adjacent so the
        // sorted final combines them, and its output stays worker_order-sorted (whose prefix is the
        // query's ORDER BY), so the query's limit above stays correct -- no extra Sort needed.
        let merged: Arc<dyn ExecutionPlan> =
            Arc::new(SortPreservingMergeExec::new(worker_order, new_cs));
        Ok(Arc::new(AggregateExec::try_new(
            AggregateMode::Final,
            final_agg.group_expr,
            final_agg.aggr_expr,
            final_agg.filter_expr,
            merged,
            final_agg.input_schema,
        )?))
    }
}

/// The `SortOptions` a `worker_sort_and_limit` descriptor entry `(asc, nulls_first)` maps to.
fn descriptor_sort_options(asc: bool, nulls_first: bool) -> SortOptions {
    SortOptions {
        descending: !asc,
        nulls_first,
    }
}

/// Builds the `worker_order` LexOrdering over an aggregate's group columns from the descriptor.
fn worker_ordering(
    group_expr: &datafusion::physical_plan::aggregates::PhysicalGroupBy,
    cols: &[(usize, bool, bool)],
) -> Result<LexOrdering, DataFusionError> {
    let names = group_expr.expr();
    let mut exprs = Vec::with_capacity(cols.len());
    for (idx, asc, nulls_first) in cols {
        let (_, name) = names.get(*idx).ok_or_else(|| {
            DataFusionError::Internal("worker_sort_and_limit column out of range".to_string())
        })?;
        exprs.push(PhysicalSortExpr {
            expr: Arc::new(Column::new(name, *idx)),
            options: descriptor_sort_options(*asc, *nulls_first),
        });
    }
    Ok(LexOrdering::new(exprs))
}

/// Rebuilds a worker subtree to bound its output to the top `fetch` groups by the total order in
/// `cols`. Two shapes:
/// - hash (`AggregateExec`) with `factor > 0`: `CoalescePartitions <- GroupByLimitAggregate` --
///   trim during aggregation, emitted unsorted for the router's hash final.
/// - sorted/inline, and hash with `factor == 0`:
///   `SortPreservingMerge(T) <- Sort(T, fetch, per partition) <- PartialAggregate` -- bound each
///   partition with a sort. For hash this is the pre-trim behavior, kept as the `factor == 0`
///   fallback so disabling the trim still bounds the worker rather than leaving it unbounded.
///
/// Returns `None` for an unrecognized subtree (no locatable partial aggregate), otherwise the
/// rebuilt subtree paired with `is_hash` (true for the hash path) so the caller doesn't have to
/// walk down to the partial aggregate a second time to decide the router shape.
fn resort_worker_subtree(
    worker_subtree: &Arc<dyn ExecutionPlan>,
    cols: &[(usize, bool, bool)],
    fetch: usize,
    group_by_limit_factor: usize,
    group_by_limit_per_partition: bool,
) -> Option<(Arc<dyn ExecutionPlan>, bool)> {
    let partial = locate_partial_aggregate(worker_subtree)?;

    // Hash path: trim during aggregation, emit unsorted for the router's hash final. Engaged only
    // when factor > 0; with it off, the hash aggregate falls through to the sorted-bounding path
    // below (a per-partition Sort(fetch) + merge), which is the pre-trim behavior -- so disabling
    // the trim bounds the worker exactly as before, not leaving it unbounded.
    if partial.as_any().is::<AggregateExec>() && group_by_limit_factor > 0 {
        // CUBESTORE_GROUP_BY_LIMIT_PER_PARTITION decides where the worker's hash table lives:
        //  - off (default): coalesce the aggregate's input to a single partition, so the worker
        //    builds ONE hash table over the merged partitions ("over merge"). Peak memory ~k, no
        //    intra-worker parallelism on the aggregation.
        //  - on: keep the raw multi-partition input, so the aggregate runs once per partition
        //    ("under merge"): N parallel hash tables -- more parallelism, peak ~N*k. The trim keeps
        //    each per-partition output to ~factor*k, and the global top-k argument (full group key)
        //    guarantees a surviving group stays within every partition's local top-k.
        let partial = if group_by_limit_per_partition {
            partial
        } else {
            partial
                .as_any()
                .downcast_ref::<AggregateExec>()
                .and_then(|agg| {
                    if agg.input().output_partitioning().partition_count() <= 1 {
                        return None;
                    }
                    let merged: Arc<dyn ExecutionPlan> =
                        Arc::new(CoalescePartitionsExec::new(agg.input().clone()));
                    partial.clone().with_new_children(vec![merged]).ok()
                })
                .unwrap_or(partial)
        };

        let order: Vec<(usize, SortOptions)> = cols
            .iter()
            .map(|(idx, asc, nulls_first)| (*idx, descriptor_sort_options(*asc, *nulls_first)))
            .collect();
        let trimmed: Arc<dyn ExecutionPlan> = match partial.as_any().downcast_ref::<AggregateExec>()
        {
            Some(agg) => match GroupByLimitAggregateExec::try_new_from_partial(
                agg,
                fetch,
                group_by_limit_factor,
                order,
            ) {
                Some(e) => Arc::new(e),
                None => partial,
            },
            None => partial,
        };

        // Emit the trimmed top-k unsorted, coalesced to one stream for the ClusterSend. The router
        // hash-combines and re-applies the top-k sort, so no worker-side sort/merge holds the whole
        // result. Under "over merge" the trimmed agg is already single-partition and this coalesce is
        // a passthrough; under "under merge" it drains the per-partition aggregates concurrently.
        return Some((Arc::new(CoalescePartitionsExec::new(trimmed)), true));
    }

    // Sorted/inline path (and the factor == 0 fallback for hash aggregates): bound each partition
    // with Sort(fetch) and merge in the total order. The per-partition `fetch` is sound because the
    // key is the full group key: a globally top-`fetch` group stays within every partition's first
    // `fetch`, so the router's sorted final still sees all its partial states (see this function's
    // doc and the module note on the total order).
    let worker_order = lex_ordering_from_cols(cols, &partial.schema())?;
    let per_partition_sort: Arc<dyn ExecutionPlan> = Arc::new(
        SortExec::new(worker_order.clone(), partial)
            .with_fetch(Some(fetch))
            .with_preserve_partitioning(true),
    );
    Some((
        Arc::new(SortPreservingMergeExec::new(
            worker_order,
            per_partition_sort,
        )),
        false,
    ))
}

/// Build a `LexOrdering` over the partial aggregate's group columns from the descriptor (indices
/// into the partial output schema). Returns `None` if a column index is out of range.
fn lex_ordering_from_cols(
    cols: &[(usize, bool, bool)],
    schema: &datafusion::arrow::datatypes::SchemaRef,
) -> Option<LexOrdering> {
    let mut exprs = Vec::with_capacity(cols.len());
    for (idx, asc, nulls_first) in cols {
        let field = schema.fields().get(*idx)?;
        exprs.push(PhysicalSortExpr {
            expr: Arc::new(Column::new(field.name(), *idx)),
            options: descriptor_sort_options(*asc, *nulls_first),
        });
    }
    Some(LexOrdering::new(exprs))
}

/// The group/aggregate state of either an `InlineAggregateExec` or a plain `AggregateExec`, when in
/// Final mode.
struct FinalAggregateInfo {
    group_expr: datafusion::physical_plan::aggregates::PhysicalGroupBy,
    aggr_expr: Vec<Arc<datafusion::physical_expr::aggregate::AggregateFunctionExpr>>,
    filter_expr: Vec<Option<Arc<dyn PhysicalExpr>>>,
    input_schema: datafusion::arrow::datatypes::SchemaRef,
    input: Arc<dyn ExecutionPlan>,
}

impl FinalAggregateInfo {
    fn extract(p: &Arc<dyn ExecutionPlan>) -> Option<Self> {
        if let Some(a) = p.as_any().downcast_ref::<InlineAggregateExec>() {
            if *a.mode() != InlineAggregateMode::Final {
                return None;
            }
            Some(Self {
                group_expr: a.group_expr().clone(),
                aggr_expr: a.aggr_expr().to_vec(),
                filter_expr: a.filter_expr().to_vec(),
                input_schema: a.input_schema.clone(),
                input: a.input().clone(),
            })
        } else if let Some(a) = p.as_any().downcast_ref::<AggregateExec>() {
            if !matches!(
                a.mode(),
                AggregateMode::Final | AggregateMode::FinalPartitioned
            ) {
                return None;
            }
            Some(Self {
                group_expr: a.group_expr().clone(),
                aggr_expr: a.aggr_expr().to_vec(),
                filter_expr: a.filter_expr().to_vec(),
                input_schema: a.input_schema().clone(),
                input: a.input().clone(),
            })
        } else {
            None
        }
    }
}

/// Descends through merge/coalesce/check-memory wrappers to the ClusterSend feeding a final
/// aggregate.
fn find_cluster_send(p: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
    let any = p.as_any();
    if any.is::<ClusterSendExec>() {
        return Some(p.clone());
    }
    if any.is::<SortPreservingMergeExec>()
        || any.is::<CoalescePartitionsExec>()
        || any.is::<CheckMemoryExec>()
    {
        return find_cluster_send(p.children().into_iter().next()?);
    }
    None
}

/// The partial aggregate inside a worker subtree, reached through merge/coalesce wrappers. Returns
/// `None` for any other shape (including an already-rewritten subtree, whose partial now sits under
/// a `SortExec` -- not a wrapper we descend -- keeping the pass idempotent).
fn locate_partial_aggregate(p: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
    let mut candidate = p.clone();
    loop {
        let any = candidate.as_any();
        if let Some(a) = any.downcast_ref::<InlineAggregateExec>() {
            return (*a.mode() == InlineAggregateMode::Partial).then_some(candidate.clone());
        }
        if let Some(a) = any.downcast_ref::<AggregateExec>() {
            return (*a.mode() == AggregateMode::Partial).then_some(candidate.clone());
        }
        if any.is::<SortPreservingMergeExec>()
            || any.is::<CoalescePartitionsExec>()
            || any.is::<CheckMemoryExec>()
        {
            candidate = candidate.children().into_iter().next()?.clone();
            continue;
        }
        return None;
    }
}

pub fn ensure_partition_merge_helper(
    p: Arc<dyn ExecutionPlan>,
    new_child: &mut bool,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    if p.as_any().is::<ClusterSendExec>()
        || p.as_any().is::<WorkerExec>()
        || p.as_any().is::<UnionExec>()
    {
        let rewritten: Arc<dyn ExecutionPlan> = if let Some(ordering) = p.output_ordering() {
            let ordering = ordering.to_vec();
            let merged_children = p
                .children()
                .into_iter()
                .map(|c| -> Arc<dyn ExecutionPlan> {
                    Arc::new(SortPreservingMergeExec::new(
                        LexOrdering::new(ordering.clone()),
                        c.clone(),
                    ))
                })
                .collect();
            let new_plan = p.clone().with_new_children(merged_children)?;
            Arc::new(SortPreservingMergeExec::new(
                LexOrdering::new(ordering),
                new_plan,
            ))
        } else {
            let merged_children = p
                .children()
                .into_iter()
                .map(|c| -> Arc<dyn ExecutionPlan> {
                    Arc::new(CoalescePartitionsExec::new(c.clone()))
                })
                .collect();
            let new_plan = p.clone().with_new_children(merged_children)?;
            Arc::new(CoalescePartitionsExec::new(new_plan))
        };
        *new_child = true;
        Ok(rewritten)
    } else {
        Ok(p)
    }
}

pub fn ensure_partition_merge(
    p: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let mut new_child = false;
    ensure_partition_merge_helper(p, &mut new_child)
}

// TODO upgrade DF: this one was handled by something else but most likely only in sorted scenario
pub fn ensure_partition_merge_with_acceptable_parent(
    parent: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    // TODO upgrade DF: Figure out the right clean way to handle this function in general --
    // possibly involving uncommenting EnforceDistribution, and having this
    // SortPreservingMergeExec/CoalescePartitionsExec wrapping the ClusterSendExec node as we
    // construct the query.

    // Special case, don't do this inside AggregateTopKExec-ClusterSendExec-Aggregate because we
    // need the partitioning: (This is gross.)
    if parent.as_any().is::<AggregateTopKExec>() {
        return Ok(parent);
    }

    let mut any_new_children = false;
    let mut new_children = Vec::new();

    for p in parent.children() {
        new_children.push(ensure_partition_merge_helper(
            p.clone(),
            &mut any_new_children,
        )?);
    }
    if any_new_children {
        parent.with_new_children(new_children)
    } else {
        Ok(parent)
    }
}

/// Add `GlobalLimitExec` behind worker node if this node has `limit` property set and applies DF
/// `LimitPushdown` optimizer. Should be executed after all optimizations which can move `Worker`
/// node or change its input.  `config` is ignored -- we pass it to DF's `LimitPushdown` optimizer,
/// which also ignores it (as of DF 46.0.1).
pub fn add_limit_to_workers(
    p: Arc<dyn ExecutionPlan>,
    config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let limit_and_reverse;
    let input;
    if let Some(w) = p.as_any().downcast_ref::<WorkerExec>() {
        limit_and_reverse = w.limit_and_reverse;
        input = &w.input;
    } else if let Some(cs) = p.as_any().downcast_ref::<ClusterSendExec>() {
        limit_and_reverse = cs.limit_and_reverse;
        input = &cs.input_for_optimizations;
    } else {
        return Ok(p);
    }

    let Some((limit, reverse)) = limit_and_reverse else {
        return Ok(p);
    };

    // A row limit must not be placed above a merge of per-partition aggregates: the merged
    // stream carries duplicate group keys from different partitions, and cutting it by rows
    // could cut off part of some group's partial states, silently corrupting that group's
    // total. Instead the limit descends through the merges and lands directly above the first
    // aggregate, where one row is one complete (within its partition) group: the union of
    // per-partition first (last, for reverse) `limit` groups covers the global first (last)
    // `limit` groups, and groups beyond that arrive complete or not at all -- the router
    // orders by the group key, not by the totals, and its own limit drops them.
    if first_aggregate_below_merges(input).is_some() {
        let new_input = limit_above_first_aggregate(input, limit, reverse)?;
        return p.with_new_children(vec![new_input]);
    }

    // Tripwire for shapes the descent doesn't recognize: a per-partition partial aggregate in
    // the subtree means a duplicate-bearing merge somewhere above it, and the row limit below
    // (LimitPushdown descends through projections into merges) could land on it. The worker
    // limit is an optimization, skipping it is always correct.
    if contains_multi_partition_partial_aggregate(input) {
        return Ok(p);
    }

    // No aggregation: plain rows, each one self-contained, a row limit is exact anywhere.
    if reverse {
        let limit = Arc::new(TailLimitExec::new(input.clone(), limit));
        p.with_new_children(vec![limit])
    } else {
        let limit = Arc::new(GlobalLimitExec::new(input.clone(), 0, Some(limit)));
        let limit_optimized = LimitPushdown::new().optimize(limit, config)?;
        p.with_new_children(vec![limit_optimized])
    }
}

/// The first aggregate reachable from `p` looking only through sort preserving merges and
/// projections.
fn first_aggregate_below_merges(mut p: &Arc<dyn ExecutionPlan>) -> Option<&Arc<dyn ExecutionPlan>> {
    loop {
        if let Some(merge) = p.as_any().downcast_ref::<SortPreservingMergeExec>() {
            p = merge.input();
        } else if let Some(projection) = p.as_any().downcast_ref::<ProjectionExec>() {
            p = projection.input();
        } else if p.as_any().is::<InlineAggregateExec>() || p.as_any().is::<AggregateExec>() {
            return Some(p);
        } else {
            return None;
        }
    }
}

/// Rebuilds the chain of merges and projections with a per-partition row limit placed directly
/// above the first aggregate. Must only be called when [first_aggregate_below_merges] found one.
fn limit_above_first_aggregate(
    p: &Arc<dyn ExecutionPlan>,
    limit: usize,
    reverse: bool,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    if let Some(merge) = p.as_any().downcast_ref::<SortPreservingMergeExec>() {
        let child = limit_above_first_aggregate(merge.input(), limit, reverse)?;
        return Ok(Arc::new(
            SortPreservingMergeExec::new(merge.expr().clone(), child).with_fetch(merge.fetch()),
        ));
    }
    if let Some(projection) = p.as_any().downcast_ref::<ProjectionExec>() {
        // The limit is a plain row count, it commutes with column renames
        let child = limit_above_first_aggregate(projection.input(), limit, reverse)?;
        return Ok(Arc::new(ProjectionExec::try_new(
            projection.expr().to_vec(),
            child,
        )?));
    }

    // The sorted streaming aggregate can additionally stop reading its input early once it has
    // emitted `limit` groups; for reverse the last groups are unknown until the input ends, so
    // there is nothing to pass into the aggregate.
    let mut node = p.clone();
    if !reverse {
        if let Some(agg) = p.as_any().downcast_ref::<InlineAggregateExec>() {
            if *agg.mode() == InlineAggregateMode::Partial {
                let agg_limit = agg.limit().map_or(limit, |l| l.min(limit));
                node = Arc::new(agg.with_limit(Some(agg_limit)));
            }
        }
    }

    Ok(if reverse {
        Arc::new(TailLimitExec::new(node, limit))
    } else if node.output_partitioning().partition_count() == 1 {
        Arc::new(GlobalLimitExec::new(node, 0, Some(limit)))
    } else {
        Arc::new(LocalLimitExec::new(node, limit))
    })
}

/// True if the plan contains a partial aggregate executed per partition: its merged output
/// carries duplicate group keys, so a row limit above such a merge is not group-aligned.
fn contains_multi_partition_partial_aggregate(p: &Arc<dyn ExecutionPlan>) -> bool {
    let is_one = if let Some(agg) = p.as_any().downcast_ref::<InlineAggregateExec>() {
        *agg.mode() == InlineAggregateMode::Partial
    } else if let Some(agg) = p.as_any().downcast_ref::<AggregateExec>() {
        *agg.mode() == AggregateMode::Partial
    } else {
        false
    };
    if is_one && p.output_partitioning().partition_count() > 1 {
        return true;
    }
    p.children()
        .into_iter()
        .any(contains_multi_partition_partial_aggregate)
}

/// Because we disable `EnforceDistribution`, and because we add `SortPreservingMergeExec` in
/// `ensure_partition_merge_with_acceptable_parent` so that Sorted ("inplace") aggregates work
/// properly (which reduces memory usage), we in some cases have unnecessary
/// `SortPreservingMergeExec` nodes underneath a `Sort` node with a different ordering.  Or,
/// perhaps, we added a `GlobalLimitExec` by `add_limit_to_workers` and we can push down the limit
/// into a _matching_ `SortPreservingMergeExec` node.
///
/// A minor complication: There may be projection nodes in between that rename things.
pub fn replace_suboptimal_merge_sorts(
    p: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    if let Some(sort) = p.as_any().downcast_ref::<SortExec>() {
        if sort.preserve_partitioning() {
            // Let's not handle this.
            return Ok(p);
        }
        let required_ordering = p
            .output_ordering()
            .cloned()
            .map(LexRequirement::from)
            .unwrap_or_default();
        let new_input =
            replace_suboptimal_merge_sorts_helper(&required_ordering, sort.fetch(), sort.input())?;
        p.with_new_children(vec![new_input])
    } else {
        Ok(p)
    }
}

/// Replaces SortPreservingMergeExec in the subtree with either a CoalescePartitions (if it doesn't
/// match the ordering) or, if it does match the sort ordering, pushes down fetch information if
/// appropriate.
fn replace_suboptimal_merge_sorts_helper(
    required_ordering: &LexRequirement,
    fetch: Option<usize>,
    node: &Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    let node_any = node.as_any();
    if let Some(spm) = node_any.downcast_ref::<SortPreservingMergeExec>() {
        // A SortPreservingMergeExec that sort_exprs is a prefix of, is an acceptable ordering.  But
        // if there is no sort_exprs at all, we just use CoalescePartitions.
        if !required_ordering.is_empty() {
            let spm_req = LexRequirement::from(
                spm.properties()
                    .output_ordering()
                    .cloned()
                    .unwrap_or(LexOrdering::default()),
            );
            if !required_ordering.is_empty()
                && spm
                    .properties()
                    .eq_properties
                    .requirements_compatible(required_ordering, &spm_req)
            {
                // Okay, we have a matching SortPreservingMergeExec node!

                let mut new_fetch: Option<usize> = fetch;
                let new_spm = if let Some(fetch) = fetch {
                    if let Some(spm_fetch) = spm.fetch() {
                        if fetch < spm_fetch {
                            Arc::new(spm.clone().with_fetch(Some(fetch)))
                        } else {
                            // spm fetch is tighter.
                            new_fetch = Some(spm_fetch);
                            node.clone()
                        }
                    } else {
                        Arc::new(spm.clone().with_fetch(Some(fetch)))
                    }
                } else {
                    node.clone()
                };

                // Pass down spm's ordering, not sort_exprs, because we didn't touch spm besides the fetch..

                let new_input = replace_suboptimal_merge_sorts_helper(
                    &spm_req,
                    new_fetch,
                    new_spm
                        .children()
                        .first()
                        .ok_or(internal_datafusion_err!("no child"))?,
                )?;

                return new_spm.with_new_children(vec![new_input]);
            }
        }
        // sort_exprs is _not_ a prefix of spm.expr()
        // Aside: if spm.expr() is a prefix of sort_exprs, maybe SortExec could take advantage.

        // So it's not an acceptable ordering.  Create a CoalescePartitions, and remove other nested SortPreservingMergeExecs.
        let new_input = replace_suboptimal_merge_sorts_helper(
            &LexRequirement::new(vec![]),
            fetch,
            spm.input(),
        )?;

        return Ok(Arc::new(CoalescePartitionsExec::new(new_input)));
    } else if let Some(proj) = node_any.downcast_ref::<ProjectionExec>() {
        // TODO: Note that ProjectionExec has a TODO comment in DF's EnforceSorting optimizer (in sort_pushdown.rs).
        if let Some(new_sort_exprs) =
            sort_exprs_underneath_projection(required_ordering, proj.expr())?
        {
            let new_input =
                replace_suboptimal_merge_sorts_helper(&new_sort_exprs, fetch, proj.input())?;
            node.clone().with_new_children(vec![new_input])
        } else {
            Ok(node.clone())
        }
    } else if let Some(u) = node_any.downcast_ref::<UnionExec>() {
        let new_children: Result<Vec<_>, DataFusionError> = u
            .inputs()
            .iter()
            .map(|child| replace_suboptimal_merge_sorts_helper(required_ordering, fetch, child))
            .collect::<Result<Vec<_>, DataFusionError>>();
        let new_children = new_children?;
        Ok(Arc::new(UnionExec::new(new_children)))
    } else {
        Ok(node.clone())
    }
}

fn sort_exprs_underneath_projection(
    sort_exprs: &LexRequirement,
    proj_expr: &[(Arc<dyn PhysicalExpr>, String)],
) -> Result<Option<LexRequirement>, DataFusionError> {
    let mut sort_expr_columns = HashSet::<usize>::new();
    for expr in sort_exprs.iter() {
        record_columns_used(&mut sort_expr_columns, expr.expr.as_ref());
    }

    // sorted() just for determinism
    let sort_expr_columns: Vec<usize> = sort_expr_columns.into_iter().sorted().collect();
    let mut replacement_map =
        HashMap::<usize, datafusion::physical_plan::expressions::Column>::with_capacity(
            sort_expr_columns.len(),
        );

    for index in sort_expr_columns {
        let proj_lookup = proj_expr.get(index).ok_or_else(|| {
            DataFusionError::Internal(
                "proj_expr lookup in sort_exprs_underneath_projection failed".to_owned(),
            )
        })?;
        let Some(column_expr) = proj_lookup
            .0
            .as_any()
            .downcast_ref::<datafusion::physical_plan::expressions::Column>()
        else {
            return Ok(None);
        };
        replacement_map.insert(index, column_expr.clone());
    }

    // Now replace the columns in the sort_exprs with our different ones.
    let mut new_sort_exprs = Vec::with_capacity(sort_exprs.len());
    for e in sort_exprs.iter() {
        let transformed = replace_columns(&replacement_map, &e.expr)?;
        new_sort_exprs.push(PhysicalSortRequirement {
            expr: transformed,
            options: e.options,
        });
    }

    Ok(Some(LexRequirement::new(new_sort_exprs)))
}

fn record_columns_used(set: &mut HashSet<usize>, expr: &dyn PhysicalExpr) {
    if let Some(column) = expr
        .as_any()
        .downcast_ref::<datafusion::physical_plan::expressions::Column>()
    {
        set.insert(column.index());
    } else {
        for child in expr.children() {
            record_columns_used(set, child.as_ref());
        }
    }
}

fn replace_columns(
    replacement_map: &HashMap<usize, datafusion::physical_plan::expressions::Column>,
    expr: &Arc<dyn PhysicalExpr>,
) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
    Ok(
        TreeNode::transform(expr.clone(), |node: Arc<dyn PhysicalExpr>| {
            if let Some(column) = node
                .as_any()
                .downcast_ref::<datafusion::physical_plan::expressions::Column>()
            {
                let replacement = replacement_map.get(&column.index()).ok_or_else(|| {
                    DataFusionError::Internal("replace_columns has bad replacement_map".to_owned())
                })?;
                Ok(Transformed::yes(Arc::new(replacement.clone())))
            } else {
                Ok(Transformed::no(node))
            }
        })?
        .data,
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::arrow::array::{Int64Array, RecordBatch};
    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    use datafusion::functions_aggregate::sum::sum_udaf;
    use datafusion::physical_expr::aggregate::AggregateExprBuilder;
    use datafusion::physical_expr::expressions::col;
    use datafusion::physical_expr::PhysicalSortExpr;
    use datafusion::physical_plan::aggregates::PhysicalGroupBy;
    use datafusion::physical_plan::collect;
    use datafusion::prelude::SessionContext;
    use datafusion_datasource::memory::MemorySourceConfig;
    use datafusion_datasource::source::DataSourceExec;
    use std::collections::BTreeMap;

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

    fn make_batch(schema: &SchemaRef, rows: &[(i64, i64)]) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int64Array::from_iter_values(rows.iter().map(|r| r.0))),
                Arc::new(Int64Array::from_iter_values(rows.iter().map(|r| r.1))),
            ],
        )
        .unwrap()
    }

    /// Memory source with each partition sorted by `k`.
    fn sorted_source(
        schema: &SchemaRef,
        partitions: Vec<Vec<RecordBatch>>,
    ) -> Arc<dyn ExecutionPlan> {
        let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(
            col("k", schema).unwrap(),
        )]);
        let source = MemorySourceConfig::try_new(&partitions, schema.clone(), None)
            .unwrap()
            .try_with_sort_information(vec![ordering])
            .unwrap();
        Arc::new(DataSourceExec::new(Arc::new(source)))
    }

    fn merge_by_k(input: Arc<dyn ExecutionPlan>) -> Arc<SortPreservingMergeExec> {
        let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(
            col("k", &input.schema()).unwrap(),
        )]);
        Arc::new(SortPreservingMergeExec::new(ordering, input))
    }

    fn sum_aggregate(
        mode: AggregateMode,
        group_col: &str,
        input: Arc<dyn ExecutionPlan>,
    ) -> Arc<dyn ExecutionPlan> {
        let schema = input.schema();
        let group_by = PhysicalGroupBy::new_single(vec![(
            col(group_col, &schema).unwrap(),
            group_col.to_string(),
        )]);
        let sum = AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema).unwrap()])
            .schema(schema.clone())
            .alias("sum_v")
            .build()
            .unwrap();
        Arc::new(
            AggregateExec::try_new(
                mode,
                group_by,
                vec![Arc::new(sum)],
                vec![None],
                input,
                schema,
            )
            .unwrap(),
        )
    }

    /// Collects plan output into per-key sums, combining duplicate keys the way a Final
    /// aggregate would.
    async fn collect_summed(plan: Arc<dyn ExecutionPlan>) -> BTreeMap<i64, i64> {
        let session = SessionContext::new();
        let batches = collect(plan, session.task_ctx()).await.unwrap();
        let mut result = BTreeMap::new();
        for batch in batches {
            let keys = batch
                .column(0)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap();
            let sums = batch
                .column(1)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap();
            for (k, s) in keys.iter().zip(sums.iter()) {
                *result.entry(k.unwrap()).or_insert(0) += s.unwrap();
            }
        }
        result
    }

    fn two_partition_source(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
        // Duplicate keys 2 and 3 across partitions
        sorted_source(
            schema,
            vec![
                vec![make_batch(schema, &[(1, 10), (2, 20), (3, 30)])],
                vec![make_batch(schema, &[(2, 21), (3, 31), (4, 40)])],
            ],
        )
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn pushes_sorted_partial_aggregate_below_merge() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let original = sum_aggregate(AggregateMode::Partial, "k", merge_by_k(source));

        let rewritten = push_sorted_partial_aggregate_below_merge(original.clone()).unwrap();

        let merge = rewritten
            .as_any()
            .downcast_ref::<SortPreservingMergeExec>()
            .expect("merge must become the root");
        let agg = merge
            .input()
            .as_any()
            .downcast_ref::<AggregateExec>()
            .expect("partial aggregate must move below the merge");
        assert_eq!(*agg.mode(), AggregateMode::Partial);
        assert!(matches!(agg.input_order_mode(), InputOrderMode::Sorted));
        assert_eq!(
            agg.properties().output_partitioning().partition_count(),
            2,
            "aggregate must run per partition"
        );
        assert_eq!(rewritten.schema(), original.schema());

        // Cross-partition duplicate keys combine to the same totals as in the original plan
        assert_eq!(
            collect_summed(rewritten).await,
            collect_summed(original).await
        );
    }

    #[test]
    fn does_not_push_hash_aggregate() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("k", DataType::Int64, false),
            Field::new("g", DataType::Int64, false),
            Field::new("v", DataType::Int64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int64Array::from(vec![1, 2])),
                Arc::new(Int64Array::from(vec![5, 4])),
                Arc::new(Int64Array::from(vec![10, 20])),
            ],
        )
        .unwrap();
        let source = sorted_source(&schema, vec![vec![batch.clone()], vec![batch]]);
        // Grouping by `g` while the input is sorted by `k` makes the aggregate hash-based
        let original = sum_aggregate(AggregateMode::Partial, "g", merge_by_k(source));
        assert!(matches!(
            original
                .as_any()
                .downcast_ref::<AggregateExec>()
                .unwrap()
                .input_order_mode(),
            InputOrderMode::Linear
        ));

        let rewritten = push_sorted_partial_aggregate_below_merge(original.clone()).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    #[test]
    fn does_not_push_non_partial_aggregate() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let original = sum_aggregate(AggregateMode::Single, "k", merge_by_k(source));

        let rewritten = push_sorted_partial_aggregate_below_merge(original.clone()).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    #[test]
    fn does_not_push_below_merge_with_fetch() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let merge = Arc::new(merge_by_k(source).as_ref().clone().with_fetch(Some(3)));
        let original = sum_aggregate(AggregateMode::Partial, "k", merge);

        let rewritten = push_sorted_partial_aggregate_below_merge(original.clone()).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    #[test]
    fn does_not_push_below_single_partition_merge() {
        let schema = test_schema();
        let source = sorted_source(
            &schema,
            vec![vec![make_batch(&schema, &[(1, 10), (2, 20)])]],
        );
        let original = sum_aggregate(AggregateMode::Partial, "k", merge_by_k(source));

        let rewritten = push_sorted_partial_aggregate_below_merge(original.clone()).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    fn inline(plan: Arc<dyn ExecutionPlan>) -> Arc<InlineAggregateExec> {
        let agg = plan.as_any().downcast_ref::<AggregateExec>().unwrap();
        Arc::new(InlineAggregateExec::try_new_from_aggregate(agg).unwrap())
    }

    fn worker(
        input: Arc<dyn ExecutionPlan>,
        limit_and_reverse: Option<(usize, bool)>,
    ) -> Arc<dyn ExecutionPlan> {
        Arc::new(WorkerExec::new(
            input,
            4096,
            limit_and_reverse,
            None,
            None,
            1,
        ))
    }

    /// Worker plan with the partial aggregate below the merge: merge of per-partition partial
    /// states can contain duplicate group keys, so the row limit must be limit * partitions
    /// while the aggregates take the group limit.
    #[test]
    fn worker_limit_above_merged_partial_aggregate_limits_groups_per_partition() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let agg = inline(sum_aggregate(AggregateMode::Partial, "k", source));
        let merged = push_sorted_partial_aggregate_below_merge_shape(agg);
        let p = worker(merged, Some((3, false)));

        let rewritten = add_limit_to_workers(p, &ConfigOptions::default()).unwrap();

        let worker = rewritten.as_any().downcast_ref::<WorkerExec>().unwrap();
        let merge = worker
            .input
            .as_any()
            .downcast_ref::<SortPreservingMergeExec>()
            .expect("merge must stay on top of per-partition aggregates");
        assert_eq!(merge.fetch(), None);
        let local_limit = merge
            .input()
            .as_any()
            .downcast_ref::<LocalLimitExec>()
            .expect("the limit must become a per-partition row limit above the aggregate");
        assert_eq!(local_limit.fetch(), 3);
        let agg = local_limit
            .input()
            .as_any()
            .downcast_ref::<InlineAggregateExec>()
            .unwrap();
        assert_eq!(agg.limit(), Some(3));
    }

    /// The limit descends below the merge even for an aggregate that is not an
    /// InlineAggregateExec: only the early-stop absorption is specific to it, the row limit
    /// placement is not.
    #[test]
    fn worker_limit_above_merged_raw_partial_aggregate_stays_below_merge() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let agg = sum_aggregate(AggregateMode::Partial, "k", source);
        let merged = merge_by_k(agg);
        let p = worker(merged, Some((3, false)));

        let rewritten = add_limit_to_workers(p, &ConfigOptions::default()).unwrap();

        let worker_node = rewritten.as_any().downcast_ref::<WorkerExec>().unwrap();
        let merge = worker_node
            .input
            .as_any()
            .downcast_ref::<SortPreservingMergeExec>()
            .unwrap();
        assert_eq!(merge.fetch(), None);
        let local_limit = merge
            .input()
            .as_any()
            .downcast_ref::<LocalLimitExec>()
            .expect("a row limit above the merge would truncate duplicate group keys");
        assert_eq!(local_limit.fetch(), 3);
        assert!(local_limit.input().as_any().is::<AggregateExec>());

        let merged = merge_by_k(sum_aggregate(
            AggregateMode::Partial,
            "k",
            two_partition_source(&schema),
        ));
        let p = worker(merged, Some((3, true)));
        let rewritten = add_limit_to_workers(p, &ConfigOptions::default()).unwrap();
        let worker_node = rewritten.as_any().downcast_ref::<WorkerExec>().unwrap();
        let merge = worker_node
            .input
            .as_any()
            .downcast_ref::<SortPreservingMergeExec>()
            .unwrap();
        let tail = merge
            .input()
            .as_any()
            .downcast_ref::<TailLimitExec>()
            .unwrap();
        assert_eq!(tail.limit, 3);
        assert!(tail.input.as_any().is::<AggregateExec>());
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn worker_reverse_limit_above_merged_partial_aggregate_tails_each_partition() {
        let schema = test_schema();
        // Keys 4..=6 are present in both partitions; per-partition tails of 3 groups are
        // {4, 5, 6} and {7, 8, 9}.
        let source = sorted_source(
            &schema,
            vec![
                vec![make_batch(
                    &schema,
                    &(1..=6).map(|k| (k, k * 10)).collect::<Vec<_>>(),
                )],
                vec![make_batch(
                    &schema,
                    &(4..=9).map(|k| (k, k * 100 + 1)).collect::<Vec<_>>(),
                )],
            ],
        );
        let baseline = collect_summed(sum_aggregate(
            AggregateMode::Partial,
            "k",
            merge_by_k(source.clone()),
        ))
        .await;

        let agg = inline(sum_aggregate(AggregateMode::Partial, "k", source));
        let merged = push_sorted_partial_aggregate_below_merge_shape(agg);
        let p = worker(merged, Some((3, true)));

        let rewritten = add_limit_to_workers(p, &ConfigOptions::default()).unwrap();

        let worker = rewritten.as_any().downcast_ref::<WorkerExec>().unwrap();
        let merge = worker
            .input
            .as_any()
            .downcast_ref::<SortPreservingMergeExec>()
            .expect("merge must stay on top of per-partition tails");
        assert_eq!(merge.fetch(), None);
        let tail = merge
            .input()
            .as_any()
            .downcast_ref::<TailLimitExec>()
            .expect("reverse limit must become a per-partition tail below the merge");
        assert_eq!(tail.limit, 3);
        let agg = tail
            .input
            .as_any()
            .downcast_ref::<InlineAggregateExec>()
            .unwrap();
        assert_eq!(
            agg.limit(),
            None,
            "tail limit can not stop the aggregate early"
        );

        // The merged stream must carry complete totals for the last 3 group keys; earlier keys
        // may arrive partial and are dropped by the router's own limit.
        let summed = collect_summed(worker.input.clone()).await;
        for key in [7, 8, 9] {
            assert_eq!(summed[&key], baseline[&key], "complete total for key {key}");
        }
    }

    /// A global aggregate doesn't use the input order: the merge below it becomes a plain
    /// coalesce, through the filter, with identical results.
    #[tokio::test(flavor = "multi_thread")]
    async fn drops_sort_merge_under_global_aggregate() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let filter: Arc<dyn ExecutionPlan> = Arc::new(
            FilterExec::try_new(
                Arc::new(datafusion::physical_plan::expressions::BinaryExpr::new(
                    col("k", &schema).unwrap(),
                    datafusion::logical_expr::Operator::Gt,
                    Arc::new(datafusion::physical_plan::expressions::Literal::new(
                        datafusion::scalar::ScalarValue::Int64(Some(1)),
                    )),
                )),
                merge_by_k(source),
            )
            .unwrap(),
        );
        let original = global_sum_aggregate(filter);

        let rewritten = drop_sort_merge_under_global_aggregate(original.clone(), false).unwrap();

        let agg = rewritten.as_any().downcast_ref::<AggregateExec>().unwrap();
        let filter = agg
            .input()
            .as_any()
            .downcast_ref::<FilterExec>()
            .expect("filter must stay in place");
        assert!(
            filter.input().as_any().is::<CoalescePartitionsExec>(),
            "the merge must become a plain coalesce"
        );

        assert_eq!(
            collect_global_sum(rewritten).await,
            collect_global_sum(original).await
        );
    }

    /// An order-sensitive aggregate stays Linear with an empty GROUP BY but still needs its
    /// input ordered, so its merge stays.
    #[test]
    fn keeps_sort_merge_under_order_sensitive_aggregate() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let merged = merge_by_k(source);
        let first = AggregateExprBuilder::new(
            datafusion::functions_aggregate::array_agg::array_agg_udaf(),
            vec![col("v", &schema).unwrap()],
        )
        .schema(schema.clone())
        .alias("vals")
        .order_by(LexOrdering::new(vec![PhysicalSortExpr::new_default(
            col("k", &schema).unwrap(),
        )]))
        .build()
        .unwrap();
        let original: Arc<dyn ExecutionPlan> = Arc::new(
            AggregateExec::try_new(
                AggregateMode::Partial,
                PhysicalGroupBy::new_single(Vec::new()),
                vec![Arc::new(first)],
                vec![None],
                merged,
                schema,
            )
            .unwrap(),
        );
        assert!(original
            .as_any()
            .downcast_ref::<AggregateExec>()
            .unwrap()
            .required_input_ordering()[0]
            .is_some());

        let rewritten = drop_sort_merge_under_global_aggregate(original.clone(), false).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    /// A grouped hash aggregate over unmerged partitions could build a hash table per
    /// partition, so its merge stays.
    #[test]
    fn keeps_sort_merge_under_grouped_hash_aggregate() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("k", DataType::Int64, false),
            Field::new("g", DataType::Int64, false),
            Field::new("v", DataType::Int64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int64Array::from(vec![1, 2])),
                Arc::new(Int64Array::from(vec![5, 4])),
                Arc::new(Int64Array::from(vec![10, 20])),
            ],
        )
        .unwrap();
        let source = sorted_source(&schema, vec![vec![batch.clone()], vec![batch]]);
        let original = sum_aggregate(AggregateMode::Partial, "g", merge_by_k(source));
        assert!(matches!(
            original
                .as_any()
                .downcast_ref::<AggregateExec>()
                .unwrap()
                .input_order_mode(),
            InputOrderMode::Linear
        ));

        let rewritten = drop_sort_merge_under_global_aggregate(original.clone(), false).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    /// With the flag on, the same grouped hash aggregate's under-scan merge is replaced by plain
    /// partition coalescing (single output partition -> one hash table), dropping the useless sort.
    #[test]
    fn coalesces_sort_merge_under_grouped_hash_aggregate_when_enabled() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("k", DataType::Int64, false),
            Field::new("g", DataType::Int64, false),
            Field::new("v", DataType::Int64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int64Array::from(vec![1, 2])),
                Arc::new(Int64Array::from(vec![5, 4])),
                Arc::new(Int64Array::from(vec![10, 20])),
            ],
        )
        .unwrap();
        let source = sorted_source(&schema, vec![vec![batch.clone()], vec![batch]]);
        let original = sum_aggregate(AggregateMode::Partial, "g", merge_by_k(source));

        let rewritten = drop_sort_merge_under_global_aggregate(original.clone(), true).unwrap();
        // The aggregate stays grouped, but its input merge becomes a plain CoalescePartitionsExec
        // (still a single output partition -> one hash table).
        assert!(!Arc::ptr_eq(&rewritten, &original));
        assert!(rewritten.as_any().is::<AggregateExec>());
        assert!(rewritten.children()[0]
            .as_any()
            .is::<CoalescePartitionsExec>());
    }

    fn global_sum_aggregate(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        let schema = input.schema();
        let sum = AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema).unwrap()])
            .schema(schema.clone())
            .alias("sum_v")
            .build()
            .unwrap();
        Arc::new(
            AggregateExec::try_new(
                AggregateMode::Partial,
                PhysicalGroupBy::new_single(Vec::new()),
                vec![Arc::new(sum)],
                vec![None],
                input,
                schema,
            )
            .unwrap(),
        )
    }

    /// Sums the partial states of a global aggregate across partitions.
    async fn collect_global_sum(plan: Arc<dyn ExecutionPlan>) -> i64 {
        let session = SessionContext::new();
        let batches = collect(plan, session.task_ctx()).await.unwrap();
        batches
            .iter()
            .flat_map(|b| {
                b.column(0)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .iter()
                    .flatten()
                    .collect::<Vec<_>>()
            })
            .sum()
    }

    #[test]
    fn keeps_sort_merge_under_sorted_aggregate() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let original = sum_aggregate(AggregateMode::Partial, "k", merge_by_k(source));

        let rewritten = drop_sort_merge_under_global_aggregate(original.clone(), false).unwrap();
        assert!(Arc::ptr_eq(&rewritten, &original));
    }

    /// The descent looks through projections: the row limit is a plain count and commutes with
    /// column renames.
    #[test]
    fn worker_limit_descends_through_projection() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let agg = inline(sum_aggregate(AggregateMode::Partial, "k", source));
        let merged = push_sorted_partial_aggregate_below_merge_shape(agg);
        let projection = Arc::new(
            ProjectionExec::try_new(
                merged
                    .schema()
                    .fields()
                    .iter()
                    .map(|f| (col(f.name(), &merged.schema()).unwrap(), f.name().clone()))
                    .collect(),
                merged,
            )
            .unwrap(),
        );
        let p = worker(projection, Some((3, false)));

        let rewritten = add_limit_to_workers(p, &ConfigOptions::default()).unwrap();

        let worker_node = rewritten.as_any().downcast_ref::<WorkerExec>().unwrap();
        let projection = worker_node
            .input
            .as_any()
            .downcast_ref::<ProjectionExec>()
            .expect("projection must stay on top");
        let merge = projection
            .input()
            .as_any()
            .downcast_ref::<SortPreservingMergeExec>()
            .unwrap();
        let local_limit = merge
            .input()
            .as_any()
            .downcast_ref::<LocalLimitExec>()
            .expect("the limit must descend through the projection to the aggregate");
        assert_eq!(local_limit.fetch(), 3);
        let agg = local_limit
            .input()
            .as_any()
            .downcast_ref::<InlineAggregateExec>()
            .unwrap();
        assert_eq!(agg.limit(), Some(3));
    }

    /// A shape the descent doesn't recognize but containing per-partition partial aggregates
    /// must not get a worker limit at all: a row limit could land above a duplicate-bearing
    /// merge.
    #[test]
    fn worker_limit_skipped_for_unrecognized_shape_with_per_partition_aggregate() {
        let schema = test_schema();
        let source = two_partition_source(&schema);
        let agg = inline(sum_aggregate(AggregateMode::Partial, "k", source));
        let coalesce: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(agg));
        let p = worker(coalesce, Some((3, false)));

        let rewritten = add_limit_to_workers(p.clone(), &ConfigOptions::default()).unwrap();
        assert!(
            Arc::ptr_eq(&rewritten, &p),
            "no limit must be added to an unrecognized shape with a per-partition aggregate"
        );
    }

    /// Single partition partial aggregate emits unique group keys, the row limit stays exact
    /// and also lets the aggregate stop early.
    #[test]
    fn worker_limit_above_single_partition_partial_aggregate_sets_aggregate_limit() {
        let schema = test_schema();
        let source = sorted_source(
            &schema,
            vec![vec![make_batch(&schema, &[(1, 10), (2, 20)])]],
        );
        let agg = inline(sum_aggregate(AggregateMode::Partial, "k", source));
        let p = worker(agg, Some((3, false)));

        let rewritten = add_limit_to_workers(p, &ConfigOptions::default()).unwrap();

        let worker = rewritten.as_any().downcast_ref::<WorkerExec>().unwrap();
        let limit = worker
            .input
            .as_any()
            .downcast_ref::<GlobalLimitExec>()
            .unwrap();
        assert_eq!(limit.fetch(), Some(3));
        let agg = limit
            .input()
            .as_any()
            .downcast_ref::<InlineAggregateExec>()
            .unwrap();
        assert_eq!(agg.limit(), Some(3));
    }

    /// Builds the post-rewrite shape merge-above-aggregate for an already converted
    /// InlineAggregateExec.
    fn push_sorted_partial_aggregate_below_merge_shape(
        agg: Arc<InlineAggregateExec>,
    ) -> Arc<dyn ExecutionPlan> {
        let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(
            col("k", &agg.schema()).unwrap(),
        )]);
        Arc::new(SortPreservingMergeExec::new(ordering, agg))
    }
}
