query/optimizer/

global_limit.rs

// Copyright 2023 Greptime Team
//
// Licensed 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.

use std::sync::Arc;

use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::filter::FilterExec;
use datafusion::physical_plan::limit::GlobalLimitExec;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{
    ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions,
};
use datafusion_common::Result as DfResult;
use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning};

use crate::dist_plan::MergeSortExec;

#[derive(Debug)]
pub struct EnsureGlobalLimitForFetch;

impl PhysicalOptimizerRule for EnsureGlobalLimitForFetch {
    fn optimize(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        _config: &ConfigOptions,
    ) -> DfResult<Arc<dyn ExecutionPlan>> {
        Self::optimize_plan(plan, ParentContext::default())
    }

    fn name(&self) -> &str {
        "EnsureGlobalLimitForFetch"
    }

    fn schema_check(&self) -> bool {
        true
    }
}

impl EnsureGlobalLimitForFetch {
    fn optimize_plan(
        plan: Arc<dyn ExecutionPlan>,
        parent: ParentContext,
    ) -> DfResult<Arc<dyn ExecutionPlan>> {
        let children = plan.children();
        let plan = if children.is_empty() {
            plan
        } else {
            let required_input_distribution = plan.input_distribution_requirements();
            let required_input_ordering = plan.required_input_ordering();
            let maintains_input_order = plan.maintains_input_order();
            let child_parent = ParentContext {
                global_fetch: provided_global_fetch(&plan),
                required_ordering: None,
                required_distribution: Distribution::UnspecifiedDistribution,
                partitioning_to_restore: None,
                preserve_hash_partitioning: false,
            };
            let children = children
                .into_iter()
                .enumerate()
                .map(|(idx, child)| {
                    let required_distribution = required_input_distribution
                        .child_distribution(idx)
                        .cloned()
                        .unwrap_or(Distribution::UnspecifiedDistribution);
                    let partitioning_to_restore =
                        partitioning_to_restore_for(child, &required_distribution)
                            .or_else(|| inherited_partitioning_to_restore(&plan, child, &parent));
                    let preserve_hash_partitioning = partitioning_to_restore.is_some();
                    let required_ordering = required_input_ordering
                        .get(idx)
                        .cloned()
                        .unwrap_or(None)
                        .or_else(|| {
                            maintains_input_order
                                .get(idx)
                                .copied()
                                .unwrap_or(false)
                                .then(|| parent.required_ordering.clone())
                                .flatten()
                        });
                    let parent = ParentContext {
                        required_ordering,
                        required_distribution,
                        partitioning_to_restore,
                        preserve_hash_partitioning,
                        ..child_parent.clone()
                    };
                    Self::optimize_plan(Arc::clone(child), parent)
                })
                .collect::<DfResult<Vec<_>>>()?;
            plan.replace_children(
                children,
                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
            )?
        };

        let Some(fetch) = plan.fetch() else {
            return Ok(plan);
        };

        if parent
            .global_fetch
            .is_some_and(|parent_fetch| parent_fetch <= fetch)
            || !plan.is::<FilterExec>()
            || plan.output_partitioning().partition_count() <= 1
        {
            return Ok(plan);
        }

        add_global_fetch(
            plan,
            fetch,
            parent.required_ordering,
            parent.partitioning_to_restore,
        )
    }
}

#[derive(Clone)]
struct ParentContext {
    global_fetch: Option<usize>,
    required_ordering: Option<OrderingRequirements>,
    required_distribution: Distribution,
    partitioning_to_restore: Option<Partitioning>,
    preserve_hash_partitioning: bool,
}

impl Default for ParentContext {
    fn default() -> Self {
        Self {
            global_fetch: None,
            required_ordering: None,
            required_distribution: Distribution::UnspecifiedDistribution,
            partitioning_to_restore: None,
            preserve_hash_partitioning: false,
        }
    }
}

fn provided_global_fetch(plan: &Arc<dyn ExecutionPlan>) -> Option<usize> {
    let fetch = plan.fetch()?;
    (plan.is::<GlobalLimitExec>()
        || plan.is::<CoalescePartitionsExec>()
        || plan.is::<SortPreservingMergeExec>()
        || plan.is::<MergeSortExec>())
    .then_some(fetch)
}

fn add_global_fetch(
    plan: Arc<dyn ExecutionPlan>,
    fetch: usize,
    required_ordering: Option<OrderingRequirements>,
    partitioning_to_restore: Option<Partitioning>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
    let plan = if required_ordering.is_some()
        && let Some(ordering) = plan.output_ordering().cloned()
    {
        Arc::new(SortPreservingMergeExec::new(ordering, plan).with_fetch(Some(fetch)))
            as Arc<dyn ExecutionPlan>
    } else {
        Arc::new(CoalescePartitionsExec::new(plan).with_fetch(Some(fetch)))
            as Arc<dyn ExecutionPlan>
    };

    restore_required_partitioning(plan, partitioning_to_restore)
}

fn restore_required_partitioning(
    plan: Arc<dyn ExecutionPlan>,
    partitioning_to_restore: Option<Partitioning>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
    let Some(partitioning) = partitioning_to_restore else {
        return Ok(plan);
    };

    if partitioning.partition_count() <= 1 || !matches!(&partitioning, Partitioning::Hash(_, _)) {
        return Ok(plan);
    }

    Ok(Arc::new(
        RepartitionExec::try_new(plan, partitioning)?.with_preserve_order(),
    ))
}

fn partitioning_to_restore_for(
    child: &Arc<dyn ExecutionPlan>,
    required_distribution: &Distribution,
) -> Option<Partitioning> {
    if !matches!(required_distribution, Distribution::KeyPartitioned(_))
        || child.output_partitioning().partition_count() <= 1
    {
        return None;
    }

    if child
        .output_partitioning()
        .satisfaction(required_distribution, child.equivalence_properties(), false)
        .is_satisfied()
    {
        Some(child.output_partitioning().clone())
    } else {
        Some(
            required_distribution
                .clone()
                .create_partitioning(child.output_partitioning().partition_count()),
        )
    }
}

fn inherited_partitioning_to_restore(
    plan: &Arc<dyn ExecutionPlan>,
    child: &Arc<dyn ExecutionPlan>,
    parent: &ParentContext,
) -> Option<Partitioning> {
    if child.output_partitioning().partition_count() <= 1
        || !matches!(child.output_partitioning(), Partitioning::Hash(_, _))
        || !matches!(plan.output_partitioning(), Partitioning::Hash(_, _))
        || plan.output_partitioning().partition_count()
            != child.output_partitioning().partition_count()
    {
        return None;
    }

    let satisfies_parent_distribution = matches!(
        parent.required_distribution,
        Distribution::KeyPartitioned(_)
    ) && plan
        .output_partitioning()
        .satisfaction(
            &parent.required_distribution,
            plan.equivalence_properties(),
            false,
        )
        .is_satisfied();

    (satisfies_parent_distribution || parent.preserve_hash_partitioning)
        .then(|| child.output_partitioning().clone())
}

#[cfg(test)]
mod tests {
    use datafusion::arrow::array::{Array, Int32Array};
    use datafusion::arrow::compute::SortOptions;
    use datafusion::arrow::datatypes::{DataType, Field, Schema};
    use datafusion::arrow::record_batch::RecordBatch;
    use datafusion::execution::TaskContext;
    use datafusion::physical_expr::expressions::{col, lit};
    use datafusion::physical_optimizer::optimizer::PhysicalOptimizer;
    use datafusion::physical_plan::aggregates::{
        AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy,
    };
    use datafusion::physical_plan::filter::FilterExecBuilder;
    use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
    use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
    use datafusion::physical_plan::projection::ProjectionExec;
    use datafusion::physical_plan::repartition::RepartitionExec;
    use datafusion::physical_plan::sorts::sort::SortExec;
    use datafusion::physical_plan::test::TestMemoryExec;
    use datafusion_common::{JoinType, NullEquality};
    use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr};

    use super::*;

    async fn optimize_and_collect_twice(
        mut plan: Arc<dyn ExecutionPlan>,
        config: &ConfigOptions,
    ) -> Vec<Vec<i32>> {
        let mut results = Vec::with_capacity(2);
        for _ in 0..2 {
            for rule in PhysicalOptimizer::new().rules {
                plan = rule.optimize(plan, config).unwrap();
            }
            let batches = datafusion::physical_plan::collect(
                Arc::clone(&plan),
                Arc::new(TaskContext::default()),
            )
            .await
            .unwrap();
            results.push(
                batches
                    .iter()
                    .flat_map(|batch| {
                        batch
                            .column(0)
                            .as_any()
                            .downcast_ref::<Int32Array>()
                            .unwrap()
                            .values()
                            .to_vec()
                    })
                    .collect(),
            );
        }
        results
    }

    #[tokio::test]
    async fn physical_optimizer_keeps_global_distinct_limit_across_two_passes() {
        for soft_limit in [false, true] {
            let mut config = ConfigOptions::new();
            config.execution.target_partitions = 3;
            config.optimizer.enable_distinct_aggregation_soft_limit = soft_limit;

            for mode in [
                AggregateMode::FinalPartitioned,
                AggregateMode::SinglePartitioned,
            ] {
                // The query limit is the only initial limit.
                let input = input_with_all_hash_partitions();
                let aggregate = match mode {
                    AggregateMode::FinalPartitioned => {
                        agg(agg(input, AggregateMode::Partial), mode)
                    }
                    AggregateMode::SinglePartitioned => agg(hash_repartition(input), mode),
                    _ => unreachable!(),
                };
                let plan =
                    Arc::new(GlobalLimitExec::new(aggregate, 0, Some(1))) as Arc<dyn ExecutionPlan>;

                let results = optimize_and_collect_twice(plan, &config).await;
                assert_eq!(
                    results.iter().map(Vec::len).collect::<Vec<_>>(),
                    vec![1, 1],
                    "soft limit enabled: {soft_limit}, mode: {mode:?}",
                );
            }
        }
    }

    #[tokio::test]
    async fn physical_optimizer_keeps_count_over_distinct_limit_across_two_passes() {
        use datafusion::datasource::MemTable;
        use datafusion::execution::context::{SessionConfig, SessionContext as DFSessionContext};
        use datafusion_common::ScalarValue;

        for soft_limit in [false, true] {
            let mut config = SessionConfig::new().with_target_partitions(3);
            config
                .options_mut()
                .optimizer
                .enable_distinct_aggregation_soft_limit = soft_limit;
            let ctx = DFSessionContext::new_with_config(config.clone());
            let schema = schema();
            let batch = RecordBatch::try_new(
                schema.clone(),
                vec![Arc::new(Int32Array::from((0..100).collect::<Vec<_>>()))],
            )
            .unwrap();
            let table = MemTable::try_new(schema, vec![vec![batch]; 3]).unwrap();
            ctx.register_table("t", Arc::new(table)).unwrap();
            let mut plan = ctx
                .sql("SELECT COUNT(*) FROM (SELECT DISTINCT a FROM t LIMIT 1) AS limited")
                .await
                .unwrap()
                .create_physical_plan()
                .await
                .unwrap();

            for pass in 0..=2 {
                if pass > 0 {
                    for rule in PhysicalOptimizer::new().rules {
                        plan = rule.optimize(plan, config.options()).unwrap();
                    }
                }
                let batches = datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx())
                    .await
                    .unwrap();
                let values = batches
                    .iter()
                    .flat_map(|batch| {
                        (0..batch.num_rows()).map(|row| {
                            ScalarValue::try_from_array(batch.column(0).as_ref(), row).unwrap()
                        })
                    })
                    .collect::<Vec<_>>();
                assert_eq!(
                    values,
                    vec![ScalarValue::Int64(Some(1))],
                    "soft limit enabled: {soft_limit}, additional optimizer passes: {pass}",
                );
            }
        }
    }

    #[tokio::test]
    async fn physical_optimizer_keeps_local_limit_per_partition_across_two_passes() {
        let mut config = ConfigOptions::new();
        config.execution.target_partitions = 3;
        let input = input_with_all_hash_partitions();
        let local_limit = Arc::new(LocalLimitExec::new(
            agg(hash_repartition(input), AggregateMode::SinglePartitioned),
            1,
        )) as Arc<dyn ExecutionPlan>;

        let results = optimize_and_collect_twice(local_limit, &config).await;
        assert_eq!(results.iter().map(Vec::len).collect::<Vec<_>>(), vec![3, 3]);
    }

    #[tokio::test]
    async fn physical_optimizer_keeps_ordered_topk_across_two_passes() {
        let schema = schema();
        let losing =
            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![100]))])
                .unwrap();
        let winning =
            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))])
                .unwrap();
        let input = Arc::new(
            TestMemoryExec::try_new(&[vec![losing], vec![winning]], schema.clone(), None).unwrap(),
        );
        let ordering = ordering(schema.as_ref(), false);
        let topk = Arc::new(
            SortExec::new(
                ordering,
                agg_with_limit_options(
                    input,
                    AggregateMode::FinalPartitioned,
                    Some(LimitOptions::new_with_order(1, false)),
                ),
            )
            .with_fetch(Some(1)),
        ) as Arc<dyn ExecutionPlan>;

        let mut config = ConfigOptions::new();
        config.execution.target_partitions = 3;
        assert_eq!(
            optimize_and_collect_twice(topk, &config).await,
            vec![vec![1], vec![1]],
        );
    }

    #[tokio::test]
    async fn physical_optimizer_keeps_ordered_offset_limit_across_two_passes() {
        let mut config = ConfigOptions::new();
        config.execution.target_partitions = 3;
        let offset = Arc::new(GlobalLimitExec::new(
            Arc::new(SortExec::new(
                ordering(schema().as_ref(), false),
                unordered_input(),
            )),
            3,
            Some(1),
        )) as Arc<dyn ExecutionPlan>;

        assert_eq!(
            optimize_and_collect_twice(offset, &config).await,
            vec![vec![2], vec![2]],
        );
    }

    #[test]
    fn adds_global_limit_for_multi_partition_filter_fetch() {
        let filter = filter_fetch(unordered_input(), 1);

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(filter, ParentContext::default()).unwrap();

        assert!(optimized.is::<CoalescePartitionsExec>());
        assert_eq!(optimized.fetch(), Some(1));
        assert_eq!(optimized.output_partitioning().partition_count(), 1);
    }

    #[test]
    fn still_visits_subtree_under_global_limit() {
        let filter = filter_fetch(unordered_input(), 5);
        let projection = Arc::new(
            ProjectionExec::try_new(
                vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())],
                filter,
            )
            .unwrap(),
        );
        let limit =
            Arc::new(GlobalLimitExec::new(projection, 0, Some(10))) as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(limit, ParentContext::default()).unwrap();
        let projection = optimized.children()[0];
        let coalesce = projection.children()[0];

        assert!(coalesce.is::<CoalescePartitionsExec>());
        assert_eq!(coalesce.fetch(), Some(5));
    }

    #[test]
    fn keeps_filter_under_parent_global_fetch() {
        let (input, ordering) = ordered_input();
        let filter = filter_fetch(input, 1);
        let merge = Arc::new(SortPreservingMergeExec::new(ordering, filter).with_fetch(Some(1)))
            as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
        let child = optimized.children()[0];

        assert!(optimized.is::<SortPreservingMergeExec>());
        assert!(child.is::<FilterExec>());
    }

    #[test]
    fn adds_tighter_global_fetch_under_looser_parent_fetch() {
        let (input, ordering) = ordered_input();
        let filter = filter_fetch(input, 5);
        let merge = Arc::new(SortPreservingMergeExec::new(ordering, filter).with_fetch(Some(10)))
            as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
        let child = optimized.children()[0];

        assert!(optimized.is::<SortPreservingMergeExec>());
        assert!(child.is::<SortPreservingMergeExec>());
        assert_eq!(child.fetch(), Some(5));
    }

    #[test]
    fn keeps_filter_under_parent_merge_sort_fetch() {
        let (input, ordering) = ordered_input();
        let filter = filter_fetch(input, 1);
        let merge = merge_sort_fetch(ordering, filter, 1);

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
        let child = optimized.children()[0];

        assert!(optimized.is::<MergeSortExec>());
        assert!(child.is::<FilterExec>());
    }

    #[test]
    fn adds_tighter_global_fetch_under_looser_merge_sort_fetch() {
        let (input, ordering) = ordered_input();
        let filter = filter_fetch(input, 5);
        let merge = merge_sort_fetch(ordering, filter, 10);

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
        let child = optimized.children()[0];

        assert!(optimized.is::<MergeSortExec>());
        assert!(child.is::<SortPreservingMergeExec>());
        assert_eq!(child.fetch(), Some(5));
        assert!(child.children()[0].is::<FilterExec>());
    }

    #[test]
    fn preserves_parent_ordering_requirement() {
        let (input, ordering) = ordered_input();
        let filter = filter_fetch(input, 1);
        let merge =
            Arc::new(SortPreservingMergeExec::new(ordering, filter)) as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
        let child = optimized.children()[0];

        assert!(optimized.is::<SortPreservingMergeExec>());
        assert!(child.is::<SortPreservingMergeExec>());
        assert_eq!(child.fetch(), Some(1));
    }

    #[test]
    fn uses_child_output_ordering_for_merge() {
        let schema = schema();
        let required_ordering = ordering(schema.as_ref(), false);
        let actual_ordering = ordering(schema.as_ref(), true);
        let batch = batch(schema.clone());
        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
        let input = TestMemoryExec::try_new(&partitions, schema, None)
            .unwrap()
            .try_with_sort_information(vec![actual_ordering.clone()])
            .unwrap();
        let filter = filter_fetch(Arc::new(input), 1);

        let optimized = add_global_fetch(
            filter,
            1,
            Some(OrderingRequirements::from(required_ordering)),
            None,
        )
        .unwrap();
        let merge = optimized.downcast_ref::<SortPreservingMergeExec>().unwrap();

        assert_eq!(merge.expr(), &actual_ordering);
    }

    #[test]
    fn preserves_inherited_ordering_requirement_through_projection() {
        let (input, ordering) = ordered_input();
        let filter = filter_fetch(input, 1);
        let projection = Arc::new(
            ProjectionExec::try_new(
                vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())],
                filter,
            )
            .unwrap(),
        );
        let merge =
            Arc::new(SortPreservingMergeExec::new(ordering, projection)) as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
        let projection = optimized.children()[0];
        let child = projection.children()[0];

        assert!(optimized.is::<SortPreservingMergeExec>());
        assert!(projection.is::<ProjectionExec>());
        assert!(child.is::<SortPreservingMergeExec>());
        assert_eq!(child.fetch(), Some(1));
    }

    #[test]
    fn restores_parent_hash_distribution_after_global_fetch() {
        let left = filter_fetch(hash_repartition(unordered_input()), 1);
        let right = hash_repartition(unordered_input());
        let on = vec![(
            col("a", left.schema().as_ref()).unwrap(),
            col("a", right.schema().as_ref()).unwrap(),
        )];
        let join = Arc::new(
            HashJoinExec::try_new(
                left,
                right,
                on,
                None,
                &JoinType::Inner,
                None,
                PartitionMode::Partitioned,
                NullEquality::NullEqualsNothing,
                false,
            )
            .unwrap(),
        ) as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
        let left = optimized.children()[0];
        let repartition = left.downcast_ref::<RepartitionExec>().unwrap();

        assert!(matches!(
            repartition.partitioning(),
            Partitioning::Hash(_, 3)
        ));
        assert!(repartition.input().is::<CoalescePartitionsExec>());
        assert_eq!(repartition.input().fetch(), Some(1));
    }

    #[test]
    fn restores_inherited_hash_distribution_through_projection() {
        let filter = filter_fetch(hash_repartition(unordered_input()), 1);
        let projection = Arc::new(
            ProjectionExec::try_new(
                vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())],
                filter,
            )
            .unwrap(),
        ) as Arc<dyn ExecutionPlan>;
        let right = hash_repartition(unordered_input());
        let on = vec![(
            col("a", projection.schema().as_ref()).unwrap(),
            col("a", right.schema().as_ref()).unwrap(),
        )];
        let join = Arc::new(
            HashJoinExec::try_new(
                projection,
                right,
                on,
                None,
                &JoinType::Inner,
                None,
                PartitionMode::Partitioned,
                NullEquality::NullEqualsNothing,
                false,
            )
            .unwrap(),
        ) as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
        let projection = optimized.children()[0];
        let repartition = projection.children()[0]
            .downcast_ref::<RepartitionExec>()
            .unwrap();

        assert!(projection.is::<ProjectionExec>());
        assert!(matches!(
            repartition.partitioning(),
            Partitioning::Hash(_, 3)
        ));
        assert!(repartition.input().is::<CoalescePartitionsExec>());
        assert_eq!(repartition.input().fetch(), Some(1));
    }

    #[test]
    fn restores_inherited_hash_distribution_through_multiple_projections() {
        let filter = filter_fetch(hash_repartition(unordered_input()), 1);
        let projection = project_a(filter);
        let projection = project_a(projection);
        let right = hash_repartition(unordered_input());
        let on = vec![(
            col("a", projection.schema().as_ref()).unwrap(),
            col("a", right.schema().as_ref()).unwrap(),
        )];
        let join = Arc::new(
            HashJoinExec::try_new(
                projection,
                right,
                on,
                None,
                &JoinType::Inner,
                None,
                PartitionMode::Partitioned,
                NullEquality::NullEqualsNothing,
                false,
            )
            .unwrap(),
        ) as Arc<dyn ExecutionPlan>;

        let optimized =
            EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
        let outer_projection = optimized.children()[0];
        let inner_projection = outer_projection.children()[0];
        let repartition = inner_projection.children()[0]
            .downcast_ref::<RepartitionExec>()
            .unwrap();

        assert!(outer_projection.is::<ProjectionExec>());
        assert!(inner_projection.is::<ProjectionExec>());
        assert!(matches!(
            repartition.partitioning(),
            Partitioning::Hash(_, 3)
        ));
        assert!(repartition.input().is::<CoalescePartitionsExec>());
        assert_eq!(repartition.input().fetch(), Some(1));
    }

    fn unordered_input() -> Arc<dyn ExecutionPlan> {
        let schema = schema();
        let batch = batch(schema.clone());
        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
        Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap())
    }

    fn input_with_all_hash_partitions() -> Arc<dyn ExecutionPlan> {
        let schema = schema();
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(Int32Array::from((0..100).collect::<Vec<_>>()))],
        )
        .unwrap();
        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
        Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap())
    }

    fn agg(input: Arc<dyn ExecutionPlan>, mode: AggregateMode) -> Arc<dyn ExecutionPlan> {
        agg_with_limit_options(input, mode, None)
    }

    fn agg_with_limit_options(
        input: Arc<dyn ExecutionPlan>,
        mode: AggregateMode,
        limit_options: Option<LimitOptions>,
    ) -> Arc<dyn ExecutionPlan> {
        let schema = input.schema();
        let group_by = PhysicalGroupBy::new_single(vec![(
            col("a", schema.as_ref()).unwrap(),
            "a".to_string(),
        )]);
        Arc::new(
            AggregateExec::try_new(mode, group_by, vec![], vec![], input, schema)
                .unwrap()
                .with_limit_options(limit_options),
        )
    }

    fn ordered_input() -> (Arc<dyn ExecutionPlan>, LexOrdering) {
        let schema = schema();
        let ordering = ordering(schema.as_ref(), false);
        let batch = batch(schema.clone());
        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
        let input = TestMemoryExec::try_new(&partitions, schema, None)
            .unwrap()
            .try_with_sort_information(vec![ordering.clone()])
            .unwrap();

        (Arc::new(input), ordering)
    }

    fn filter_fetch(input: Arc<dyn ExecutionPlan>, fetch: usize) -> Arc<dyn ExecutionPlan> {
        Arc::new(
            FilterExecBuilder::new(lit(true), input)
                .with_fetch(Some(fetch))
                .build()
                .unwrap(),
        )
    }

    fn merge_sort_fetch(
        ordering: LexOrdering,
        input: Arc<dyn ExecutionPlan>,
        fetch: usize,
    ) -> Arc<dyn ExecutionPlan> {
        Arc::new(MergeSortExec::new(ordering, input, Some(fetch)))
    }

    fn hash_repartition(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        let partitioning = Partitioning::Hash(vec![col("a", input.schema().as_ref()).unwrap()], 3);
        Arc::new(RepartitionExec::try_new(input, partitioning).unwrap())
    }

    fn project_a(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        Arc::new(
            ProjectionExec::try_new(
                vec![(col("a", input.schema().as_ref()).unwrap(), "a".to_string())],
                input,
            )
            .unwrap(),
        )
    }

    fn schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
    }

    fn batch(schema: Arc<Schema>) -> RecordBatch {
        RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap()
    }

    fn ordering(schema: &Schema, descending: bool) -> LexOrdering {
        LexOrdering::new([PhysicalSortExpr::new(
            col("a", schema).unwrap(),
            SortOptions {
                descending,
                nulls_first: descending,
            },
        )])
        .unwrap()
    }
}