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

use std::{any::Any, fmt::Formatter, mem, sync::Arc};

use arrow::{
    array::{Array, ArrayRef, Int32Array},
    compute::concat_batches,
    datatypes::SchemaRef,
    record_batch::{RecordBatch, RecordBatchOptions},
};
use datafusion::{
    common::{Result, Statistics},
    execution::context::TaskContext,
    physical_expr::{EquivalenceProperties, PhysicalExprRef, PhysicalSortExpr},
    physical_plan::{
        DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
        SendableRecordBatchStream,
        execution_plan::{Boundedness, EmissionType},
        metrics::{ExecutionPlanMetricsSet, MetricsSet, Time},
    },
};
use datafusion_ext_commons::{arrow::cast::cast, downcast_any};
use futures::StreamExt;
use once_cell::sync::OnceCell;

use crate::{
    common::execution_context::{ExecutionContext, WrappedRecordBatchSender},
    window::{WindowExpr, WindowFunctionProcessor, window_context::WindowContext},
};

#[derive(Debug)]
pub struct WindowExec {
    input: Arc<dyn ExecutionPlan>,
    context: Arc<WindowContext>,
    metrics: ExecutionPlanMetricsSet,
    props: OnceCell<PlanProperties>,
}

impl WindowExec {
    // NOTE:
    // WindowExec now supports spark's WindowExec and WindowGroupLimitExec:
    // for normal WindowExec:
    //   group_limit = None
    //   output_window_cols = true
    //
    // for partial WindowGroupLimitExec
    //   group_limit = Some(K)
    //   output_window_cols = false
    //
    // for final WindowGroupLimitExec:
    //   group_limit = Some(K)
    //   output_window_cols = true
    pub fn try_new(
        input: Arc<dyn ExecutionPlan>,
        window_exprs: Vec<WindowExpr>,
        partition_spec: Vec<PhysicalExprRef>,
        order_spec: Vec<PhysicalSortExpr>,
        group_limit: Option<usize>,
        output_window_cols: bool,
    ) -> Result<Self> {
        let context = Arc::new(WindowContext::try_new(
            input.schema(),
            window_exprs,
            partition_spec,
            order_spec,
            group_limit,
            output_window_cols,
        )?);
        Ok(Self {
            input,
            context,
            metrics: ExecutionPlanMetricsSet::new(),
            props: OnceCell::new(),
        })
    }

    pub fn window_context(&self) -> &WindowContext {
        &self.context
    }

    pub fn with_output_window_cols(&self, output_window_cols: bool) -> Self {
        Self {
            input: self.input.clone(),
            context: Arc::new(
                WindowContext::try_new(
                    self.input.schema(),
                    self.context.window_exprs.clone(),
                    self.context.partition_spec.clone(),
                    self.context.order_spec.clone(),
                    self.context.group_limit,
                    output_window_cols,
                )
                .expect("failed to create window context"),
            ),
            metrics: ExecutionPlanMetricsSet::new(),
            props: OnceCell::new(),
        }
    }
}

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

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

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

    fn schema(&self) -> SchemaRef {
        self.context.output_schema.clone()
    }

    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::try_new(
            children[0].clone(),
            self.context.window_exprs.clone(),
            self.context.partition_spec.clone(),
            self.context.order_spec.clone(),
            self.context.group_limit,
            self.context.output_window_cols,
        )?))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        // combine WindowGroupLimitExec -> WindowExec
        if let Ok(window_group_limit) = downcast_any!(&self.input, WindowExec)
            && window_group_limit.window_context().group_limit.is_some()
        {
            let combined = Arc::new(Self {
                input: window_group_limit.input.clone(),
                context: Arc::new(WindowContext::try_new(
                    self.input.schema(),
                    self.context.window_exprs.clone(),
                    self.context.partition_spec.clone(),
                    self.context.order_spec.clone(),
                    window_group_limit.context.group_limit,
                    true,
                )?),
                metrics: self.metrics.clone(),
                props: OnceCell::new(),
            });
            return combined.execute(partition, context);
        }

        // Streaming rank-like/agg windows are evaluated batch by batch. Functions that
        // need complete partition context, such as cume_dist, are handled in
        // execute_window().
        let exec_ctx = ExecutionContext::new(context, partition, self.schema(), &self.metrics);
        let input = exec_ctx.execute_with_input_stats(&self.input)?;
        let coalesced = exec_ctx.coalesce_with_default_batch_size(input);
        let window_ctx = self.context.clone();
        execute_window(coalesced, exec_ctx, window_ctx)
    }

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

    fn statistics(&self) -> Result<Statistics> {
        todo!()
    }
}

fn execute_window(
    mut input: SendableRecordBatchStream,
    exec_ctx: Arc<ExecutionContext>,
    window_ctx: Arc<WindowContext>,
) -> Result<SendableRecordBatchStream> {
    // start processing input batches
    Ok(exec_ctx
        .clone()
        .output_with_sender("Window", |sender| async move {
            let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
            sender.exclude_time(&elapsed_compute);

            let mut processors = window_ctx
                .window_exprs
                .iter()
                .map(|expr: &WindowExpr| expr.create_processor(&window_ctx))
                .collect::<Result<Vec<_>>>()?;

            if window_ctx.requires_full_partition() {
                // Functions like percent_rank/lead need a complete window partition,
                // but can still stream partition-by-partition when partition_spec exists.
                if !window_ctx.has_partition() {
                    let mut staging_batches = vec![];
                    while let Some(batch) = input.next().await.transpose()? {
                        staging_batches.push(batch);
                    }

                    flush_window_batches(
                        &mut staging_batches,
                        &window_ctx,
                        &exec_ctx,
                        &elapsed_compute,
                        processors.as_mut_slice(),
                        &sender,
                    )
                    .await?;
                    return Ok(());
                }

                let mut staging_batches = vec![];
                let mut current_partition: Option<Box<[u8]>> = None;
                while let Some(batch) = input.next().await.transpose()? {
                    if batch.num_rows() == 0 {
                        continue;
                    }

                    let partition_rows = window_ctx.get_partition_rows(&batch)?;
                    let mut partition_start = 0usize;
                    while partition_start < batch.num_rows() {
                        let partition_key: Box<[u8]> =
                            partition_rows.row(partition_start).as_ref().into();
                        let mut partition_end = partition_start + 1;
                        while partition_end < batch.num_rows()
                            && partition_rows.row(partition_end).as_ref() == partition_key.as_ref()
                        {
                            partition_end += 1;
                        }

                        if current_partition.as_deref() != Some(partition_key.as_ref()) {
                            flush_window_batches(
                                &mut staging_batches,
                                &window_ctx,
                                &exec_ctx,
                                &elapsed_compute,
                                processors.as_mut_slice(),
                                &sender,
                            )
                            .await?;
                            current_partition = Some(partition_key);
                        }

                        staging_batches
                            .push(batch.slice(partition_start, partition_end - partition_start));
                        partition_start = partition_end;
                    }
                }

                flush_window_batches(
                    &mut staging_batches,
                    &window_ctx,
                    &exec_ctx,
                    &elapsed_compute,
                    processors.as_mut_slice(),
                    &sender,
                )
                .await?;
                return Ok(());
            }

            while let Some(batch) = input.next().await.transpose()? {
                let _timer = elapsed_compute.timer();
                let output_batch =
                    process_window_batch(batch, &window_ctx, processors.as_mut_slice())?;
                exec_ctx
                    .baseline_metrics()
                    .record_output(output_batch.num_rows());
                sender.send(output_batch).await;
            }
            Ok(())
        }))
}

async fn flush_window_batches(
    staging_batches: &mut Vec<RecordBatch>,
    window_ctx: &WindowContext,
    exec_ctx: &ExecutionContext,
    elapsed_compute: &Time,
    processors: &mut [Box<dyn WindowFunctionProcessor>],
    sender: &Arc<WrappedRecordBatchSender>,
) -> Result<()> {
    if staging_batches.is_empty() {
        return Ok(());
    }

    let _timer = elapsed_compute.timer();
    let input_batches = mem::take(staging_batches);
    let batch = concat_batches(&window_ctx.input_schema, &input_batches)?;
    let output_batch = process_window_batch(batch, window_ctx, processors)?;
    exec_ctx
        .baseline_metrics()
        .record_output(output_batch.num_rows());
    sender.send(output_batch).await;
    Ok(())
}

fn process_window_batch(
    mut batch: RecordBatch,
    window_ctx: &WindowContext,
    processors: &mut [Box<dyn WindowFunctionProcessor>],
) -> Result<RecordBatch> {
    let mut window_cols: Vec<ArrayRef> = processors
        .iter_mut()
        .map(|processor| processor.process_batch(window_ctx, &batch))
        .collect::<Result<_>>()?;

    if let Some(group_limit) = window_ctx.group_limit {
        assert_eq!(window_cols.len(), 1);
        let limited = arrow::compute::kernels::cmp::lt_eq(
            &window_cols[0],
            &Int32Array::new_scalar(group_limit as i32),
        )?;
        window_cols[0] = arrow::compute::filter(&window_cols[0], &limited)?;
        batch = arrow::compute::filter_record_batch(&batch, &limited)?;
    }

    let outputs: Vec<ArrayRef> = batch
        .columns()
        .iter()
        .cloned()
        .chain(if window_ctx.output_window_cols {
            window_cols
        } else {
            vec![]
        })
        .zip(window_ctx.output_schema.fields())
        .map(|(array, field)| {
            if array.data_type() != field.data_type() {
                return cast(&array, field.data_type());
            }
            Ok(array.clone())
        })
        .collect::<Result<_>>()?;
    Ok(RecordBatch::try_new_with_options(
        window_ctx.output_schema.clone(),
        outputs,
        &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
    )?)
}

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

    use arrow::{array::*, datatypes::*, record_batch::RecordBatch};
    use datafusion::{
        assert_batches_eq,
        common::{Result, ScalarValue},
        physical_expr::{
            PhysicalSortExpr,
            expressions::{Column, Literal},
        },
        physical_plan::{ExecutionPlan, test::TestMemoryExec},
        prelude::SessionContext,
    };

    use crate::{
        agg::AggFunction,
        window::{WindowExpr, WindowFunction, WindowRankType},
        window_exec::WindowExec,
    };

    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,
        )?))
    }

    fn build_nullable_utf8_table(
        a: (&str, &Vec<i32>),
        b: (&str, &Vec<i32>),
        c: (&str, &Vec<Option<&str>>),
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let schema = Schema::new(vec![
            Field::new(a.0, DataType::Int32, false),
            Field::new(b.0, DataType::Int32, false),
            Field::new(c.0, DataType::Utf8, true),
        ]);

        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(StringArray::from(c.1.clone())),
            ],
        )?;
        let schema = batch.schema();
        Ok(Arc::new(TestMemoryExec::try_new(
            &[vec![batch]],
            schema,
            None,
        )?))
    }

    #[tokio::test]
    async fn test_window() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        // test window
        let input = build_table(
            ("a1", &vec![1, 1, 1, 1, 2, 3, 3]),
            ("b1", &vec![1, 2, 2, 3, 4, 1, 1]),
            ("c1", &vec![0, 0, 0, 0, 0, 0, 0]),
        )?;
        let window_exprs = vec![
            WindowExpr::new(
                WindowFunction::RankLike(WindowRankType::RowNumber),
                vec![],
                Arc::new(Field::new("b1_row_number", DataType::Int32, false)),
                DataType::Int32,
            ),
            WindowExpr::new(
                WindowFunction::RankLike(WindowRankType::Rank),
                vec![],
                Arc::new(Field::new("b1_rank", DataType::Int32, false)),
                DataType::Int32,
            ),
            WindowExpr::new(
                WindowFunction::RankLike(WindowRankType::DenseRank),
                vec![],
                Arc::new(Field::new("b1_dense_rank", DataType::Int32, false)),
                DataType::Int32,
            ),
            WindowExpr::new(
                WindowFunction::Agg(AggFunction::Sum),
                vec![Arc::new(Column::new("b1", 1))],
                Arc::new(Field::new("b1_sum", DataType::Int64, false)),
                DataType::Int64,
            ),
        ];
        let window = Arc::new(WindowExec::try_new(
            input.clone(),
            window_exprs.clone(),
            vec![Arc::new(Column::new("a1", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("b1", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);
        let stream = window.execute(0, task_ctx.clone())?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+----+----+----+---------------+---------+---------------+--------+",
            "| a1 | b1 | c1 | b1_row_number | b1_rank | b1_dense_rank | b1_sum |",
            "+----+----+----+---------------+---------+---------------+--------+",
            "| 1  | 1  | 0  | 1             | 1       | 1             | 1      |",
            "| 1  | 2  | 0  | 2             | 2       | 2             | 3      |",
            "| 1  | 2  | 0  | 3             | 2       | 2             | 5      |",
            "| 1  | 3  | 0  | 4             | 4       | 3             | 8      |",
            "| 2  | 4  | 0  | 1             | 1       | 1             | 4      |",
            "| 3  | 1  | 0  | 1             | 1       | 1             | 1      |",
            "| 3  | 1  | 0  | 2             | 1       | 1             | 2      |",
            "+----+----+----+---------------+---------+---------------+--------+",
        ];
        assert_batches_eq!(expected, &batches);

        // test window without partition by clause
        let input = build_table(
            ("a1", &vec![1, 3, 3, 1, 1, 1, 2]),
            ("b1", &vec![1, 1, 1, 2, 2, 3, 4]),
            ("c1", &vec![0, 0, 0, 0, 0, 0, 0]),
        )?;
        let window_exprs = vec![
            WindowExpr::new(
                WindowFunction::RankLike(WindowRankType::RowNumber),
                vec![],
                Arc::new(Field::new("b1_row_number", DataType::Int32, false)),
                DataType::Int32,
            ),
            WindowExpr::new(
                WindowFunction::RankLike(WindowRankType::Rank),
                vec![],
                Arc::new(Field::new("b1_rank", DataType::Int32, false)),
                DataType::Int32,
            ),
            WindowExpr::new(
                WindowFunction::RankLike(WindowRankType::DenseRank),
                vec![],
                Arc::new(Field::new("b1_dense_rank", DataType::Int32, false)),
                DataType::Int32,
            ),
            WindowExpr::new(
                WindowFunction::Agg(AggFunction::Sum),
                vec![Arc::new(Column::new("b1", 1))],
                Arc::new(Field::new("b1_sum", DataType::Int64, false)),
                DataType::Int64,
            ),
        ];
        let window = Arc::new(WindowExec::try_new(
            input.clone(),
            window_exprs.clone(),
            vec![],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("b1", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);
        let stream = window.execute(0, task_ctx.clone())?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+----+----+----+---------------+---------+---------------+--------+",
            "| a1 | b1 | c1 | b1_row_number | b1_rank | b1_dense_rank | b1_sum |",
            "+----+----+----+---------------+---------+---------------+--------+",
            "| 1  | 1  | 0  | 1             | 1       | 1             | 1      |",
            "| 3  | 1  | 0  | 2             | 1       | 1             | 2      |",
            "| 3  | 1  | 0  | 3             | 1       | 1             | 3      |",
            "| 1  | 2  | 0  | 4             | 4       | 2             | 5      |",
            "| 1  | 2  | 0  | 5             | 4       | 2             | 7      |",
            "| 1  | 3  | 0  | 6             | 6       | 3             | 10     |",
            "| 2  | 4  | 0  | 7             | 7       | 4             | 14     |",
            "+----+----+----+---------------+---------+---------------+--------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }

    #[tokio::test]
    async fn test_nth_value_window() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        let input = build_nullable_utf8_table(
            ("grp", &vec![1, 1, 1, 2, 2]),
            ("id", &vec![1, 2, 3, 1, 2]),
            ("v", &vec![None, Some("b"), Some("c"), Some("x"), None]),
        )?;
        let window_exprs = vec![
            WindowExpr::new(
                WindowFunction::NthValue {
                    ignore_nulls: false,
                },
                vec![
                    Arc::new(Column::new("v", 2)),
                    Arc::new(Literal::new(ScalarValue::Int32(Some(2)))),
                ],
                Arc::new(Field::new("nth_value_all", DataType::Utf8, true)),
                DataType::Utf8,
            ),
            WindowExpr::new(
                WindowFunction::NthValue { ignore_nulls: true },
                vec![
                    Arc::new(Column::new("v", 2)),
                    Arc::new(Literal::new(ScalarValue::Int32(Some(2)))),
                ],
                Arc::new(Field::new("nth_value_ignore_nulls", DataType::Utf8, true)),
                DataType::Utf8,
            ),
        ];
        let window = Arc::new(WindowExec::try_new(
            input,
            window_exprs,
            vec![Arc::new(Column::new("grp", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("id", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);
        let stream = window.execute(0, task_ctx)?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+-----+----+---+---------------+------------------------+",
            "| grp | id | v | nth_value_all | nth_value_ignore_nulls |",
            "+-----+----+---+---------------+------------------------+",
            "| 1   | 1  |   |               |                        |",
            "| 1   | 2  | b | b             |                        |",
            "| 1   | 3  | c | b             | c                      |",
            "| 2   | 1  | x |               |                        |",
            "| 2   | 2  |   |               |                        |",
            "+-----+----+---+---------------+------------------------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }

    #[tokio::test]
    async fn test_percent_rank_window() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        let input = build_table(
            ("grp", &vec![1, 1, 1, 2]),
            ("id", &vec![1, 1, 2, 5]),
            ("v", &vec![10, 20, 30, 40]),
        )?;
        let window_exprs = vec![WindowExpr::new(
            WindowFunction::PercentRank,
            vec![],
            Arc::new(Field::new("percent_rank", DataType::Float64, false)),
            DataType::Float64,
        )];
        let window = Arc::new(WindowExec::try_new(
            input,
            window_exprs,
            vec![Arc::new(Column::new("grp", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("id", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);
        let stream = window.execute(0, task_ctx)?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+-----+----+----+--------------+",
            "| grp | id | v  | percent_rank |",
            "+-----+----+----+--------------+",
            "| 1   | 1  | 10 | 0.0          |",
            "| 1   | 1  | 20 | 0.0          |",
            "| 1   | 2  | 30 | 1.0          |",
            "| 2   | 5  | 40 | 0.0          |",
            "+-----+----+----+--------------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }

    #[tokio::test]
    async fn test_percent_rank_window_across_batches() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        let batch1 = build_table_i32(
            ("grp", &vec![1, 1]),
            ("id", &vec![1, 2]),
            ("v", &vec![10, 20]),
        )?;
        let batch2 = build_table_i32(
            ("grp", &vec![1, 1, 1, 2]),
            ("id", &vec![3, 4, 5, 1]),
            ("v", &vec![30, 40, 50, 60]),
        )?;
        let schema = batch1.schema();
        let input = Arc::new(TestMemoryExec::try_new(
            &[vec![batch1, batch2]],
            schema,
            None,
        )?);

        let window_exprs = vec![WindowExpr::new(
            WindowFunction::PercentRank,
            vec![],
            Arc::new(Field::new("percent_rank", DataType::Float64, false)),
            DataType::Float64,
        )];
        let window = Arc::new(WindowExec::try_new(
            input,
            window_exprs,
            vec![Arc::new(Column::new("grp", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("id", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);

        let stream = window.execute(0, task_ctx)?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+-----+----+----+--------------+",
            "| grp | id | v  | percent_rank |",
            "+-----+----+----+--------------+",
            "| 1   | 1  | 10 | 0.0          |",
            "| 1   | 2  | 20 | 0.25         |",
            "| 1   | 3  | 30 | 0.5          |",
            "| 1   | 4  | 40 | 0.75         |",
            "| 1   | 5  | 50 | 1.0          |",
            "| 2   | 1  | 60 | 0.0          |",
            "+-----+----+----+--------------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }

    #[tokio::test]
    async fn test_window_group_limit() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        // test window
        let input = build_table(
            ("a1", &vec![1, 1, 1, 1, 2, 3, 3]),
            ("b1", &vec![1, 2, 2, 3, 4, 1, 1]),
            ("c1", &vec![0, 0, 0, 0, 0, 0, 0]),
        )?;
        let window_exprs = vec![WindowExpr::new(
            WindowFunction::RankLike(WindowRankType::RowNumber),
            vec![],
            Arc::new(Field::new("b1_row_number", DataType::Int32, false)),
            DataType::Int32,
        )];
        let window = Arc::new(WindowExec::try_new(
            input.clone(),
            window_exprs.clone(),
            vec![Arc::new(Column::new("a1", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("b1", 1)),
                options: Default::default(),
            }],
            Some(2),
            true,
        )?);
        let stream = window.execute(0, task_ctx.clone())?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+----+----+----+---------------+",
            "| a1 | b1 | c1 | b1_row_number |",
            "+----+----+----+---------------+",
            "| 1  | 1  | 0  | 1             |",
            "| 1  | 2  | 0  | 2             |",
            "| 2  | 4  | 0  | 1             |",
            "| 3  | 1  | 0  | 1             |",
            "| 3  | 1  | 0  | 2             |",
            "+----+----+----+---------------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }

    #[tokio::test]
    async fn test_window_lead_across_batches() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        let batch1 = build_table_i32(
            ("a1", &vec![1, 1]),
            ("b1", &vec![10, 20]),
            ("c1", &vec![0, 0]),
        )?;
        let batch2 = build_table_i32(
            ("a1", &vec![1, 2]),
            ("b1", &vec![30, 40]),
            ("c1", &vec![0, 0]),
        )?;
        let schema = batch1.schema();
        let input = Arc::new(TestMemoryExec::try_new(
            &[vec![batch1, batch2]],
            schema,
            None,
        )?);

        let window_exprs = vec![WindowExpr::new(
            WindowFunction::Lead,
            vec![
                Arc::new(Column::new("b1", 1)),
                Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
                Arc::new(Literal::new(ScalarValue::Int32(Some(-1)))),
            ],
            Arc::new(Field::new("b1_lead", DataType::Int32, false)),
            DataType::Int32,
        )];

        let window = Arc::new(WindowExec::try_new(
            input,
            window_exprs,
            vec![Arc::new(Column::new("a1", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("b1", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);

        let stream = window.execute(0, task_ctx)?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+----+----+----+---------+",
            "| a1 | b1 | c1 | b1_lead |",
            "+----+----+----+---------+",
            "| 1  | 10 | 0  | 20      |",
            "| 1  | 20 | 0  | 30      |",
            "| 1  | 30 | 0  | -1      |",
            "| 2  | 40 | 0  | -1      |",
            "+----+----+----+---------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }

    #[tokio::test]
    async fn test_cume_dist_window() -> Result<(), Box<dyn std::error::Error>> {
        let session_ctx = SessionContext::new();
        let task_ctx = session_ctx.task_ctx();

        let input = build_table(
            ("grp", &vec![1, 1, 1, 2]),
            ("id", &vec![1, 1, 2, 5]),
            ("v", &vec![10, 20, 30, 40]),
        )?;
        let window_exprs = vec![WindowExpr::new(
            WindowFunction::CumeDist,
            vec![],
            Arc::new(Field::new("cume_dist", DataType::Float64, false)),
            DataType::Float64,
        )];
        let window = Arc::new(WindowExec::try_new(
            input,
            window_exprs,
            vec![Arc::new(Column::new("grp", 0))],
            vec![PhysicalSortExpr {
                expr: Arc::new(Column::new("id", 1)),
                options: Default::default(),
            }],
            None,
            true,
        )?);
        let stream = window.execute(0, task_ctx)?;
        let batches = datafusion::physical_plan::common::collect(stream).await?;
        let expected = vec![
            "+-----+----+----+--------------------+",
            "| grp | id | v  | cume_dist          |",
            "+-----+----+----+--------------------+",
            "| 1   | 1  | 10 | 0.6666666666666666 |",
            "| 1   | 1  | 20 | 0.6666666666666666 |",
            "| 1   | 2  | 30 | 1.0                |",
            "| 2   | 5  | 40 | 1.0                |",
            "+-----+----+----+--------------------+",
        ];
        assert_batches_eq!(expected, &batches);
        Ok(())
    }
}
