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

//! Defines the External shuffle repartition plan

use std::{
    any::Any,
    collections::{HashSet, vec_deque::VecDeque},
    fmt::Formatter,
    io::{Cursor, Read, Write},
    sync::{
        Arc, Weak,
        atomic::{AtomicUsize, Ordering::SeqCst},
    },
};

use arrow::{
    array::ArrayRef,
    datatypes::{Schema, SchemaRef},
    record_batch::{RecordBatch, RecordBatchOptions},
    row::{RowConverter, Rows, SortField},
};
use async_trait::async_trait;
use auron_memmgr::{
    MemConsumer, MemConsumerInfo, MemManager,
    spill::{Spill, SpillCompressedReader, SpillCompressedWriter, try_new_spill},
};
use bytesize::ByteSize;
use datafusion::{
    common::{DataFusionError, Result, Statistics, utils::proxy::VecAllocExt},
    execution::context::TaskContext,
    physical_expr::{
        EquivalenceProperties, PhysicalExprRef, PhysicalSortExpr, expressions::Column,
    },
    physical_plan::{
        DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
        SendableRecordBatchStream,
        execution_plan::{Boundedness, EmissionType},
        metrics::{ExecutionPlanMetricsSet, MetricsSet},
    },
};
use datafusion_ext_commons::{
    algorithm::loser_tree::{ComparableForLoserTree, LoserTree},
    arrow::{
        array_size::BatchSize,
        selection::{BatchInterleaver, create_batch_interleaver, take_batch},
    },
    compute_suggested_batch_size_for_kway_merge, compute_suggested_batch_size_for_output,
    downcast_any,
    io::{read_len, read_one_batch, write_len, write_one_batch},
};
use futures::StreamExt;
use itertools::Itertools;
use once_cell::sync::OnceCell;
use parking_lot::Mutex;

use crate::common::{
    execution_context::{
        ExecutionContext, WrappedRecordBatchSender, WrappedRecordBatchWithKeyRowsSender,
        WrappedSenderTrait,
    },
    key_rows_output::{RecordBatchWithKeyRows, SendableRecordBatchWithKeyRowsStream},
    timer_helper::TimerHelper,
};

// reserve memory for each spill
// estimated size: bufread=64KB + lz4dec.src=64KB + lz4dec.dest=64KB
const SPILL_OFFHEAP_MEM_COST: usize = 200000;

// max number of sorted batches to merge
const NUM_MAX_MERGING_BATCHES: usize = 32;

#[derive(Debug)]
pub struct SortExec {
    input: Arc<dyn ExecutionPlan>,
    exprs: Vec<PhysicalSortExpr>,
    limit: Option<usize>,
    offset: usize,
    metrics: ExecutionPlanMetricsSet,
    record_output: bool,
    props: OnceCell<PlanProperties>,
}

impl SortExec {
    pub fn new(
        input: Arc<dyn ExecutionPlan>,
        exprs: Vec<PhysicalSortExpr>,
        limit: Option<usize>,
        offset: usize,
    ) -> Self {
        let metrics = ExecutionPlanMetricsSet::new();
        Self {
            input,
            exprs,
            limit,
            offset,
            metrics,
            record_output: true,
            props: OnceCell::new(),
        }
    }

    pub fn sort_exprs(&self) -> &[PhysicalSortExpr] {
        &self.exprs
    }
}

pub fn create_default_ascending_sort_exec(
    input: Arc<dyn ExecutionPlan>,
    key_exprs: &[PhysicalExprRef],
    execution_plan_metrics: Option<ExecutionPlanMetricsSet>,
    record_output: bool,
) -> Arc<dyn ExecutionPlan> {
    let mut sort_exec = SortExec::new(
        input,
        key_exprs
            .iter()
            .map(|e| PhysicalSortExpr {
                expr: e.clone(),
                options: Default::default(),
            })
            .collect(),
        None,
        0,
    );
    if let Some(execution_plan_metrics) = execution_plan_metrics {
        sort_exec.metrics = execution_plan_metrics;
        sort_exec.record_output = record_output;
    }
    Arc::new(sort_exec)
}

impl DisplayAs for SortExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
        let exprs = self
            .exprs
            .iter()
            .map(|e| e.to_string())
            .collect::<Vec<_>>()
            .join(", ");
        write!(f, "SortExec: {exprs}")
    }
}

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

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }

    fn properties(&self) -> &PlanProperties {
        self.props.get_or_init(|| {
            PlanProperties::new(
                EquivalenceProperties::new(self.schema()),
                self.input.output_partitioning().clone(),
                EmissionType::Both,
                Boundedness::Bounded,
            )
        })
    }

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

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        Ok(Arc::new(Self::new(
            children[0].clone(),
            self.exprs.clone(),
            self.limit,
            self.offset,
        )))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        let projection: Vec<usize> = (0..self.schema().fields().len()).collect();
        self.execute_projected_without_key_rows(partition, context, &projection)
    }

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

    fn statistics(&self) -> Result<Statistics> {
        Statistics::with_fetch(
            self.input.partition_statistics(None)?,
            self.schema(),
            self.limit,
            self.offset,
            1,
        )
    }
}

impl SortExec {
    fn init_sorter(
        &self,
        exec_ctx: Arc<ExecutionContext>,
        projection: &[usize],
    ) -> Result<Arc<ExternalSorter>> {
        let prune_sort_keys_from_batch = Arc::new(PruneSortKeysFromBatch::try_new(
            self.input.schema(),
            projection,
            &self.exprs,
        )?);
        let mut sorter = Arc::new(ExternalSorter {
            exec_ctx,
            mem_consumer_info: None,
            weak: Weak::new(),
            prune_sort_keys_from_batch: prune_sort_keys_from_batch.clone(),
            skip: self.offset,
            limit: self.limit.unwrap_or(usize::MAX),
            record_output: self.record_output,
            in_mem_blocks: Default::default(),
            spills: Default::default(),
            num_total_rows: Default::default(),
            mem_total_size: Default::default(),
        });
        MemManager::register_consumer(sorter.clone(), true);

        unsafe {
            // safety: set weak reference to sorter
            let weak = Arc::downgrade(&sorter);
            let sorter_mut = Arc::get_mut_unchecked(&mut sorter);
            sorter_mut.weak = weak;
        }
        Ok(sorter)
    }

    /// Execute with projection without outputting key rows
    fn execute_projected_without_key_rows(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
        projection: &[usize],
    ) -> Result<SendableRecordBatchStream> {
        let exec_ctx =
            ExecutionContext::new(context.clone(), partition, self.schema(), &self.metrics);
        let sorter = self.init_sorter(exec_ctx.clone(), projection)?;

        let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
        let input = exec_ctx.execute_with_input_stats(&self.input)?;
        let mut coalesced = exec_ctx.coalesce_with_default_batch_size(input);

        // Implement proper key rows output by using the existing sorter with key rows
        // support
        let output = exec_ctx
            .clone()
            .output_with_sender("Sort", move |sender| async move {
                let _timer = elapsed_compute.timer();
                sender.exclude_time(&elapsed_compute);

                while let Some(batch) = elapsed_compute
                    .exclude_timer_async(coalesced.next())
                    .await
                    .transpose()?
                {
                    sorter.insert_batch(batch).await?;
                }
                sorter.output(sender).await?;
                Ok(())
            });
        Ok(exec_ctx.coalesce_with_default_batch_size(output))
    }

    /// Execute with projection and output key rows for optimization
    pub fn execute_with_key_rows(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchWithKeyRowsStream> {
        let projection: Vec<usize> = (0..self.schema().fields().len()).collect();
        self.execute_projected_with_key_rows(partition, context, &projection)
    }

    /// Execute with projection and output key rows for optimization
    pub fn execute_projected_with_key_rows(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
        projection: &[usize],
    ) -> Result<SendableRecordBatchWithKeyRowsStream> {
        let exec_ctx =
            ExecutionContext::new(context.clone(), partition, self.schema(), &self.metrics);
        let sorter = self.init_sorter(exec_ctx.clone(), projection)?;

        let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
        let input = exec_ctx.execute_with_input_stats(&self.input)?;
        let mut coalesced = exec_ctx.coalesce_with_default_batch_size(input);

        // Implement proper key rows output by using the existing sorter with key rows
        // support
        let output = exec_ctx.clone().output_with_sender_with_key_rows(
            "Sort",
            &self.exprs,
            move |sender| async move {
                let _timer = elapsed_compute.timer();
                sender.exclude_time(&elapsed_compute);

                while let Some(batch) = elapsed_compute
                    .exclude_timer_async(coalesced.next())
                    .await
                    .transpose()?
                {
                    sorter.insert_batch(batch).await?;
                }
                sorter.output(sender).await?;
                Ok(())
            },
        );
        Ok(output)
    }
}

struct LevelSpill {
    block: SpillSortedBlock,
    level: usize,
}

struct ExternalSorter {
    exec_ctx: Arc<ExecutionContext>,
    mem_consumer_info: Option<Weak<MemConsumerInfo>>,
    weak: Weak<Self>,
    prune_sort_keys_from_batch: Arc<PruneSortKeysFromBatch>,
    skip: usize,
    limit: usize,
    record_output: bool,
    in_mem_blocks: Arc<Mutex<Vec<InMemSortedBlock>>>,
    spills: Arc<Mutex<Vec<LevelSpill>>>,
    num_total_rows: AtomicUsize,
    mem_total_size: AtomicUsize,
}

impl ExternalSorter {
    fn sub_batch_size(&self) -> usize {
        compute_suggested_batch_size_for_kway_merge(self.mem_total_size(), self.num_total_rows())
    }

    fn output_batch_size(&self) -> usize {
        compute_suggested_batch_size_for_output(self.mem_total_size(), self.num_total_rows())
    }

    fn pruned_schema(&self) -> SchemaRef {
        self.prune_sort_keys_from_batch.pruned_schema()
    }
}

#[async_trait]
impl MemConsumer for ExternalSorter {
    fn name(&self) -> &str {
        "ExternalSorter"
    }

    fn set_consumer_info(&mut self, consumer_info: Weak<MemConsumerInfo>) {
        self.mem_consumer_info = Some(consumer_info);
    }

    fn get_consumer_info(&self) -> &Weak<MemConsumerInfo> {
        self.mem_consumer_info
            .as_ref()
            .expect("consumer info not set")
    }

    async fn spill(&self) -> Result<()> {
        let spills = self.spills.clone();
        let blocks = std::mem::take(&mut *self.in_mem_blocks.lock());
        let self_arc = self.to_arc();

        tokio::task::spawn_blocking(move || {
            let mut spills = spills.lock();
            let spill = try_new_spill(self_arc.exec_ctx.spill_metrics())?;
            let merged_block = merge_blocks::<_, SqueezeKeyCollector>(
                self_arc.clone(),
                blocks,
                SpillSortedBlockBuilder::new(self_arc.pruned_schema(), spill),
            )?;
            spills.push(LevelSpill {
                block: merged_block,
                level: 0,
            });

            // merge if there are too many spills
            let mut levels = std::mem::take(&mut *spills)
                .into_iter()
                .sorted_unstable_by_key(|spill| spill.level)
                .chunk_by(|spill| spill.level)
                .into_iter()
                .map(|(_, spills)| {
                    spills
                        .into_iter()
                        .map(|spill| spill.block)
                        .collect::<Vec<_>>()
                })
                .collect::<Vec<_>>();
            levels.push(vec![]); // add a new level for possible merging target

            for level in 0..levels.len() {
                if levels[level].len() >= NUM_MAX_MERGING_BATCHES {
                    let spill = try_new_spill(self_arc.exec_ctx.spill_metrics())?;
                    let merged = merge_blocks::<_, SqueezeKeyCollector>(
                        self_arc.clone(),
                        std::mem::take(&mut levels[level]),
                        SpillSortedBlockBuilder::new(self_arc.pruned_schema(), spill),
                    )?;
                    levels[level + 1].push(merged);
                } else {
                    spills.extend(
                        std::mem::take(&mut levels[level])
                            .into_iter()
                            .map(|block| LevelSpill { block, level }),
                    );
                }
            }
            Ok::<_, DataFusionError>(())
        })
        .await
        .expect("tokio error")?;

        self.update_mem_used(0).await?;
        Ok(())
    }
}

trait SortedBlock: Send {
    fn is_single_batch(&self) -> bool;
    fn next_batch(&mut self) -> Result<Option<RecordBatch>>;
    fn next_key(&mut self) -> Result<()>;
    fn cur_key(&self) -> &[u8];
    fn is_equal_to_prev_key(&self) -> bool;
    fn mem_used(&self) -> usize;
}

struct InMemSortedBlock {
    sorted_keys: VecDeque<InMemRowsKeyCollector>,
    sorted_batches: VecDeque<RecordBatch>,
    mem_used: usize,
    cur_row_idx: usize,
}

impl SortedBlock for InMemSortedBlock {
    fn is_single_batch(&self) -> bool {
        self.sorted_batches.len() == 1
    }

    fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
        let batch = self.sorted_batches.pop_front();
        if let Some(batch) = batch {
            self.mem_used -= batch.get_batch_mem_size();
            self.mem_used -= self.sorted_keys[0].mem_size();
            self.sorted_keys.pop_front().expect("missing key");
            self.cur_row_idx = usize::MAX;
            Ok(Some(batch))
        } else {
            Ok(None)
        }
    }

    fn next_key(&mut self) -> Result<()> {
        self.cur_row_idx += 1;
        Ok(())
    }

    fn cur_key(&self) -> &[u8] {
        self.sorted_keys[0].key(self.cur_row_idx)
    }

    fn is_equal_to_prev_key(&self) -> bool {
        self.cur_row_idx > 0 && {
            let key0 = self.sorted_keys[0].key(self.cur_row_idx - 1);
            let key1 = self.sorted_keys[0].key(self.cur_row_idx);
            key0 == key1
        }
    }

    fn mem_used(&self) -> usize {
        self.mem_used
    }
}

struct SpillSortedBlock {
    pruned_schema: SchemaRef,
    spill_reader: SpillCompressedReader<'static>,
    cur_key_reader: SortedKeysReader,
    _spill: Box<dyn Spill>,
}

impl SortedBlock for SpillSortedBlock {
    fn is_single_batch(&self) -> bool {
        false
    }

    fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
        if let Some((num_rows, cols)) = read_one_batch(&mut self.spill_reader, &self.pruned_schema)?
        {
            let batch = RecordBatch::try_new_with_options(
                self.pruned_schema.clone(),
                cols,
                &RecordBatchOptions::new().with_row_count(Some(num_rows)),
            )?;
            Ok(Some(batch))
        } else {
            Ok(None)
        }
    }

    fn next_key(&mut self) -> Result<()> {
        Ok(self.cur_key_reader.next_key(&mut self.spill_reader)?)
    }

    fn cur_key(&self) -> &[u8] {
        &self.cur_key_reader.cur_key
    }

    fn is_equal_to_prev_key(&self) -> bool {
        self.cur_key_reader.is_equal_to_prev
    }

    fn mem_used(&self) -> usize {
        0
    }
}

trait SortedBlockBuilder<B: SortedBlock, KS: KeyCollector> {
    fn add_batch_and_keys(&mut self, batch: RecordBatch, keys: KS) -> Result<()>;
    fn finish(self) -> Result<B>;
}

#[derive(Default)]
struct InMemSortedBlockBuilder {
    sorted_keys: VecDeque<InMemRowsKeyCollector>,
    sorted_batches: VecDeque<RecordBatch>,
    mem_used: usize,
}

impl SortedBlockBuilder<InMemSortedBlock, InMemRowsKeyCollector> for InMemSortedBlockBuilder {
    fn add_batch_and_keys(
        &mut self,
        batch: RecordBatch,
        keys: InMemRowsKeyCollector,
    ) -> Result<()> {
        self.mem_used += batch.get_batch_mem_size();
        self.mem_used += keys.mem_size();
        self.sorted_keys.push_back(keys);
        self.sorted_batches.push_back(batch);
        Ok(())
    }

    fn finish(mut self) -> Result<InMemSortedBlock> {
        self.sorted_keys
            .push_front(InMemRowsKeyCollector::default()); // for first pop
        Ok(InMemSortedBlock {
            sorted_keys: self.sorted_keys,
            sorted_batches: self.sorted_batches,
            cur_row_idx: 0,
            mem_used: self.mem_used,
        })
    }
}

struct SpillSortedBlockBuilder {
    pruned_schema: SchemaRef,
    spill: Box<dyn Spill>,
    spill_writer: SpillCompressedWriter<'static>,
}

impl SpillSortedBlockBuilder {
    fn new(pruned_schema: SchemaRef, mut spill: Box<dyn Spill>) -> Self {
        let spill_writer = unsafe {
            // safety: bypass lifetime check, spill writer has the same lifetime as spill
            std::mem::transmute(spill.get_compressed_writer())
        };
        Self {
            pruned_schema,
            spill,
            spill_writer,
        }
    }
}

impl SortedBlockBuilder<SpillSortedBlock, SqueezeKeyCollector> for SpillSortedBlockBuilder {
    fn add_batch_and_keys(&mut self, batch: RecordBatch, keys: SqueezeKeyCollector) -> Result<()> {
        write_one_batch(batch.num_rows(), batch.columns(), &mut self.spill_writer)?;
        self.spill_writer.write_all(&keys.store)?;
        Ok(())
    }

    fn finish(self) -> Result<SpillSortedBlock> {
        let spill = self.spill;
        self.spill_writer.finish()?;

        let spill_reader = unsafe {
            // safety: bypass lifetime check, spill reader has the same lifetime as spill
            std::mem::transmute(spill.get_compressed_reader())
        };
        Ok(SpillSortedBlock {
            pruned_schema: self.pruned_schema,
            spill_reader,
            cur_key_reader: SortedKeysReader::default(),
            _spill: spill,
        })
    }
}

impl Drop for ExternalSorter {
    fn drop(&mut self) {
        MemManager::deregister_consumer(self);
    }
}

impl ExternalSorter {
    async fn insert_batch(self: &Arc<Self>, batch: RecordBatch) -> Result<()> {
        if batch.num_rows() == 0 {
            return Ok(());
        }
        self.num_total_rows.fetch_add(batch.num_rows(), SeqCst);
        self.mem_total_size
            .fetch_add(batch.get_batch_mem_size(), SeqCst);

        // sort keys
        // NOTE: we use stable merge sort for longer keys due to less comparison
        let (keys, batch) = self.prune_sort_keys_from_batch.prune(batch)?;
        let sorted_indices = if keys.size() / keys.num_rows() <= 8 {
            (0..keys.num_rows() as u32).sorted_unstable_by_key(|&row_idx| unsafe {
                // safety: bypass boundary and lifetime checking
                std::mem::transmute::<_, &'static [u8]>(
                    keys.row_unchecked(row_idx as usize).as_ref(),
                )
            })
        } else {
            (0..keys.num_rows() as u32).sorted_by_key(|&row_idx| unsafe {
                // safety: bypass boundary and lifetime checking
                std::mem::transmute::<_, &'static [u8]>(
                    keys.row_unchecked(row_idx as usize).as_ref(),
                )
            })
        }
        .take(self.limit)
        .collect::<Vec<_>>();

        // build keys
        let mut key_collector = InMemRowsKeyCollector::default();
        key_collector.reserve(keys.num_rows(), keys.size());
        for &row_idx in &sorted_indices {
            key_collector.add_key(keys.row(row_idx as usize).as_ref());
        }
        key_collector.freeze();

        // build batch
        let sorted_batch = if !self.prune_sort_keys_from_batch.is_all_pruned() {
            take_batch(batch, sorted_indices)?
        } else {
            create_zero_column_batch(batch.num_rows())
        };

        // add to in-mem blocks
        let mut block_builder = InMemSortedBlockBuilder::default();
        block_builder.add_batch_and_keys(sorted_batch, key_collector)?;
        self.in_mem_blocks.lock().push(block_builder.finish()?);

        let mem_used = self.mem_used();
        self.update_mem_used(mem_used).await?;
        Ok(())
    }

    async fn output(self: &Arc<Self>, sender: Arc<impl WrappedSenderTrait>) -> Result<()> {
        // if external merging is required, we need to spill in-mem data
        // to free memory as soon as possible
        let has_spill = !self.spills.lock().is_empty();
        let has_in_mem_data = !self.in_mem_blocks.lock().is_empty();
        if has_spill && has_in_mem_data {
            log::info!(
                "{} spills rest in-mem data to disk, size={}",
                self.name(),
                ByteSize(self.mem_used() as u64),
            );
            self.force_spill().await?;
        }
        self.set_spillable(false);

        let output_batch_size = self.output_batch_size();
        let spills = self.spills.lock().drain(..).collect::<Vec<_>>();
        log::info!(
            "{} starts outputting ({} spills + in_mem: {})",
            self.name(),
            spills.len(),
            ByteSize(self.mem_used() as u64)
        );

        // no spills -- output in-mem batches
        if spills.is_empty() {
            let in_mem_blocks = std::mem::take(&mut *self.in_mem_blocks.lock());
            if !in_mem_blocks.is_empty() {
                let mut merger = Merger::try_new(self.clone(), in_mem_blocks)?;
                if self.skip > 0 {
                    let _ = merger.skip_rows::<InMemRowsKeyCollector>(self.skip, output_batch_size);
                }
                while let Some((key_collector, pruned_batch)) =
                    merger.next::<InMemRowsKeyCollector>(output_batch_size)?
                {
                    if self.record_output {
                        self.exec_ctx
                            .baseline_metrics()
                            .record_output(pruned_batch.num_rows());
                    }
                    send_output_batch(
                        sender.as_ref(),
                        pruned_batch,
                        &self.prune_sort_keys_from_batch,
                        key_collector,
                    )
                    .await?;
                }
            }
            self.update_mem_used(0).await?;
            return Ok(());
        }

        let spill_blocks = spills.into_iter().map(|spill| spill.block).collect();
        let mut merger = Merger::try_new(self.to_arc(), spill_blocks)?;
        if self.skip > 0 {
            let _ = merger.skip_rows::<InMemRowsKeyCollector>(self.skip, output_batch_size);
        }
        while let Some((key_collector, pruned_batch)) =
            merger.next::<InMemRowsKeyCollector>(output_batch_size)?
        {
            self.update_mem_used(merger.cursors_mem_used()).await?;
            if self.record_output {
                self.exec_ctx
                    .baseline_metrics()
                    .record_output(pruned_batch.num_rows());
            }
            send_output_batch(
                sender.as_ref(),
                pruned_batch,
                &self.prune_sort_keys_from_batch,
                key_collector,
            )
            .await?;
        }
        self.update_mem_used(0).await?;
        Ok(())
    }

    fn to_arc(&self) -> Arc<Self> {
        self.weak.upgrade().expect("ExternalSorter.weak not set")
    }

    fn num_total_rows(&self) -> usize {
        self.num_total_rows.load(SeqCst)
    }

    fn mem_total_size(&self) -> usize {
        self.mem_total_size.load(SeqCst)
    }

    fn mem_used(&self) -> usize {
        let in_mem_blocks = self.in_mem_blocks.lock();
        in_mem_blocks
            .iter()
            .map(|block| block.mem_used())
            .sum::<usize>()
    }
}

async fn send_output_batch(
    sender: &impl WrappedSenderTrait,
    pruned_batch: RecordBatch,
    prune_sort_keys_from_batch: &PruneSortKeysFromBatch,
    key_collector: impl KeyCollector,
) -> Result<()> {
    if let Ok(sender) = downcast_any!(sender, WrappedRecordBatchSender) {
        let batch = prune_sort_keys_from_batch.restore(pruned_batch, key_collector)?;
        sender.send(batch).await
    } else if let Ok(sender) = downcast_any!(sender, WrappedRecordBatchWithKeyRowsSender) {
        let key_rows = Arc::new(key_collector.into_rows(
            pruned_batch.num_rows(),
            &prune_sort_keys_from_batch.sort_row_converter.lock(),
        )?);
        let batch =
            prune_sort_keys_from_batch.restore_from_existed_key_rows(pruned_batch, &key_rows)?;
        sender
            .send(RecordBatchWithKeyRows { batch, key_rows })
            .await;
    }
    Ok(())
}

struct SortedBlockCursor<B: SortedBlock> {
    id: usize,
    input: B,
    cur_batch_num_rows: usize,
    cur_batches: Vec<RecordBatch>,
    cur_batches_changed: bool,
    cur_key_row_idx: usize,
    cur_batch_idx: usize,
    cur_row_idx: usize,
    cur_mem_used: usize,
    finished: bool,
}

impl<B: SortedBlock> ComparableForLoserTree for SortedBlockCursor<B> {
    #[inline(always)]
    fn lt(&self, other: &Self) -> bool {
        if other.finished || self.finished {
            return other.finished;
        }
        self.input.cur_key() < other.input.cur_key()
    }
}

impl<B: SortedBlock> SortedBlockCursor<B> {
    fn try_from_block(id: usize, block: B) -> Result<Self> {
        let mut new = Self {
            id,
            input: block,
            cur_batch_num_rows: 0,
            cur_batches: vec![],
            cur_batches_changed: false,
            cur_key_row_idx: 0,
            cur_batch_idx: 0,
            cur_row_idx: 0,
            cur_mem_used: 0,
            finished: false,
        };
        new.next_key()?; // load first record
        Ok(new)
    }

    fn block(&self) -> &B {
        &self.input
    }

    // forwards to next key and returns current key
    fn next_key(&mut self) -> Result<()> {
        assert!(
            !self.finished,
            "calling next_key() on finished sort spill cursor"
        );

        if self.cur_key_row_idx >= self.cur_batches.last().map(|b| b.num_rows()).unwrap_or(0)
            && !self.load_next_batch()?
        {
            self.finished = true;
            return Ok(());
        }
        self.input.next_key()?;
        self.cur_key_row_idx += 1;
        Ok(())
    }

    fn load_next_batch(&mut self) -> Result<bool> {
        if let Some(batch) = self.input.next_batch()? {
            self.cur_batch_num_rows = batch.num_rows();
            self.cur_mem_used += batch.get_batch_mem_size();
            self.cur_batches.push(batch);
            self.cur_key_row_idx = 0;
            self.cur_batches_changed = true;
            return Ok(true);
        }
        self.finished = true;
        Ok(false)
    }

    fn next_row(&mut self) -> (usize, usize) {
        let batch_idx = self.cur_batch_idx;
        let row_idx = self.cur_row_idx;

        self.cur_row_idx += 1;
        if self.cur_row_idx >= self.cur_batches[self.cur_batch_idx].num_rows() {
            self.cur_batch_idx += 1;
            self.cur_row_idx = 0;
        }
        (batch_idx, row_idx)
    }

    fn clear_finished_batches(&mut self) {
        if self.cur_batch_idx > 0 {
            for batch in self.cur_batches.drain(..self.cur_batch_idx) {
                self.cur_batches_changed = true;
                self.cur_mem_used -= batch.get_batch_mem_size();
            }
            self.cur_batch_idx = 0;
        }
    }
}

struct Merger<B: SortedBlock> {
    sorter: Arc<ExternalSorter>,
    cursors: LoserTree<SortedBlockCursor<B>>,
    batch_interleaver: BatchInterleaver,
    single_batch_mode: bool,
    batches_base_indices: Vec<usize>,
    num_total_output_rows: usize,
}

impl<B: SortedBlock> Merger<B> {
    fn try_new(sorter: Arc<ExternalSorter>, blocks: Vec<B>) -> Result<Self> {
        let single_batch_mode = blocks.iter().all(|block| block.is_single_batch());
        let cursors = LoserTree::new(
            blocks
                .into_iter()
                .enumerate()
                .map(|(id, block)| SortedBlockCursor::try_from_block(id, block))
                .collect::<Result<Vec<_>>>()?,
        );
        let batch_interleaver = if single_batch_mode {
            let batches = cursors
                .values()
                .iter()
                .map(|cursor| cursor.cur_batches[0].clone())
                .collect::<Vec<_>>();
            create_batch_interleaver(&batches, true)?
        } else {
            create_batch_interleaver(&[RecordBatch::new_empty(sorter.pruned_schema())], true)?
        };

        Ok(Self {
            sorter,
            cursors,
            batch_interleaver,
            single_batch_mode,
            batches_base_indices: vec![],
            num_total_output_rows: 0,
        })
    }

    fn cursors_mem_used(&self) -> usize {
        self.cursors.len() * SPILL_OFFHEAP_MEM_COST
            + self
                .cursors
                .values()
                .iter()
                .map(|cursor| cursor.cur_mem_used)
                .sum::<usize>()
    }

    fn next<KC: KeyCollector>(&mut self, batch_size: usize) -> Result<Option<(KC, RecordBatch)>> {
        let max_num_rows = batch_size.min(self.sorter.limit - self.num_total_output_rows);
        let mut num_rows = 0;

        // collect merged records to staging
        let mut min_cursor = self.cursors.peek_mut();
        let mut key_collector = KC::default();
        let mut cursor_ids = vec![];

        key_collector.reserve(num_rows, 0);
        while !min_cursor.finished && num_rows < max_num_rows {
            if !self.sorter.prune_sort_keys_from_batch.is_all_pruned() {
                cursor_ids.push(min_cursor.id);
            }
            key_collector.add_key(&min_cursor.block().cur_key());
            min_cursor.next_key()?;
            num_rows += 1;

            // fetch next min key from loser tree only if it is different from previous key
            if min_cursor.finished || !min_cursor.block().is_equal_to_prev_key() {
                min_cursor.adjust();
            }
        }

        if num_rows == 0 {
            return Ok(None);
        }
        key_collector.freeze();
        drop(min_cursor);

        // collect pruned columns
        let pruned_batch = if !self.sorter.prune_sort_keys_from_batch.is_all_pruned() {
            if !self.single_batch_mode
                && self
                    .cursors
                    .values()
                    .iter()
                    .any(|cursor| cursor.cur_batches_changed)
            {
                let mut batches_base_indices = vec![];
                let mut batches = vec![];
                for cursor in self.cursors.values_mut() {
                    batches_base_indices.push(batches.len());
                    batches.extend(cursor.cur_batches.clone());
                    cursor.cur_batches_changed = false;
                }
                self.batch_interleaver = create_batch_interleaver(&batches, true)?;
                self.batches_base_indices = batches_base_indices;
            }
            let batch_interleaver = &self.batch_interleaver;

            // interleave output batch
            let staging_indices = if !self.single_batch_mode {
                cursor_ids
                    .into_iter()
                    .map(|cursor_id| {
                        let cursor = &mut self.cursors.values_mut()[cursor_id];
                        let (batch_idx, row_idx) = cursor.next_row();
                        (self.batches_base_indices[cursor.id] + batch_idx, row_idx)
                    })
                    .collect::<Vec<_>>()
            } else {
                cursor_ids
                    .into_iter()
                    .map(|cursor_id| {
                        let cursor = &mut self.cursors.values_mut()[cursor_id];
                        let (_, row_idx) = cursor.next_row();
                        (cursor.id, row_idx)
                    })
                    .collect::<Vec<_>>()
            };
            batch_interleaver(&staging_indices)?
        } else {
            create_zero_column_batch(num_rows)
        };
        self.num_total_output_rows += num_rows;

        for cursor in self.cursors.values_mut() {
            cursor.clear_finished_batches();
        }
        Ok(Some((key_collector, pruned_batch)))
    }

    pub fn skip_rows<KC: KeyCollector>(
        &mut self,
        skip: usize,
        suggested_batch_size: usize,
    ) -> Result<()> {
        let mut remaining = skip;
        while remaining > 0 {
            let batch_size = remaining.min(suggested_batch_size);
            if self.next::<KC>(batch_size)?.is_none() {
                break;
            }
            remaining -= batch_size;
        }
        Ok(())
    }
}

fn merge_blocks<B: SortedBlock, KC: KeyCollector>(
    sorter: Arc<ExternalSorter>,
    blocks: Vec<impl SortedBlock>,
    mut block_builder: impl SortedBlockBuilder<B, KC>,
) -> Result<B> {
    assert!(blocks.len() >= 1);
    let sub_batch_size = sorter.sub_batch_size();
    let mut merger = Merger::try_new(sorter, blocks)?;
    while let Some((key_collector, pruned_batch)) = merger.next::<KC>(sub_batch_size)? {
        block_builder.add_batch_and_keys(pruned_batch, key_collector)?;
    }
    Ok(block_builder.finish()?)
}

fn create_zero_column_batch(num_rows: usize) -> RecordBatch {
    static EMPTY_SCHEMA: OnceCell<SchemaRef> = OnceCell::new();
    let empty_schema = EMPTY_SCHEMA
        .get_or_init(|| Arc::new(Schema::empty()))
        .clone();
    RecordBatch::try_new_with_options(
        empty_schema,
        vec![],
        &RecordBatchOptions::new().with_row_count(Some(num_rows)),
    )
    .expect("failed to create empty RecordBatch")
}

struct PruneSortKeysFromBatch {
    input_projection: Vec<usize>,
    sort_row_converter: Arc<Mutex<RowConverter>>,
    key_exprs: Vec<PhysicalSortExpr>,
    key_cols: HashSet<usize>,
    restored_col_mappers: Vec<ColMapper>,
    restored_schema: SchemaRef,
    pruned_schema: SchemaRef,
}

#[derive(Clone, Copy)]
enum ColMapper {
    FromPrunedBatch(usize),
    FromKey(usize),
}

impl PruneSortKeysFromBatch {
    fn try_new(
        input_schema: SchemaRef,
        input_projection: &[usize],
        exprs: &[PhysicalSortExpr],
    ) -> Result<Self> {
        let sort_row_converter = Arc::new(Mutex::new(RowConverter::new(
            exprs
                .iter()
                .map(|expr: &PhysicalSortExpr| {
                    Ok(SortField::new_with_options(
                        expr.expr.data_type(&input_schema)?,
                        expr.options,
                    ))
                })
                .collect::<Result<Vec<SortField>>>()?,
        )?));
        let input_projected_schema = Arc::new(input_schema.project(input_projection)?);

        let mut relation = vec![];
        for (expr_idx, expr) in exprs.iter().enumerate() {
            if let Some(col) = expr.expr.as_any().downcast_ref::<Column>() {
                relation.push((expr_idx, col.index()));
            }
        }

        // compute pruned col indices
        let pruned_cols = relation
            .iter()
            .map(|(_, col_idx)| *col_idx)
            .collect::<HashSet<_>>();

        // compute schema after pruning
        let mut fields_after_pruning = vec![];
        for field_idx in 0..input_projected_schema.fields().len() {
            if !pruned_cols.contains(&field_idx) {
                fields_after_pruning.push(input_projected_schema.field(field_idx).clone());
            }
        }
        let pruned_schema = Arc::new(Schema::new(fields_after_pruning));

        // compute col mappers for restoring
        let restored_schema = input_projected_schema;
        let mut restored_col_mappers = vec![];
        let mut num_pruned_cols = 0;
        for col_idx in 0..restored_schema.fields().len() {
            if let Some(expr_idx) = relation.iter().find(|kv| kv.1 == col_idx).map(|kv| kv.0) {
                restored_col_mappers.push(ColMapper::FromKey(expr_idx));
                num_pruned_cols += 1;
            } else {
                restored_col_mappers.push(ColMapper::FromPrunedBatch(col_idx - num_pruned_cols));
            }
        }

        Ok(Self {
            input_projection: input_projection.to_vec(),
            sort_row_converter,
            key_exprs: exprs.to_vec(),
            key_cols: pruned_cols,
            restored_col_mappers,
            pruned_schema,
            restored_schema,
        })
    }

    fn is_all_pruned(&self) -> bool {
        self.pruned_schema.fields().is_empty()
    }

    fn pruned_schema(&self) -> SchemaRef {
        self.pruned_schema.clone()
    }

    fn restored_schema(&self) -> SchemaRef {
        self.restored_schema.clone()
    }

    fn prune(&self, batch: RecordBatch) -> Result<(Rows, RecordBatch)> {
        // compute key rows
        let key_cols: Vec<ArrayRef> = self
            .key_exprs
            .iter()
            .map(|expr| {
                expr.expr
                    .evaluate(&batch)
                    .and_then(|cv| cv.into_array(batch.num_rows()))
            })
            .collect::<Result<_>>()?;
        let key_rows = self.sort_row_converter.lock().convert_columns(&key_cols)?;

        let retained_cols = batch
            .project(&self.input_projection)?
            .columns()
            .iter()
            .enumerate()
            .filter(|(col_idx, _)| !self.key_cols.contains(col_idx))
            .map(|(_, col)| col)
            .cloned()
            .collect();
        let pruned_batch = RecordBatch::try_new_with_options(
            self.pruned_schema(),
            retained_cols,
            &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
        )?;
        Ok((key_rows, pruned_batch))
    }

    fn restore<'a, KC: KeyCollector>(
        &self,
        pruned_batch: RecordBatch,
        key_collector: KC,
    ) -> Result<RecordBatch> {
        let num_rows = pruned_batch.num_rows();
        let key_rows = key_collector.into_rows(num_rows, &self.sort_row_converter.lock())?;
        self.restore_from_existed_key_rows(pruned_batch, &key_rows)
    }

    fn restore_from_existed_key_rows(
        &self,
        pruned_batch: RecordBatch,
        key_rows: &Rows,
    ) -> Result<RecordBatch> {
        let key_cols = OnceCell::new();

        // restore batch
        let mut restored_fields = vec![];
        for &map in &self.restored_col_mappers {
            match map {
                ColMapper::FromPrunedBatch(idx) => {
                    restored_fields.push(pruned_batch.column(idx).clone());
                }
                ColMapper::FromKey(idx) => {
                    let sort_row_converter = self.sort_row_converter.clone();
                    let key_cols = key_cols.get_or_try_init(move || {
                        sort_row_converter.lock().convert_rows(key_rows.iter())
                    })?;
                    restored_fields.push(key_cols[idx].clone());
                }
            }
        }
        Ok(RecordBatch::try_new_with_options(
            self.restored_schema(),
            restored_fields,
            &RecordBatchOptions::new().with_row_count(Some(pruned_batch.num_rows())),
        )?)
    }
}

trait KeyCollector: Default {
    fn reserve(&mut self, num_rows: usize, data_size: usize);
    fn add_key(&mut self, key: &[u8]);
    fn freeze(&mut self);
    fn mem_size(&self) -> usize;
    fn into_rows(self, num_rows: usize, row_converter: &RowConverter) -> Result<Rows>;
}

#[derive(Default, Clone)]
struct InMemRowsKeyCollector {
    buffer: Vec<u8>,
    offsets: Vec<usize>,
}

impl InMemRowsKeyCollector {
    fn key(&self, i: usize) -> &[u8] {
        unsafe {
            // safety: performance critical, bypass bounds check
            let start = *self.offsets.get_unchecked(i);
            let end = *self.offsets.get_unchecked(i + 1);
            self.buffer.get_unchecked(start..end)
        }
    }
}

impl KeyCollector for InMemRowsKeyCollector {
    fn reserve(&mut self, num_rows: usize, data_size: usize) {
        self.offsets.reserve(num_rows + 1);
        self.buffer.reserve(data_size);
    }

    fn add_key(&mut self, key: &[u8]) {
        self.offsets.push(self.buffer.len());
        self.buffer.extend_from_slice(key);
    }

    fn freeze(&mut self) {
        self.offsets.push(self.buffer.len());
        self.offsets.shrink_to_fit();
        self.buffer.shrink_to_fit();
    }

    fn mem_size(&self) -> usize {
        self.offsets.allocated_size() + self.buffer.allocated_size()
    }

    fn into_rows(self, num_rows: usize, row_converter: &RowConverter) -> Result<Rows> {
        assert_eq!(self.offsets.len() - 1, num_rows);
        let mut rows = row_converter.empty_rows(0, 0);
        unsafe {
            // safety: access rows.buffer/offsets
            struct XRows {
                buffer: Vec<u8>,
                offsets: Vec<usize>,
            }
            let xrows = std::mem::transmute::<_, &mut XRows>(&mut rows);
            xrows.buffer = self.buffer;
            xrows.offsets = self.offsets;
        }
        Ok(rows)
    }
}

#[derive(Default)]
struct SqueezeKeyCollector {
    sorted_key_writer: SortedKeysWriter,
    store: Vec<u8>,
}

impl KeyCollector for SqueezeKeyCollector {
    fn reserve(&mut self, _num_rows: usize, _data_size: usize) {
        // do nothing because we cannot get squeezed data size at this moment
    }

    fn add_key(&mut self, key: &[u8]) {
        self.sorted_key_writer
            .write_key(key, &mut self.store)
            .expect("failed to write key");
    }

    fn freeze(&mut self) {
        self.store.shrink_to_fit();
    }

    fn mem_size(&self) -> usize {
        self.store.allocated_size()
    }

    fn into_rows(self, num_rows: usize, row_converter: &RowConverter) -> Result<Rows> {
        let mut sorted_key_reader = SortedKeysReader::default();
        let mut r = Cursor::new(self.store);
        let mut simple_key_collector = InMemRowsKeyCollector::default();

        for _ in 0..num_rows {
            sorted_key_reader.next_key(&mut r)?;
            simple_key_collector.add_key(&sorted_key_reader.cur_key);
        }
        simple_key_collector.into_rows(num_rows, row_converter)
    }
}

#[derive(Default)]
struct SortedKeysWriter {
    cur_key: Vec<u8>,
}

impl SortedKeysWriter {
    fn write_key(&mut self, key: &[u8], w: &mut impl Write) -> std::io::Result<()> {
        let prefix_len = common_prefix_len(&self.cur_key, key);
        let suffix_len = key.len() - prefix_len;
        if prefix_len == key.len() && suffix_len == 0 {
            write_len(0, w)?; // indicates same record
        } else {
            let write_buf = unsafe {
                // safety: bypass bounds check in performance critical code
                self.cur_key
                    .reserve(key.len().saturating_sub(self.cur_key.len()));
                self.cur_key.set_len(key.len());
                let write_buf = self
                    .cur_key
                    .get_unchecked_mut(prefix_len..key.len())
                    .as_mut();
                write_buf.copy_from_slice(key.get_unchecked(prefix_len..));
                write_buf
            };
            write_len(suffix_len + 1, w)?;
            write_len(prefix_len, w)?;
            w.write_all(write_buf)?;
        }
        Ok(())
    }
}

#[derive(Default)]
struct SortedKeysReader {
    cur_key: Vec<u8>,
    is_equal_to_prev: bool,
}

impl SortedKeysReader {
    fn next_key(&mut self, r: &mut impl Read) -> std::io::Result<()> {
        let b = read_len(r)?;
        if b > 0 {
            self.is_equal_to_prev = false;
            let suffix_len = b - 1;
            let prefix_len = read_len(r)?;
            let read_buf = unsafe {
                // safety: bypass bounds check in performance critical code
                let new_key_len = prefix_len + suffix_len;
                self.cur_key
                    .reserve(new_key_len.saturating_sub(self.cur_key.len()));
                self.cur_key.set_len(new_key_len);
                self.cur_key
                    .get_unchecked_mut(prefix_len..new_key_len)
                    .as_mut()
            };
            r.read_exact(read_buf)?;
        } else {
            self.is_equal_to_prev = true;
        }
        Ok(())
    }
}

fn common_prefix_len(a: &[u8], b: &[u8]) -> usize {
    let min_len = a.len().min(b.len());
    let mut lcp = 0;

    macro_rules! pread {
        ($v:expr, $off:expr, $nb:expr) => {{
            // safety: bypass bounds check in performance critical code
            unsafe { std::ptr::read_unaligned($v.as_ptr().add($off) as *const [u8; $nb]) }
        }};
    }

    while lcp + 8 <= min_len && pread!(a, lcp, 8) == pread!(b, lcp, 8) {
        lcp += 8;
    }
    lcp += (lcp + 4 <= min_len && pread!(a, lcp, 4) == pread!(b, lcp, 4)) as usize * 4;
    lcp += (lcp + 2 <= min_len && pread!(a, lcp, 2) == pread!(b, lcp, 2)) as usize * 2;
    lcp += (lcp + 1 <= min_len && pread!(a, lcp, 1) == pread!(b, lcp, 1)) as usize;
    lcp
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use arrow::{
        array::Int32Array,
        compute::SortOptions,
        datatypes::{DataType, Field, Schema},
        record_batch::RecordBatch,
    };
    use auron_memmgr::MemManager;
    use datafusion::{
        assert_batches_eq,
        common::Result,
        physical_expr::{PhysicalSortExpr, expressions::Column},
        physical_plan::{ExecutionPlan, common, test::TestMemoryExec},
        prelude::SessionContext,
    };

    use super::common_prefix_len;
    use crate::sort_exec::SortExec;

    fn build_table_i32(
        a: (&str, &Vec<i32>),
        b: (&str, &Vec<i32>),
        c: (&str, &Vec<i32>),
    ) -> Result<RecordBatch> {
        let schema = Schema::new(vec![
            Field::new(a.0, DataType::Int32, false),
            Field::new(b.0, DataType::Int32, false),
            Field::new(c.0, DataType::Int32, false),
        ]);

        let batch = RecordBatch::try_new(
            Arc::new(schema),
            vec![
                Arc::new(Int32Array::from(a.1.clone())),
                Arc::new(Int32Array::from(b.1.clone())),
                Arc::new(Int32Array::from(c.1.clone())),
            ],
        )?;
        Ok(batch)
    }

    fn build_table(
        a: (&str, &Vec<i32>),
        b: (&str, &Vec<i32>),
        c: (&str, &Vec<i32>),
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let batch = build_table_i32(a, b, c)?;
        let schema = batch.schema();
        Ok(Arc::new(TestMemoryExec::try_new(
            &[vec![batch]],
            schema,
            None,
        )?))
    }

    // Verify that `common_prefix_len` correctly computes prefixes of small fixed
    // sizes. This ensures compiler optimizations do not incorrectly transform
    // the function for small lengths.
    #[test]
    fn has_common_prefix_len_been_optimized_into_bogusness() {
        let cpl = common_prefix_len(&[0; 8], &[0; 8]);
        assert_eq!(cpl, 8);
        let cpl = common_prefix_len(&[0; 4], &[0; 4]);
        assert_eq!(cpl, 4);
        let cpl = common_prefix_len(&[0; 2], &[0; 2]);
        assert_eq!(cpl, 2);
        let cpl = common_prefix_len(&[0; 1], &[0; 1]);
        assert_eq!(cpl, 1);
    }

    #[tokio::test]
    async fn test_sort_i32() -> Result<()> {
        MemManager::init(100);
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();
        let input = build_table(
            ("a", &vec![9, 8, 7, 6, 5, 4, 3, 2, 1, 0]),
            ("b", &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
            ("c", &vec![5, 6, 7, 8, 9, 0, 1, 2, 3, 4]),
        )?;
        let sort_exprs = vec![PhysicalSortExpr {
            expr: Arc::new(Column::new("a", 0)),
            options: SortOptions::default(),
        }];

        let sort = SortExec::new(input, sort_exprs, Some(6), 0);
        let output = sort.execute(0, task_ctx)?;
        let batches = common::collect(output).await?;
        let expected = vec![
            "+---+---+---+",
            "| a | b | c |",
            "+---+---+---+",
            "| 0 | 9 | 4 |",
            "| 1 | 8 | 3 |",
            "| 2 | 7 | 2 |",
            "| 3 | 6 | 1 |",
            "| 4 | 5 | 0 |",
            "| 5 | 4 | 9 |",
            "+---+---+---+",
        ];
        assert_batches_eq!(expected, &batches);

        Ok(())
    }

    #[tokio::test]
    async fn test_sort_i32_with_skip() -> Result<()> {
        MemManager::init(100);
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();
        let input = build_table(
            ("a", &vec![9, 8, 7, 6, 5, 4, 3, 2, 1, 0]),
            ("b", &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
            ("c", &vec![5, 6, 7, 8, 9, 0, 1, 2, 3, 4]),
        )?;
        let sort_exprs = vec![PhysicalSortExpr {
            expr: Arc::new(Column::new("a", 0)),
            options: SortOptions::default(),
        }];

        let sort = SortExec::new(input, sort_exprs, Some(8), 3);
        let output = sort.execute(0, task_ctx)?;
        let batches = common::collect(output).await?;
        let expected = vec![
            "+---+---+---+",
            "| a | b | c |",
            "+---+---+---+",
            "| 3 | 6 | 1 |",
            "| 4 | 5 | 0 |",
            "| 5 | 4 | 9 |",
            "| 6 | 3 | 8 |",
            "| 7 | 2 | 7 |",
            "+---+---+---+",
        ];
        assert_batches_eq!(expected, &batches);

        Ok(())
    }
}

#[cfg(test)]
mod fuzztest {
    #![allow(deprecated)]
    use std::{sync::Arc, time::Instant};

    use arrow::{
        array::{ArrayRef, StringArray, UInt32Array},
        compute::{SortOptions, concat_batches},
        record_batch::RecordBatch,
    };
    use auron_memmgr::MemManager;
    use datafusion::{
        common::{Result, stats::Precision},
        physical_expr::{LexOrdering, PhysicalSortExpr, expressions::Column},
        physical_plan::{ExecutionPlan, test::TestMemoryExec},
        prelude::{SessionConfig, SessionContext},
    };

    use crate::sort_exec::SortExec;

    #[tokio::test]
    async fn fuzztest_in_mem_sorting() -> Result<()> {
        let time_start = Instant::now();
        fuzztest_with_mem_conf(1000000000).await?;
        eprintln!("fuzztest_in_mem_sorting_time: {:?}", time_start.elapsed());
        Ok(())
    }

    #[tokio::test]
    async fn fuzztest_external_sorting() -> Result<()> {
        let time_start = Instant::now();
        fuzztest_with_mem_conf(10000).await?;
        eprintln!("fuzztest_external_sorting_time: {:?}", time_start.elapsed());
        Ok(())
    }

    async fn fuzztest_with_mem_conf(mem: usize) -> Result<()> {
        MemManager::init(mem);
        let session_ctx =
            SessionContext::new_with_config(SessionConfig::new().with_batch_size(10000));
        let task_ctx = session_ctx.task_ctx();
        let n = 1234567;

        // generate random batch for fuzzying
        let mut batches = vec![];
        let mut num_rows = 0;
        while num_rows < n {
            let rand_key1: ArrayRef = Arc::new(
                std::iter::repeat_with(|| Some(format!("{}", rand::random::<u32>())))
                    .take((n - num_rows).min(10000))
                    .collect::<StringArray>(),
            );
            let rand_key2: ArrayRef = Arc::new(
                std::iter::repeat_with(|| rand::random::<u32>())
                    .take((n - num_rows).min(10000))
                    .collect::<UInt32Array>(),
            );
            let rand_val1: ArrayRef = Arc::new(
                std::iter::repeat_with(|| rand::random::<u32>())
                    .take((n - num_rows).min(10000))
                    .collect::<UInt32Array>(),
            );
            let rand_val2: ArrayRef = Arc::new(
                std::iter::repeat_with(|| rand::random::<u32>())
                    .take((n - num_rows).min(10000))
                    .collect::<UInt32Array>(),
            );
            let batch = RecordBatch::try_from_iter_with_nullable(vec![
                ("k1", rand_key1, true),
                ("k2", rand_key2, true),
                ("v1", rand_val1, true),
                ("v2", rand_val2, true),
            ])?;
            num_rows += batch.num_rows();
            batches.push(batch);
        }
        let schema = batches[0].schema();
        let sort_exprs = vec![
            PhysicalSortExpr {
                expr: Arc::new(Column::new("k1", 0)),
                options: SortOptions::default(),
            },
            PhysicalSortExpr {
                expr: Arc::new(Column::new("k2", 1)),
                options: SortOptions::default(),
            },
        ];

        let input = Arc::new(TestMemoryExec::try_new(
            &[batches.clone()],
            schema.clone(),
            None,
        )?);
        let sort = Arc::new(SortExec::new(input, sort_exprs.clone(), None, 0));
        let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?;
        let a = concat_batches(&schema, &output)?;
        let a_row_count = sort.clone().statistics()?.num_rows;

        let input = Arc::new(TestMemoryExec::try_new(
            &[batches.clone()],
            schema.clone(),
            None,
        )?);
        let sort = Arc::new(datafusion::physical_plan::sorts::sort::SortExec::new(
            LexOrdering::new(sort_exprs.iter().cloned()).expect("invalid sort exprs"),
            input,
        ));
        let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?;
        let b = concat_batches(&schema, &output)?;
        let b_row_count = sort.clone().statistics()?.num_rows;

        assert_eq!(a.num_rows(), b.num_rows());
        assert_eq!(a_row_count, b_row_count);
        assert_eq!(a_row_count, Precision::Exact(n));
        assert!(a == b);
        Ok(())
    }
}
