datafusion_physical_plan/windows/

bounded_window_agg_exec.rs

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

//! Stream and channel implementations for window function expressions.
//! The executor given here uses bounded memory (does not maintain all
//! the input data seen so far), which makes it appropriate when processing
//! infinite inputs.

use std::cmp::{Ordering, min};
use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use super::utils::create_schema;
use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use crate::statistics::{ChildStats, StatisticsArgs};
use crate::stream::EmptyRecordBatchStream;
use crate::windows::{
    calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs,
    window_equivalence_properties,
};
use crate::{
    ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution,
    ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements,
    InputOrderMode, PlanProperties, RecordBatchStream, ReplaceChildrenOptions,
    SendableRecordBatchStream, Statistics, WindowExpr, validate_child_count,
};

use arrow::compute::take_record_batch;
use arrow::{
    array::{Array, ArrayRef, RecordBatchOptions, UInt32Array, UInt32Builder},
    compute::{concat, concat_batches, sort_to_indices, take_arrays},
    datatypes::SchemaRef,
    record_batch::RecordBatch,
};
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::utils::{
    evaluate_partition_ranges, get_at_indices, get_row_at_idx,
};
use datafusion_common::{
    HashMap, Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err,
};
use datafusion_execution::TaskContext;
use datafusion_expr::ColumnarValue;
use datafusion_expr::window_state::{PartitionBatchState, WindowAggState};
use datafusion_physical_expr::window::{
    PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowEvalContext,
    WindowState,
};
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use datafusion_physical_expr_common::sort_expr::{
    OrderingRequirements, PhysicalSortExpr,
};

use crate::execution_plan::CardinalityEffect;
use datafusion_common::hash_utils::RandomState;
use futures::stream::Stream;
use futures::{StreamExt, ready};
use hashbrown::hash_table::HashTable;
use indexmap::IndexMap;
use log::debug;

/// Callback receiver for per-partition window state.
///
/// `state` is the result of [`Accumulator::state`], which is a `&mut self`
/// call whose trait doc states "this function should not be called twice."
/// Several built-in aggregates (`median`, `percentile_cont`, `string_agg`,
/// `min_max_bytes`/`min_max_struct`) `std::mem::take` their internal
/// buffers to build that state — so `state` is a destructive read, not a
/// snapshot. The exec fires this at most once per group; a callee that
/// needs the value beyond the callback must retain it (e.g. clone into
/// owned storage).
///
/// [`Accumulator::state`]: datafusion_expr::Accumulator::state
pub trait WindowStateObserver: Send + Sync {
    /// Invoked once per (output-partition-index, window-expression,
    /// PARTITION BY tuple) as each PARTITION BY group closes, for every
    /// aggregate window expression on the exec. Non-aggregate window
    /// functions (e.g. `row_number`, `rank`, `lead`/`lag`) do not fire this
    /// callback.
    ///
    /// # Arguments
    ///
    /// * `partition_idx` - Output partition index of the [`BoundedWindowAggExec`]
    ///   stream firing this callback.
    /// * `window_expr` - The window expression whose state just closed.
    /// * `partition_key` - The PARTITION BY tuple that just closed.
    /// * `state` - [`Accumulator::state`] for the closed group of
    ///   `window_expr`. See the trait-level doc for the destructive-read
    ///   contract.
    ///
    /// [`Accumulator::state`]: datafusion_expr::Accumulator::state
    fn finalize_window_aggregate(
        &self,
        partition_idx: usize,
        window_expr: &Arc<dyn WindowExpr>,
        partition_key: &PartitionKey,
        state: Vec<ScalarValue>,
    ) -> Result<()>;
}

/// Window execution plan
#[derive(Clone)]
pub struct BoundedWindowAggExec {
    /// Input plan
    input: Arc<dyn ExecutionPlan>,
    /// Window function expression
    window_expr: Vec<Arc<dyn WindowExpr>>,
    /// Schema after the window is run
    schema: SchemaRef,
    /// Execution metrics
    metrics: ExecutionPlanMetricsSet,
    /// Describes how the input is ordered relative to the partition keys
    pub input_order_mode: InputOrderMode,
    /// Partition by indices that define ordering
    // For example, if input ordering is ORDER BY a, b and window expression
    // contains PARTITION BY b, a; `ordered_partition_by_indices` would be 1, 0.
    // Similarly, if window expression contains PARTITION BY a, b; then
    // `ordered_partition_by_indices` would be 0, 1.
    // See `get_ordered_partition_by_indices` for more details.
    ordered_partition_by_indices: Vec<usize>,
    /// Cache holding plan properties like equivalences, output partitioning etc.
    cache: Arc<PlanProperties>,
    /// If `can_rerepartition` is false, partition_keys is always empty.
    can_repartition: bool,
    /// Invoked at partition-close to publish finalized per-partition window
    /// state. Storage and multi-group handling are the caller's; the exec is
    /// a pure event source.
    state_observer: Option<Arc<dyn WindowStateObserver>>,
}

impl std::fmt::Debug for BoundedWindowAggExec {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BoundedWindowAggExec")
            .field("input", &self.input)
            .field("window_expr", &self.window_expr)
            .field("schema", &self.schema)
            .field("metrics", &self.metrics)
            .field("input_order_mode", &self.input_order_mode)
            .field(
                "ordered_partition_by_indices",
                &self.ordered_partition_by_indices,
            )
            .field("cache", &self.cache)
            .field("can_repartition", &self.can_repartition)
            .field(
                "state_observer",
                &self.state_observer.as_ref().map(|_| "..."),
            )
            .finish()
    }
}

impl BoundedWindowAggExec {
    /// Create a new execution plan for window aggregates
    pub fn try_new(
        window_expr: Vec<Arc<dyn WindowExpr>>,
        input: Arc<dyn ExecutionPlan>,
        input_order_mode: InputOrderMode,
        can_repartition: bool,
    ) -> Result<Self> {
        let schema = create_schema(&input.schema(), &window_expr)?;
        let schema = Arc::new(schema);
        let partition_by_exprs = window_expr[0].partition_by();
        let ordered_partition_by_indices = match &input_order_mode {
            InputOrderMode::Sorted => {
                let indices = get_ordered_partition_by_indices(
                    window_expr[0].partition_by(),
                    &input,
                )?;
                if indices.len() == partition_by_exprs.len() {
                    indices
                } else {
                    (0..partition_by_exprs.len()).collect::<Vec<_>>()
                }
            }
            InputOrderMode::PartiallySorted(ordered_indices) => ordered_indices.clone(),
            InputOrderMode::Linear => {
                vec![]
            }
        };
        let cache = Self::compute_properties(&input, &schema, &window_expr)?;
        Ok(Self {
            input,
            window_expr,
            schema,
            metrics: ExecutionPlanMetricsSet::new(),
            input_order_mode,
            ordered_partition_by_indices,
            cache: Arc::new(cache),
            can_repartition,
            state_observer: None,
        })
    }

    /// Install (or clear) a [`WindowStateObserver`] that receives each
    /// PARTITION BY group's finalized window state at partition close.
    ///
    /// Errors when `observer` is `Some` and any window expression on this
    /// exec has a non-ever-expanding frame (i.e. its start bound is not
    /// `UNBOUNDED PRECEDING`). Those frames use `SlidingAggregateWindowExpr`
    /// under the hood, whose accumulator calls `retract_batch` — at
    /// partition close the accumulator holds only the last frame's rows,
    /// not the partition aggregate, so the observed state would silently
    /// misrepresent the group.
    pub fn with_state_observer(
        mut self,
        observer: Option<Arc<dyn WindowStateObserver>>,
    ) -> Result<Self> {
        if observer.is_some() {
            for expr in &self.window_expr {
                if !expr.get_window_frame().is_ever_expanding() {
                    return exec_err!(
                        "cannot install WindowStateObserver on BoundedWindowAggExec \
                         with a sliding aggregate window frame (start != \
                         UNBOUNDED PRECEDING) for `{}`; sliding accumulator state \
                         is frame-only, not the partition aggregate",
                        expr.name()
                    );
                }
            }
        }
        self.state_observer = observer;
        Ok(self)
    }

    /// The currently-installed [`WindowStateObserver`], if any. Optimizer
    /// rules that rebuild this exec via
    /// [`crate::windows::get_best_fitting_window`] or a direct `try_new`
    /// call must read this and reinstall it on the new exec, otherwise a
    /// caller-installed observer is silently dropped by the rewrite.
    pub fn state_observer(&self) -> Option<&Arc<dyn WindowStateObserver>> {
        self.state_observer.as_ref()
    }

    /// Window expressions
    pub fn window_expr(&self) -> &[Arc<dyn WindowExpr>] {
        &self.window_expr
    }

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

    /// Return the output sort order of partition keys: For example
    /// OVER(PARTITION BY a, ORDER BY b) -> would give sorting of the column a
    // We are sure that partition by columns are always at the beginning of sort_keys
    // Hence returned `PhysicalSortExpr` corresponding to `PARTITION BY` columns can be used safely
    // to calculate partition separation points
    pub fn partition_by_sort_keys(&self) -> Result<Vec<PhysicalSortExpr>> {
        let partition_by = self.window_expr()[0].partition_by();
        get_partition_by_sort_exprs(
            &self.input,
            partition_by,
            &self.ordered_partition_by_indices,
        )
    }

    /// Initializes the appropriate [`PartitionSearcher`] implementation from
    /// the state.
    fn get_search_algo(&self) -> Result<Box<dyn PartitionSearcher>> {
        let partition_by_sort_keys = self.partition_by_sort_keys()?;
        let ordered_partition_by_indices = self.ordered_partition_by_indices.clone();
        let input_schema = self.input().schema();
        Ok(match &self.input_order_mode {
            InputOrderMode::Sorted => {
                // In Sorted mode, all partition by columns should be ordered.
                if self.window_expr()[0].partition_by().len()
                    != ordered_partition_by_indices.len()
                {
                    return exec_err!(
                        "All partition by columns should have an ordering in Sorted mode."
                    );
                }
                Box::new(SortedSearch {
                    partition_by_sort_keys,
                    ordered_partition_by_indices,
                    input_schema,
                })
            }
            InputOrderMode::Linear | InputOrderMode::PartiallySorted(_) => Box::new(
                LinearSearch::new(ordered_partition_by_indices, input_schema),
            ),
        })
    }

    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
    fn compute_properties(
        input: &Arc<dyn ExecutionPlan>,
        schema: &SchemaRef,
        window_exprs: &[Arc<dyn WindowExpr>],
    ) -> Result<PlanProperties> {
        // Calculate equivalence properties:
        let eq_properties = window_equivalence_properties(schema, input, window_exprs)?;

        // As we can have repartitioning using the partition keys, this can
        // be either one or more than one, depending on the presence of
        // repartitioning.
        let output_partitioning = input.output_partitioning().clone();

        // Construct properties cache
        Ok(PlanProperties::new(
            eq_properties,
            output_partitioning,
            // TODO: Emission type and boundedness information can be enhanced here
            input.pipeline_behavior(),
            input.boundedness(),
        ))
    }

    pub fn partition_keys(&self) -> Vec<Arc<dyn PhysicalExpr>> {
        if !self.can_repartition {
            vec![]
        } else {
            let all_partition_keys = self
                .window_expr()
                .iter()
                .map(|expr| expr.partition_by().to_vec())
                .collect::<Vec<_>>();

            all_partition_keys
                .into_iter()
                .min_by_key(|s| s.len())
                .unwrap_or_else(Vec::new)
        }
    }

    fn statistics_helper(&self, statistics: Statistics) -> Result<Statistics> {
        let win_cols = self.window_expr.len();
        let input_cols = self.input.schema().fields().len();
        // TODO stats: some windowing function will maintain invariants such as min, max...
        let mut column_statistics = Vec::with_capacity(win_cols + input_cols);
        // copy stats of the input to the beginning of the schema.
        column_statistics.extend(statistics.column_statistics);
        for _ in 0..win_cols {
            column_statistics.push(ColumnStatistics::new_unknown())
        }
        Ok(Statistics {
            num_rows: statistics.num_rows,
            column_statistics,
            total_byte_size: Precision::Absent,
        })
    }
}

impl DisplayAs for BoundedWindowAggExec {
    fn fmt_as(
        &self,
        t: DisplayFormatType,
        f: &mut std::fmt::Formatter,
    ) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, "BoundedWindowAggExec: ")?;
                let g: Vec<String> = self
                    .window_expr
                    .iter()
                    .map(|e| {
                        let field = match e.field() {
                            Ok(f) => f.to_string(),
                            Err(e) => format!("{e:?}"),
                        };
                        format!(
                            "{}: {}, frame: {}",
                            e.name().to_owned(),
                            field,
                            e.get_window_frame()
                        )
                    })
                    .collect();
                let mode = &self.input_order_mode;
                write!(f, "wdw=[{}], mode=[{:?}]", g.join(", "), mode)?;
            }
            DisplayFormatType::TreeRender => {
                let g: Vec<String> = self
                    .window_expr
                    .iter()
                    .map(|e| e.name().to_owned().to_string())
                    .collect();
                writeln!(f, "select_list={}", g.join(", "))?;

                let mode = &self.input_order_mode;
                writeln!(f, "mode={mode:?}")?;
            }
        }
        Ok(())
    }
}

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

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

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

    fn apply_expressions(
        &self,
        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
    ) -> Result<TreeNodeRecursion> {
        let expressions = self.window_expr.iter().flat_map(|window_expr| {
            let expressions = window_expr.all_expressions();
            expressions
                .args
                .into_iter()
                .chain(expressions.partition_by_exprs)
                .chain(expressions.order_by_exprs)
        });
        crate::apply_expression_roots(expressions, f)
    }

    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
        let partition_bys = self.window_expr()[0].partition_by();
        let order_keys = self.window_expr()[0].order_by();
        let partition_bys = self
            .ordered_partition_by_indices
            .iter()
            .map(|idx| &partition_bys[*idx]);
        vec![calc_requirements(partition_bys, order_keys)]
    }

    fn required_input_distribution(&self) -> Vec<Distribution> {
        self.input_distribution_requirements().into_per_child()
    }

    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
        if self.partition_keys().is_empty() {
            debug!("No partition defined for BoundedWindowAggExec!!!");
            InputDistributionRequirements::new(vec![Distribution::SinglePartition])
        } else {
            InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
                self.partition_keys(),
            )])
        }
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        vec![true]
    }

    fn replace_children(
        self: Arc<Self>,
        mut children: Vec<Arc<dyn ExecutionPlan>>,
        options: ReplaceChildrenOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        validate_child_count!(self, children);
        match options.children_properties {
            ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
                input: children.swap_remove(0),
                metrics: ExecutionPlanMetricsSet::new(),
                ..Self::clone(&*self)
            })),
            ChildrenPropertiesMode::Recompute => {
                let new = BoundedWindowAggExec::try_new(
                    self.window_expr.clone(),
                    Arc::clone(&children[0]),
                    self.input_order_mode.clone(),
                    self.can_repartition,
                )?
                .with_state_observer(self.state_observer.clone())?;
                Ok(Arc::new(new))
            }
        }
    }

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

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

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        let input = self.input.execute(partition, context)?;
        let search_mode = self.get_search_algo()?;
        let stream = Box::pin(BoundedWindowAggStream::new(
            Arc::clone(&self.schema),
            self.window_expr.clone(),
            input,
            BaselineMetrics::new(&self.metrics, partition),
            search_mode,
            partition,
            self.state_observer.clone(),
        )?);
        Ok(stream)
    }

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

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

    fn statistics_from_inputs(
        &self,
        input_stats: &[Arc<Statistics>],
        _args: &StatisticsArgs,
    ) -> Result<Arc<Statistics>> {
        let input_stat = input_stats[0].as_ref().clone();
        Ok(Arc::new(self.statistics_helper(input_stat)?))
    }

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

    #[cfg(feature = "proto")]
    fn try_to_proto(
        &self,
        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
        use super::proto::encode_physical_window_expr;
        use datafusion_proto_common::protobuf_common::EmptyMessage;
        use datafusion_proto_models::protobuf;
        use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode;

        // Exhaustive destructure: adding a field to `BoundedWindowAggExec`
        // without deciding how it is serialized is a compile error, not a
        // silent round-trip gap.
        let Self {
            input,
            window_expr,
            // Derived at construction by `create_schema` from the input schema
            // and the window expressions.
            schema: _,
            // Runtime execution state, rebuilt empty on decode.
            metrics: _,
            input_order_mode,
            // Derived at construction from `input_order_mode` and the window
            // expressions' PARTITION BY.
            ordered_partition_by_indices: _,
            // Derived at construction by `Self::compute_properties`.
            cache: _,
            // No wire field of its own; it is folded into `partition_keys`
            // below, since `partition_keys()` returns an empty vec when this is
            // false and the decoder recovers it as `!partition_keys.is_empty()`.
            can_repartition: _,
            // Runtime callback installed after planning; not part of the wire
            // format. Any decoder that needs it must reinstall via
            // `with_state_observer`.
            state_observer: _,
        } = self;

        let input = ctx.encode_child(input)?;
        let window_expr = window_expr
            .iter()
            .map(|expr| encode_physical_window_expr(expr, ctx))
            .collect::<Result<Vec<_>>>()?;
        let partition_keys = self
            .partition_keys()
            .iter()
            .map(|expr| ctx.encode_expr(expr))
            .collect::<Result<Vec<_>>>()?;
        // A `Some(input_order_mode)` is what tells the shared `Window` decode
        // arm to rebuild a `BoundedWindowAggExec` rather than a `WindowAggExec`.
        let input_order_mode = match input_order_mode {
            InputOrderMode::Linear => ProtoInputOrderMode::Linear(EmptyMessage {}),
            InputOrderMode::PartiallySorted(columns) => {
                ProtoInputOrderMode::PartiallySorted(
                    protobuf::PartiallySortedInputOrderMode {
                        columns: columns.iter().map(|column| *column as u64).collect(),
                    },
                )
            }
            InputOrderMode::Sorted => ProtoInputOrderMode::Sorted(EmptyMessage {}),
        };

        Ok(Some(protobuf::PhysicalPlanNode {
            physical_plan_type: Some(
                protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new(
                    protobuf::WindowAggExecNode {
                        input: Some(Box::new(input)),
                        window_expr,
                        partition_keys,
                        input_order_mode: Some(input_order_mode),
                    },
                )),
            ),
        }))
    }
}

/// Trait that specifies how we search for (or calculate) partitions. It has two
/// implementations: [`SortedSearch`] and [`LinearSearch`].
trait PartitionSearcher: Send {
    /// This method constructs output columns using the result of each window expression
    /// (each entry in the output vector comes from a window expression).
    /// Executor when producing output concatenates `input_buffer` (corresponding section), and
    /// result of this function to generate output `RecordBatch`. `input_buffer` is used to determine
    /// which sections of the window expression results should be used to generate output.
    /// `partition_buffers` contains corresponding section of the `RecordBatch` for each partition.
    /// `window_agg_states` stores per partition state for each window expression.
    /// None case means that no result is generated
    /// `Some(Vec<ArrayRef>)` is the result of each window expression.
    fn calculate_out_columns(
        &mut self,
        input_buffer: &RecordBatch,
        window_agg_states: &[PartitionWindowAggStates],
        partition_buffers: &mut PartitionBatches,
        window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Option<Vec<ArrayRef>>>;

    /// Determine whether `[InputOrderMode]` is `[InputOrderMode::Linear]` or not.
    fn is_mode_linear(&self) -> bool {
        false
    }

    // Constructs corresponding batches for each partition for the record_batch.
    fn evaluate_partition_batches(
        &mut self,
        record_batch: &RecordBatch,
        window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Vec<(PartitionKey, RecordBatch)>>;

    /// Prunes the state.
    fn prune(&mut self, _n_out: usize) {}

    /// Marks the partition as done if we are sure that corresponding partition
    /// cannot receive any more values.
    fn mark_partition_end(&self, partition_buffers: &mut PartitionBatches);

    /// Updates `input_buffer` and `partition_buffers` with the new `record_batch`.
    fn update_partition_batch(
        &mut self,
        input_buffer: &mut RecordBatch,
        record_batch: RecordBatch,
        window_expr: &[Arc<dyn WindowExpr>],
        partition_buffers: &mut PartitionBatches,
    ) -> Result<()> {
        if record_batch.num_rows() == 0 {
            return Ok(());
        }
        let partition_batches =
            self.evaluate_partition_batches(&record_batch, window_expr)?;
        for (partition_row, partition_batch) in partition_batches {
            if let Some(partition_batch_state) = partition_buffers.get_mut(&partition_row)
            {
                partition_batch_state.extend(&partition_batch)?
            } else {
                let options = RecordBatchOptions::new()
                    .with_row_count(Some(partition_batch.num_rows()));
                // Use input_schema for the buffer schema, not `record_batch.schema()`
                // as it may not have the "correct" schema in terms of output
                // nullability constraints. For details, see the following issue:
                // https://github.com/apache/datafusion/issues/9320
                let partition_batch = RecordBatch::try_new_with_options(
                    Arc::clone(self.input_schema()),
                    partition_batch.columns().to_vec(),
                    &options,
                )?;
                let partition_batch_state =
                    PartitionBatchState::new_with_batch(partition_batch);
                partition_buffers.insert(partition_row, partition_batch_state);
            }
        }

        self.mark_partition_end(partition_buffers);

        *input_buffer = if input_buffer.num_rows() == 0 {
            record_batch
        } else {
            concat_batches(self.input_schema(), [input_buffer, &record_batch])?
        };

        Ok(())
    }

    fn input_schema(&self) -> &SchemaRef;
}

/// This object encapsulates the algorithm state for a simple linear scan
/// algorithm for computing partitions.
pub struct LinearSearch {
    /// Keeps the hash of input buffer calculated from PARTITION BY columns.
    /// Its length is equal to the `input_buffer` length.
    input_buffer_hashes: VecDeque<u64>,
    /// Used during hash value calculation.
    random_state: RandomState,
    /// Input ordering and partition by key ordering need not be the same, so
    /// this vector stores the mapping between them. For instance, if the input
    /// is ordered by a, b and the window expression contains a PARTITION BY b, a
    /// clause, this attribute stores [1, 0].
    ordered_partition_by_indices: Vec<usize>,
    /// We use this [`HashTable`] to calculate unique partitions for each new
    /// RecordBatch. First entry in the tuple is the hash value, the second
    /// entry is the unique ID for each partition (increments from 0 to n).
    row_map_batch: HashTable<(u64, usize)>,
    /// We use this [`HashTable`] to calculate the output columns that we can
    /// produce at each cycle. First entry in the tuple is the hash value, the
    /// second entry is the unique ID for each partition (increments from 0 to n).
    /// The third entry stores how many new outputs are calculated for the
    /// corresponding partition.
    row_map_out: HashTable<(u64, usize, usize)>,
    input_schema: SchemaRef,
}

impl PartitionSearcher for LinearSearch {
    /// This method constructs output columns using the result of each window expression.
    // Assume input buffer is         |      Partition Buffers would be (Where each partition and its data is separated)
    // a, 2                           |      a, 2
    // b, 2                           |      a, 2
    // a, 2                           |      a, 2
    // b, 2                           |
    // a, 2                           |      b, 2
    // b, 2                           |      b, 2
    // b, 2                           |      b, 2
    //                                |      b, 2
    // Also assume we happen to calculate 2 new values for a, and 3 for b (To be calculate missing values we may need to consider future values).
    // Partition buffers effectively will be
    // a, 2, 1
    // a, 2, 2
    // a, 2, (missing)
    //
    // b, 2, 1
    // b, 2, 2
    // b, 2, 3
    // b, 2, (missing)
    // When partition buffers are mapped back to the original record batch. Result becomes
    // a, 2, 1
    // b, 2, 1
    // a, 2, 2
    // b, 2, 2
    // a, 2, (missing)
    // b, 2, 3
    // b, 2, (missing)
    // This function calculates the column result of window expression(s) (First 4 entry of 3rd column in the above section.)
    // 1
    // 1
    // 2
    // 2
    // Above section corresponds to calculated result which can be emitted without breaking input buffer ordering.
    fn calculate_out_columns(
        &mut self,
        input_buffer: &RecordBatch,
        window_agg_states: &[PartitionWindowAggStates],
        partition_buffers: &mut PartitionBatches,
        window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Option<Vec<ArrayRef>>> {
        let partition_output_indices = self.calc_partition_output_indices(
            input_buffer,
            window_agg_states,
            window_expr,
        )?;

        let n_window_col = window_agg_states.len();
        let mut new_columns = vec![vec![]; n_window_col];
        // Size of all_indices can be at most input_buffer.num_rows():
        let mut all_indices = UInt32Builder::with_capacity(input_buffer.num_rows());
        for (row, indices) in partition_output_indices {
            let length = indices.len();
            for (idx, window_agg_state) in window_agg_states.iter().enumerate() {
                let partition = &window_agg_state[&row];
                let values = Arc::clone(&partition.state.out_col.slice(0, length));
                new_columns[idx].push(values);
            }
            let partition_batch_state = &mut partition_buffers[&row];
            // Store how many rows are generated for each partition
            partition_batch_state.n_out_row = length;
            // For each row keep corresponding index in the input record batch
            all_indices.append_slice(&indices);
        }
        let all_indices = all_indices.finish();
        if all_indices.is_empty() {
            // We couldn't generate any new value, return early:
            return Ok(None);
        }

        // Concatenate results for each column by converting `Vec<Vec<ArrayRef>>`
        // to Vec<ArrayRef> where inner `Vec<ArrayRef>`s are converted to `ArrayRef`s.
        let new_columns = new_columns
            .iter()
            .map(|items| {
                concat(&items.iter().map(|e| e.as_ref()).collect::<Vec<_>>())
                    .map_err(|e| arrow_datafusion_err!(e))
            })
            .collect::<Result<Vec<_>>>()?;
        // We should emit columns according to row index ordering.
        let sorted_indices = sort_to_indices(&all_indices, None, None)?;
        // Construct new column according to row ordering. This fixes ordering
        take_arrays(&new_columns, &sorted_indices, None)
            .map(Some)
            .map_err(|e| arrow_datafusion_err!(e))
    }

    fn evaluate_partition_batches(
        &mut self,
        record_batch: &RecordBatch,
        window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Vec<(PartitionKey, RecordBatch)>> {
        let partition_bys =
            evaluate_partition_by_column_values(record_batch, window_expr)?;
        // NOTE: In Linear or PartiallySorted modes, we are sure that
        //       `partition_bys` are not empty.
        let (mut keys, permutation, bounds) =
            self.compute_partition_permutation(&partition_bys, record_batch)?;
        if keys.len() == 1 {
            // The batch contains a single partition, so the gather below
            // would be an identity permutation; use the batch as-is.
            let key = keys.remove(0);
            return Ok(vec![(key, record_batch.clone())]);
        }
        // Reorder the batch with a single `take` so that each partition's
        // rows become contiguous, then hand each partition a zero-copy slice
        // of the result. The slices share the gathered batch's buffers;
        // `PartitionBatchState::extend` copies out of them the next time the
        // partition receives rows.
        let gathered = take_record_batch(record_batch, &UInt32Array::from(permutation))?;
        Ok(keys
            .into_iter()
            .zip(bounds.windows(2))
            .map(|(key, bound)| (key, gathered.slice(bound[0], bound[1] - bound[0])))
            .collect())
    }

    fn prune(&mut self, n_out: usize) {
        // Delete hashes for the rows that are outputted.
        self.input_buffer_hashes.drain(0..n_out);
    }

    fn mark_partition_end(&self, partition_buffers: &mut PartitionBatches) {
        // We should be in the `PartiallySorted` case, otherwise we can not
        // tell when we are at the end of a given partition.
        if !self.ordered_partition_by_indices.is_empty()
            && let Some((last_row, _)) = partition_buffers.last()
        {
            let last_sorted_cols = self
                .ordered_partition_by_indices
                .iter()
                .map(|idx| last_row[*idx].clone())
                .collect::<Vec<_>>();
            for (row, partition_batch_state) in partition_buffers.iter_mut() {
                let sorted_cols = self
                    .ordered_partition_by_indices
                    .iter()
                    .map(|idx| &row[*idx]);
                // All the partitions other than `last_sorted_cols` are done.
                // We are sure that we will no longer receive values for these
                // partitions (arrival of a new value would violate ordering).
                partition_batch_state.is_end = !sorted_cols.eq(&last_sorted_cols);
            }
        }
    }

    fn is_mode_linear(&self) -> bool {
        self.ordered_partition_by_indices.is_empty()
    }

    fn input_schema(&self) -> &SchemaRef {
        &self.input_schema
    }
}

impl LinearSearch {
    /// Initialize a new [`LinearSearch`] partition searcher.
    fn new(ordered_partition_by_indices: Vec<usize>, input_schema: SchemaRef) -> Self {
        LinearSearch {
            input_buffer_hashes: VecDeque::new(),
            random_state: Default::default(),
            ordered_partition_by_indices,
            row_map_batch: HashTable::with_capacity(256),
            row_map_out: HashTable::with_capacity(256),
            input_schema,
        }
    }

    /// Splits the rows of `batch` by partition, according to the PARTITION BY
    /// expression results in `columns`. Returns the distinct partition keys
    /// in first-appearance order, a permutation of the row indices of
    /// `batch` that groups each partition's rows together, and the
    /// boundaries of each partition's run of rows within that permutation:
    /// partition `p` occupies `permutation[bounds[p]..bounds[p + 1]]`, and
    /// its indices are in ascending (stream) order.
    fn compute_partition_permutation(
        &mut self,
        columns: &[ArrayRef],
        batch: &RecordBatch,
    ) -> Result<(Vec<PartitionKey>, Vec<u32>, Vec<usize>)> {
        let num_rows = batch.num_rows();
        let mut batch_hashes = vec![0; num_rows];
        create_hashes(columns, &self.random_state, &mut batch_hashes)?;
        self.input_buffer_hashes.extend(&batch_hashes);
        // reset row_map for new calculation
        self.row_map_batch.clear();
        let mut keys: Vec<PartitionKey> = vec![];
        // Partition id of each row, in row order:
        let mut row_partition_ids = Vec::with_capacity(num_rows);
        // Number of rows in each partition:
        let mut counts: Vec<usize> = vec![];
        for (hash, row_idx) in batch_hashes.into_iter().zip(0u32..) {
            let entry = self.row_map_batch.find_mut(hash, |(_, group_idx)| {
                let row = get_row_at_idx(columns, row_idx as usize).unwrap();
                // Handle hash collisions with an equality check:
                row == keys[*group_idx]
            });
            let group_idx = if let Some((_, group_idx)) = entry {
                *group_idx
            } else {
                let group_idx = keys.len();
                self.row_map_batch
                    .insert_unique(hash, (hash, group_idx), |(hash, _)| *hash);
                keys.push(get_row_at_idx(columns, row_idx as usize)?);
                counts.push(0);
                group_idx
            };
            row_partition_ids.push(group_idx);
            counts[group_idx] += 1;
        }
        // A prefix sum over the counts gives each partition's run boundaries
        // in the permutation.
        let mut bounds = Vec::with_capacity(counts.len() + 1);
        let mut total = 0;
        bounds.push(0);
        for count in counts {
            total += count;
            bounds.push(total);
        }
        // Scatter each row's index into its partition's run. Visiting rows
        // in ascending order keeps each run in ascending row order.
        let mut cursors: Vec<usize> = bounds[..bounds.len() - 1].to_vec();
        let mut permutation = vec![0u32; num_rows];
        for (row_idx, group_idx) in row_partition_ids.into_iter().enumerate() {
            permutation[cursors[group_idx]] = row_idx as u32;
            cursors[group_idx] += 1;
        }
        Ok((keys, permutation, bounds))
    }

    /// Calculates partition keys and result indices for each partition.
    /// The return value is a vector of tuples where the first entry stores
    /// the partition key (unique for each partition) and the second entry
    /// stores indices of the rows for which the partition is constructed.
    fn calc_partition_output_indices(
        &mut self,
        input_buffer: &RecordBatch,
        window_agg_states: &[PartitionWindowAggStates],
        window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Vec<(PartitionKey, Vec<u32>)>> {
        let partition_by_columns =
            evaluate_partition_by_column_values(input_buffer, window_expr)?;
        // Reset the row_map state:
        self.row_map_out.clear();
        let mut partition_indices: Vec<(PartitionKey, Vec<u32>)> = vec![];
        for (hash, row_idx) in self.input_buffer_hashes.iter().zip(0u32..) {
            let entry = self.row_map_out.find_mut(*hash, |(_, group_idx, _)| {
                let row =
                    get_row_at_idx(&partition_by_columns, row_idx as usize).unwrap();
                row == partition_indices[*group_idx].0
            });
            if let Some((_, group_idx, n_out)) = entry {
                let (_, indices) = &mut partition_indices[*group_idx];
                if indices.len() >= *n_out {
                    break;
                }
                indices.push(row_idx);
            } else {
                let row = get_row_at_idx(&partition_by_columns, row_idx as usize)?;
                let min_out = window_agg_states
                    .iter()
                    .map(|window_agg_state| {
                        window_agg_state
                            .get(&row)
                            .map(|partition| partition.state.out_col.len())
                            .unwrap_or(0)
                    })
                    .min()
                    .unwrap_or(0);
                if min_out == 0 {
                    break;
                }
                self.row_map_out.insert_unique(
                    *hash,
                    (*hash, partition_indices.len(), min_out),
                    |(hash, _, _)| *hash,
                );
                partition_indices.push((row, vec![row_idx]));
            }
        }
        Ok(partition_indices)
    }
}

/// This object encapsulates the algorithm state for sorted searching
/// when computing partitions.
pub struct SortedSearch {
    /// Stores partition by columns and their ordering information
    partition_by_sort_keys: Vec<PhysicalSortExpr>,
    /// Input ordering and partition by key ordering need not be the same, so
    /// this vector stores the mapping between them. For instance, if the input
    /// is ordered by a, b and the window expression contains a PARTITION BY b, a
    /// clause, this attribute stores [1, 0].
    ordered_partition_by_indices: Vec<usize>,
    input_schema: SchemaRef,
}

impl PartitionSearcher for SortedSearch {
    /// This method constructs new output columns using the result of each window expression.
    fn calculate_out_columns(
        &mut self,
        _input_buffer: &RecordBatch,
        window_agg_states: &[PartitionWindowAggStates],
        partition_buffers: &mut PartitionBatches,
        _window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Option<Vec<ArrayRef>>> {
        let n_out = self.calculate_n_out_row(window_agg_states, partition_buffers);
        if n_out == 0 {
            Ok(None)
        } else {
            window_agg_states
                .iter()
                .map(|map| get_aggregate_result_out_column(map, n_out).map(Some))
                .collect()
        }
    }

    fn evaluate_partition_batches(
        &mut self,
        record_batch: &RecordBatch,
        _window_expr: &[Arc<dyn WindowExpr>],
    ) -> Result<Vec<(PartitionKey, RecordBatch)>> {
        let num_rows = record_batch.num_rows();
        // Calculate result of partition by column expressions
        let partition_columns = self
            .partition_by_sort_keys
            .iter()
            .map(|elem| elem.evaluate_to_sort_column(record_batch))
            .collect::<Result<Vec<_>>>()?;
        // Reorder `partition_columns` such that its ordering matches input ordering.
        let partition_columns_ordered =
            get_at_indices(&partition_columns, &self.ordered_partition_by_indices)?;
        let partition_points =
            evaluate_partition_ranges(num_rows, &partition_columns_ordered)?;
        let partition_bys = partition_columns
            .into_iter()
            .map(|arr| arr.values)
            .collect::<Vec<ArrayRef>>();

        partition_points
            .iter()
            .map(|range| {
                let row = get_row_at_idx(&partition_bys, range.start)?;
                let len = range.end - range.start;
                let slice = record_batch.slice(range.start, len);
                Ok((row, slice))
            })
            .collect::<Result<Vec<_>>>()
    }

    fn mark_partition_end(&self, partition_buffers: &mut PartitionBatches) {
        // In Sorted case. We can mark all partitions besides last partition as ended.
        // We are sure that those partitions will never receive any values.
        // (Otherwise ordering invariant is violated.)
        let n_partitions = partition_buffers.len();
        for (idx, (_, partition_batch_state)) in partition_buffers.iter_mut().enumerate()
        {
            partition_batch_state.is_end |= idx < n_partitions - 1;
        }
    }

    fn input_schema(&self) -> &SchemaRef {
        &self.input_schema
    }
}

impl SortedSearch {
    /// Calculates how many rows we can output.
    fn calculate_n_out_row(
        &mut self,
        window_agg_states: &[PartitionWindowAggStates],
        partition_buffers: &mut PartitionBatches,
    ) -> usize {
        // Different window aggregators may produce results at different rates.
        // We produce the overall batch result only as fast as the slowest one.
        let mut counts = vec![];
        let out_col_counts = window_agg_states.iter().map(|window_agg_state| {
            // Store how many elements are generated for the current
            // window expression:
            let mut cur_window_expr_out_result_len = 0;
            // We iterate over `window_agg_state`, which is an IndexMap.
            // Iterations follow the insertion order, hence we preserve
            // sorting when partition columns are sorted.
            let mut per_partition_out_results = HashMap::new();
            for (row, WindowState { state, .. }) in window_agg_state.iter() {
                cur_window_expr_out_result_len += state.out_col.len();
                let count = per_partition_out_results.entry(row).or_insert(0);
                if *count < state.out_col.len() {
                    *count = state.out_col.len();
                }
                // If we do not generate all results for the current
                // partition, we do not generate results for next
                // partition --  otherwise we will lose input ordering.
                if state.n_row_result_missing > 0 {
                    break;
                }
            }
            counts.push(per_partition_out_results);
            cur_window_expr_out_result_len
        });
        argmin(out_col_counts).map_or(0, |(min_idx, minima)| {
            let mut slowest_partition = counts.swap_remove(min_idx);
            for (partition_key, partition_batch) in partition_buffers.iter_mut() {
                if let Some(count) = slowest_partition.remove(partition_key) {
                    partition_batch.n_out_row = count;
                }
            }
            minima
        })
    }
}

/// Calculates partition by expression results for each window expression
/// on `record_batch`.
fn evaluate_partition_by_column_values(
    record_batch: &RecordBatch,
    window_expr: &[Arc<dyn WindowExpr>],
) -> Result<Vec<ArrayRef>> {
    window_expr[0]
        .partition_by()
        .iter()
        .map(|item| match item.evaluate(record_batch)? {
            ColumnarValue::Array(array) => Ok(array),
            ColumnarValue::Scalar(scalar) => {
                scalar.to_array_of_size(record_batch.num_rows())
            }
        })
        .collect()
}

/// Stream for the bounded window aggregation plan.
pub struct BoundedWindowAggStream {
    schema: SchemaRef,
    input: SendableRecordBatchStream,
    /// The record batch executor receives as input (i.e. the columns needed
    /// while calculating aggregation results).
    input_buffer: RecordBatch,
    /// Each partition's rows, accumulated across input batches. All window
    /// expressions calculate their results against these shared rows without
    /// copying.
    partition_buffers: PartitionBatches,
    /// An executor can run multiple window expressions if the PARTITION BY
    /// and ORDER BY sections are same. We keep state of the each window
    /// expression inside `window_agg_states`.
    window_agg_states: Vec<PartitionWindowAggStates>,
    finished: bool,
    window_expr: Vec<Arc<dyn WindowExpr>>,
    baseline_metrics: BaselineMetrics,
    /// Search mode for partition columns. This determines the algorithm with
    /// which we group each partition.
    search_mode: Box<dyn PartitionSearcher>,
    /// In `Linear` mode, a single-row batch containing the most recent input
    /// row (whichever partition that row belongs to); `None` in other modes
    /// and before the first non-empty batch arrives. Since in `Linear` mode
    /// the input is sorted by the first ORDER BY column, no future input row
    /// -- in any partition -- can precede this row in that column. Every
    /// partition's evaluation consults this bound to decide whether pending
    /// window frames can be finalized before the partition receives more
    /// data (which in turn allows buffered state to be pruned). Note that
    /// only the first ORDER BY column provides this guarantee. As a counter
    /// example, consider `PARTITION BY b, ORDER BY a, c` when the input is
    /// sorted by `[a, b, c]`: the mode will be `Linear`, but the last row of
    /// the input is the "last" data in terms of `[a, b, c]`, not in terms of
    /// the ordering requirement `[a, c]`. Hence, only column `a` can serve
    /// as a guarantee of the "last" data across partitions. In the `Sorted`
    /// and `PartiallySorted` modes, the leading ordering separates
    /// partitions, so finished partitions are pruned eagerly instead and no
    /// such bound is needed.
    most_recent_row: Option<RecordBatch>,
    /// Output partition index this stream serves; passed as the first
    /// argument to [`WindowStateObserver::finalize_window_aggregate`].
    partition_idx: usize,
    /// If set, invoked from [`Self::publish_finalized_states`] with the
    /// finalized per-window-expression state for every partition key that is
    /// about to be dropped.
    state_observer: Option<Arc<dyn WindowStateObserver>>,
}

impl BoundedWindowAggStream {
    /// Fire `observer` once per (window expression, partition key) for every
    /// group whose [`WindowAggState::is_end`] is true. Always mutates when
    /// called: [`datafusion_expr::Accumulator::state`] requires `&mut`, which
    /// propagates up here. The caller is responsible for deciding whether to
    /// fire (i.e. checking whether an observer is installed).
    ///
    /// Exactly-once per group is enforced by [`WindowState::aggregate_state`],
    /// which errors on second call; the `published` early-skip below avoids reaching the error.
    fn publish_finalized_states(
        &mut self,
        observer: &dyn WindowStateObserver,
    ) -> Result<()> {
        let partition_idx = self.partition_idx;
        for (expr_idx, per_expr) in self.window_agg_states.iter_mut().enumerate() {
            let window_expr = &self.window_expr[expr_idx];
            for (key, ws) in per_expr.iter_mut() {
                if ws.published || !ws.state.is_end {
                    continue;
                }
                if let Some(state) = ws.aggregate_state()? {
                    observer.finalize_window_aggregate(
                        partition_idx,
                        window_expr,
                        key,
                        state,
                    )?;
                }
            }
        }
        Ok(())
    }

    /// Prunes sections of the state that are no longer needed when calculating
    /// results (as determined by window frame boundaries and number of results generated).
    // For instance, if first `n` (not necessarily same with `n_out`) elements are no longer needed to
    // calculate window expression result (outside the window frame boundary) we retract first `n` elements
    // from the corresponding partition's batch in `self.partition_buffers`.
    // For instance, if `n_out` number of rows are calculated, we can remove
    // first `n_out` rows from `self.input_buffer`.
    fn prune_state(&mut self, n_out: usize) -> Result<()> {
        // Prune `self.window_agg_states`:
        self.prune_out_columns();
        // Prune `self.partition_buffers`:
        self.prune_partition_batches();
        // Prune `self.input_buffer`:
        self.prune_input_batch(n_out)?;
        // Prune internal state of search algorithm.
        self.search_mode.prune(n_out);
        Ok(())
    }
}

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

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

impl BoundedWindowAggStream {
    /// Create a new BoundedWindowAggStream
    fn new(
        schema: SchemaRef,
        window_expr: Vec<Arc<dyn WindowExpr>>,
        input: SendableRecordBatchStream,
        baseline_metrics: BaselineMetrics,
        search_mode: Box<dyn PartitionSearcher>,
        partition_idx: usize,
        state_observer: Option<Arc<dyn WindowStateObserver>>,
    ) -> Result<Self> {
        let state = window_expr.iter().map(|_| IndexMap::default()).collect();
        let empty_batch = RecordBatch::new_empty(Arc::clone(&schema));
        Ok(Self {
            schema,
            input,
            input_buffer: empty_batch,
            partition_buffers: IndexMap::default(),
            window_agg_states: state,
            finished: false,
            window_expr,
            baseline_metrics,
            search_mode,
            most_recent_row: None,
            partition_idx,
            state_observer,
        })
    }

    fn compute_aggregates(&mut self) -> Result<Option<RecordBatch>> {
        // calculate window cols
        let eval_ctx = WindowEvalContext::default()
            .with_most_recent_row(self.most_recent_row.as_ref());
        for (cur_window_expr, state) in
            self.window_expr.iter().zip(&mut self.window_agg_states)
        {
            cur_window_expr.evaluate_stateful(
                &self.partition_buffers,
                state,
                &eval_ctx,
            )?;
        }

        // Fire before `calculate_out_columns`: on causal frames every row
        // already streamed out, so at EOS that call returns `None` and the
        // prune path is skipped — the final partition would otherwise be
        // dropped unobserved.
        if let Some(observer) = self.state_observer.clone() {
            self.publish_finalized_states(observer.as_ref())?;
        }

        let schema = Arc::clone(&self.schema);
        let window_expr_out = self.search_mode.calculate_out_columns(
            &self.input_buffer,
            &self.window_agg_states,
            &mut self.partition_buffers,
            &self.window_expr,
        )?;
        if let Some(window_expr_out) = window_expr_out {
            let n_out = window_expr_out[0].len();
            // right append new columns to corresponding section in the original input buffer.
            let columns_to_show = self
                .input_buffer
                .columns()
                .iter()
                .map(|elem| elem.slice(0, n_out))
                .chain(window_expr_out)
                .collect::<Vec<_>>();
            let n_generated = columns_to_show[0].len();
            self.prune_state(n_generated)?;
            Ok(Some(RecordBatch::try_new(schema, columns_to_show)?))
        } else {
            Ok(None)
        }
    }

    #[inline]
    fn poll_next_inner(
        &mut self,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<RecordBatch>>> {
        if self.finished {
            return Poll::Ready(None);
        }

        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
        match ready!(self.input.poll_next_unpin(cx)) {
            Some(Ok(batch)) => {
                // Start the timer for compute time within this operator. It will be
                // stopped when dropped.
                let _timer = elapsed_compute.timer();

                if self.search_mode.is_mode_linear() && batch.num_rows() > 0 {
                    self.most_recent_row = Some(get_last_row_batch(&batch)?);
                }
                self.search_mode.update_partition_batch(
                    &mut self.input_buffer,
                    batch,
                    &self.window_expr,
                    &mut self.partition_buffers,
                )?;
                if let Some(batch) = self.compute_aggregates()? {
                    return Poll::Ready(Some(Ok(batch)));
                }
                self.poll_next_inner(cx)
            }
            Some(Err(e)) => Poll::Ready(Some(Err(e))),
            None => {
                let _timer = elapsed_compute.timer();

                self.finished = true;
                // Release the input pipeline's resources before computing the
                // final aggregates.
                let input_schema = self.input.schema();
                self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
                for (_, partition_batch_state) in self.partition_buffers.iter_mut() {
                    partition_batch_state.is_end = true;
                }
                if let Some(batch) = self.compute_aggregates()? {
                    return Poll::Ready(Some(Ok(batch)));
                }
                Poll::Ready(None)
            }
        }
    }

    /// Removes partitions that have ended. For the remaining partitions,
    /// drops buffered rows that no window expression will need again.
    fn prune_partition_batches(&mut self) {
        // Check that per-state and per-partition end-flags are consistent;
        // otherwise, the pruning code below might produce inconsistent state.
        #[cfg(debug_assertions)]
        for window_agg_state in self.window_agg_states.iter() {
            for (partition_row, WindowState { state, .. }) in window_agg_state.iter() {
                debug_assert_eq!(
                    state.is_end, self.partition_buffers[partition_row].is_end,
                    "window state's recorded end flag is out of sync with its partition"
                );
            }
        }

        // Remove partitions which we know already ended (is_end flag is true).
        // Since the retain method preserves insertion order, we still have
        // ordering in between partitions after removal.
        self.partition_buffers
            .retain(|_, partition_batch_state| !partition_batch_state.is_end);
        // Likewise, drop per-window-expression state for ended partitions.
        for window_agg_state in self.window_agg_states.iter_mut() {
            window_agg_state.retain(|_, WindowState { state, .. }| !state.is_end);
        }

        // Calculate how many rows to prune from each partition's batch. For a
        // single window expression, rows before min(window_frame_range.start,
        // last_calculated_index) are prunable: their results are already
        // calculated, and frame boundaries never move backwards, so no future
        // frame can include them. All window expressions share the partition
        // batch, so a row can only be pruned once every expression is done with
        // it: the count to prune is the minimum across expressions. A partition
        // missing from the map has nothing to prune.
        let mut n_prune_each_partition = HashMap::new();
        if let Some((first, rest)) = self.window_agg_states.split_first() {
            // First window expression seeds the prune-count map
            for (partition_row, WindowState { state, .. }) in first.iter() {
                let n_prune =
                    min(state.window_frame_range.start, state.last_calculated_index);
                if n_prune > 0 {
                    n_prune_each_partition.insert(partition_row.clone(), n_prune);
                }
            }
            // Take the per-partition min of the prune-count for each
            // additional window expression
            for window_agg_state in rest {
                n_prune_each_partition.retain(|partition_row, current| {
                    let Some(WindowState { state, .. }) =
                        window_agg_state.get(partition_row)
                    else {
                        return false;
                    };
                    let n_prune =
                        min(state.window_frame_range.start, state.last_calculated_index);
                    *current = min(*current, n_prune);
                    *current > 0
                });
            }
        }

        // Drop the prunable prefix of each partition's buffered batch:
        for (partition_row, n_prune) in n_prune_each_partition.iter() {
            debug_assert!(
                *n_prune > 0,
                "prune-count map must only contain positive entries"
            );
            let pb_state = &mut self.partition_buffers[partition_row];

            let batch = &pb_state.record_batch;
            pb_state.record_batch = batch.slice(*n_prune, batch.num_rows() - n_prune);

            // Update state indices since we have pruned some rows from the beginning:
            for window_agg_state in self.window_agg_states.iter_mut() {
                window_agg_state[partition_row].state.prune_state(*n_prune);
            }
        }
    }

    /// Prunes the section of the input batch whose aggregate results
    /// are calculated and emitted.
    fn prune_input_batch(&mut self, n_out: usize) -> Result<()> {
        // Prune first n_out rows from the input_buffer
        let n_to_keep = self.input_buffer.num_rows() - n_out;
        let batch_to_keep = self
            .input_buffer
            .columns()
            .iter()
            .map(|elem| elem.slice(n_out, n_to_keep))
            .collect::<Vec<_>>();
        self.input_buffer = RecordBatch::try_new_with_options(
            self.input_buffer.schema(),
            batch_to_keep,
            &RecordBatchOptions::new().with_row_count(Some(n_to_keep)),
        )?;
        Ok(())
    }

    /// Prunes emitted parts from WindowAggState `out_col` field.
    fn prune_out_columns(&mut self) {
        // We store generated columns for each window expression in the `out_col`
        // field of `WindowAggState`. Given how many rows are emitted, we remove
        // these sections from state.
        for partition_window_agg_states in self.window_agg_states.iter_mut() {
            // If `is_end` is set, directly remove the entry; this shrinks the
            // hash map.
            partition_window_agg_states
                .retain(|_, partition_batch_state| !partition_batch_state.state.is_end);
        }
        // Only partitions that emitted rows since the previous pruning pass
        // have output columns to shrink. Their emitted-row counts are
        // consumed and reset here, so partitions that emitted nothing keep
        // a count of zero and are passed over without any hash lookups.
        for (partition_key, partition_batch) in self.partition_buffers.iter_mut() {
            let n_emitted = partition_batch.n_out_row;
            if n_emitted == 0 {
                continue;
            }
            partition_batch.n_out_row = 0;
            for partition_window_agg_states in self.window_agg_states.iter_mut() {
                if let Some(WindowState { state, .. }) =
                    partition_window_agg_states.get_mut(partition_key)
                {
                    let out_col = &mut state.out_col;
                    let n_to_keep = out_col.len() - n_emitted;
                    *out_col = out_col.slice(n_emitted, n_to_keep);
                }
            }
        }
    }
}

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

// Gets the index of minimum entry, returns None if empty.
fn argmin<T: PartialOrd>(data: impl Iterator<Item = T>) -> Option<(usize, T)> {
    data.enumerate()
        .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(Ordering::Equal))
}

/// Calculates the section we can show results for expression
fn get_aggregate_result_out_column(
    partition_window_agg_states: &PartitionWindowAggStates,
    len_to_show: usize,
) -> Result<ArrayRef> {
    let mut result = None;
    let mut running_length = 0;
    let mut batches_to_concat = vec![];
    // We assume that iteration order is according to insertion order
    for (
        _,
        WindowState {
            state: WindowAggState { out_col, .. },
            ..
        },
    ) in partition_window_agg_states
    {
        if running_length < len_to_show {
            let n_to_use = min(len_to_show - running_length, out_col.len());
            let slice_to_use = if n_to_use == out_col.len() {
                // avoid slice when the entire column is used
                Arc::clone(out_col)
            } else {
                out_col.slice(0, n_to_use)
            };
            batches_to_concat.push(slice_to_use);
            running_length += n_to_use;
        } else {
            break;
        }
    }

    if !batches_to_concat.is_empty() {
        let array_refs: Vec<&dyn Array> =
            batches_to_concat.iter().map(|a| a.as_ref()).collect();
        result = Some(concat(&array_refs)?);
    }

    if running_length != len_to_show {
        return exec_err!(
            "Generated row number should be {len_to_show}, it is {running_length}"
        );
    }
    result.ok_or_else(|| exec_datafusion_err!("Should contain something"))
}

/// Constructs a batch from the last row of batch in the argument.
pub(crate) fn get_last_row_batch(batch: &RecordBatch) -> Result<RecordBatch> {
    if batch.num_rows() == 0 {
        return exec_err!("Latest batch should have at least 1 row");
    }
    Ok(batch.slice(batch.num_rows() - 1, 1))
}

#[cfg(test)]
mod tests {
    use std::pin::Pin;
    use std::sync::Arc;
    use std::task::{Context, Poll};
    use std::time::Duration;

    use crate::common::collect;
    use crate::execution_plan::CardinalityEffect;
    use crate::expressions::PhysicalSortExpr;
    use crate::projection::{ProjectionExec, ProjectionExpr};
    use crate::streaming::{PartitionStream, StreamingTableExec};
    use crate::test::TestMemoryExec;
    use crate::windows::bounded_window_agg_exec::WindowStateObserver;
    use crate::windows::{
        BoundedWindowAggExec, InputOrderMode, create_udwf_window_expr, create_window_expr,
    };
    use crate::{ExecutionPlan, WindowExpr, displayable, execute_stream};

    use arrow::array::{
        RecordBatch,
        builder::{Int64Builder, UInt64Builder},
    };
    use arrow::compute::SortOptions;
    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    use datafusion_common::test_util::batches_to_string;
    use datafusion_common::{Result, ScalarValue, exec_datafusion_err};
    use datafusion_execution::config::SessionConfig;
    use datafusion_execution::{
        RecordBatchStream, SendableRecordBatchStream, TaskContext,
    };
    use datafusion_expr::{
        WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
    };
    use datafusion_functions_aggregate::count::count_udaf;
    use datafusion_functions_aggregate::sum::sum_udaf;
    use datafusion_functions_window::nth_value::last_value_udwf;
    use datafusion_functions_window::nth_value::nth_value_udwf;
    use datafusion_physical_expr::expressions::{Column, Literal, col};
    use datafusion_physical_expr::window::{PartitionKey, StandardWindowExpr};
    use datafusion_physical_expr::{LexOrdering, PhysicalExpr};

    use futures::future::Shared;
    use futures::{FutureExt, Stream, StreamExt, pin_mut, ready};
    use insta::assert_snapshot;
    use itertools::Itertools;
    use tokio::time::timeout;

    #[derive(Debug, Clone)]
    struct TestStreamPartition {
        schema: SchemaRef,
        batches: Vec<RecordBatch>,
        idx: usize,
        state: PolingState,
        sleep_duration: Duration,
        send_exit: bool,
    }

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

        fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
            // We create an iterator from the record batches and map them into Ok values,
            // converting the iterator into a futures::stream::Stream
            Box::pin(self.clone())
        }
    }

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

        fn poll_next(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Option<Self::Item>> {
            self.poll_next_inner(cx)
        }
    }

    #[derive(Debug, Clone)]
    enum PolingState {
        Sleep(Shared<futures::future::BoxFuture<'static, ()>>),
        BatchReturn,
    }

    impl TestStreamPartition {
        fn poll_next_inner(
            self: &mut Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Option<Result<RecordBatch>>> {
            loop {
                match &mut self.state {
                    PolingState::BatchReturn => {
                        // Wait for self.sleep_duration before sending any new data
                        let f = tokio::time::sleep(self.sleep_duration).boxed().shared();
                        self.state = PolingState::Sleep(f);
                        let input_batch = if let Some(batch) =
                            self.batches.clone().get(self.idx)
                        {
                            batch.clone()
                        } else if self.send_exit {
                            // Send None to signal end of data
                            return Poll::Ready(None);
                        } else {
                            // Go to sleep mode
                            let f =
                                tokio::time::sleep(self.sleep_duration).boxed().shared();
                            self.state = PolingState::Sleep(f);
                            continue;
                        };
                        self.idx += 1;
                        return Poll::Ready(Some(Ok(input_batch)));
                    }
                    PolingState::Sleep(future) => {
                        pin_mut!(future);
                        ready!(future.poll_unpin(cx));
                        self.state = PolingState::BatchReturn;
                    }
                }
            }
        }
    }

    impl RecordBatchStream for TestStreamPartition {
        fn schema(&self) -> SchemaRef {
            Arc::clone(&self.schema)
        }
    }

    fn bounded_window_exec_pb_latent_range(
        input: Arc<dyn ExecutionPlan>,
        n_future_range: usize,
        hash: &str,
        order_by: &str,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let schema = input.schema();
        let window_fn = WindowFunctionDefinition::AggregateUDF(count_udaf());
        let col_expr =
            Arc::new(Column::new(schema.fields[0].name(), 0)) as Arc<dyn PhysicalExpr>;
        let args = vec![col_expr];
        let partitionby_exprs = vec![col(hash, &schema)?];
        let orderby_exprs = vec![PhysicalSortExpr {
            expr: col(order_by, &schema)?,
            options: SortOptions::default(),
        }];
        let window_frame = WindowFrame::new_bounds(
            WindowFrameUnits::Range,
            WindowFrameBound::CurrentRow,
            WindowFrameBound::Following(ScalarValue::UInt64(Some(n_future_range as u64))),
        );
        let fn_name = format!(
            "{window_fn}({args:?}) PARTITION BY: [{partitionby_exprs:?}], ORDER BY: [{orderby_exprs:?}]"
        );
        let input_order_mode = InputOrderMode::Linear;
        Ok(Arc::new(BoundedWindowAggExec::try_new(
            vec![create_window_expr(
                &window_fn,
                fn_name,
                &args,
                &partitionby_exprs,
                &orderby_exprs,
                Arc::new(window_frame),
                input.schema(),
                false,
                false,
                None,
            )?],
            input,
            input_order_mode,
            true,
        )?))
    }

    fn projection_exec(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
        let schema = input.schema();
        let exprs = input
            .schema()
            .fields
            .iter()
            .enumerate()
            .map(|(idx, field)| {
                let name = if field.name().len() > 20 {
                    format!("col_{idx}")
                } else {
                    field.name().clone()
                };
                let expr = col(field.name(), &schema).unwrap();
                (expr, name)
            })
            .collect::<Vec<_>>();
        let proj_exprs: Vec<ProjectionExpr> = exprs
            .into_iter()
            .map(|(expr, alias)| ProjectionExpr { expr, alias })
            .collect();
        Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?))
    }

    fn task_context_helper() -> TaskContext {
        let task_ctx = TaskContext::default();
        // Create session context with config
        let session_config = SessionConfig::new()
            .with_batch_size(1)
            .with_target_partitions(2)
            .with_round_robin_repartition(false);
        task_ctx.with_session_config(session_config)
    }

    fn task_context() -> Arc<TaskContext> {
        Arc::new(task_context_helper())
    }

    pub async fn collect_stream(
        mut stream: SendableRecordBatchStream,
        results: &mut Vec<RecordBatch>,
    ) -> Result<()> {
        while let Some(item) = stream.next().await {
            results.push(item?);
        }
        Ok(())
    }

    /// Execute the [ExecutionPlan] and collect the results in memory
    pub async fn collect_with_timeout(
        plan: Arc<dyn ExecutionPlan>,
        context: Arc<TaskContext>,
        timeout_duration: Duration,
    ) -> Result<Vec<RecordBatch>> {
        let stream = execute_stream(plan, context)?;
        let mut results = vec![];

        // Execute the asynchronous operation with a timeout
        if timeout(timeout_duration, collect_stream(stream, &mut results))
            .await
            .is_ok()
        {
            return Err(exec_datafusion_err!("shouldn't have completed"));
        };

        Ok(results)
    }

    fn test_schema() -> SchemaRef {
        Arc::new(Schema::new(vec![
            Field::new("sn", DataType::UInt64, true),
            Field::new("hash", DataType::Int64, true),
        ]))
    }

    fn schema_orders(schema: &SchemaRef) -> Result<Vec<LexOrdering>> {
        let orderings = vec![
            [PhysicalSortExpr {
                expr: col("sn", schema)?,
                options: SortOptions {
                    descending: false,
                    nulls_first: false,
                },
            }]
            .into(),
        ];
        Ok(orderings)
    }

    fn is_integer_division_safe(lhs: usize, rhs: usize) -> bool {
        let res = lhs / rhs;
        res * rhs == lhs
    }
    fn generate_batches(
        schema: &SchemaRef,
        n_row: usize,
        n_chunk: usize,
    ) -> Result<Vec<RecordBatch>> {
        let mut batches = vec![];
        assert!(n_row > 0);
        assert!(n_chunk > 0);
        assert!(is_integer_division_safe(n_row, n_chunk));
        let hash_replicate = 4;

        let chunks = (0..n_row)
            .chunks(n_chunk)
            .into_iter()
            .map(|elem| elem.into_iter().collect::<Vec<_>>())
            .collect::<Vec<_>>();

        // Send 2 RecordBatches at the source
        for sn_values in chunks {
            let mut sn1_array = UInt64Builder::with_capacity(sn_values.len());
            let mut hash_array = Int64Builder::with_capacity(sn_values.len());

            for sn in sn_values {
                sn1_array.append_value(sn as u64);
                let hash_value = (2 - (sn / hash_replicate)) as i64;
                hash_array.append_value(hash_value);
            }

            let batch = RecordBatch::try_new(
                Arc::clone(schema),
                vec![Arc::new(sn1_array.finish()), Arc::new(hash_array.finish())],
            )?;
            batches.push(batch);
        }
        Ok(batches)
    }

    fn generate_never_ending_source(
        n_rows: usize,
        chunk_length: usize,
        n_partition: usize,
        is_infinite: bool,
        send_exit: bool,
        per_batch_wait_duration_in_millis: u64,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        assert!(n_partition > 0);

        // We use same hash value in the table. This makes sure that
        // After hashing computation will continue in only in one of the output partitions
        // In this case, data flow should still continue
        let schema = test_schema();
        let orderings = schema_orders(&schema)?;

        // Source waits per_batch_wait_duration_in_millis ms before sending other batch
        let per_batch_wait_duration =
            Duration::from_millis(per_batch_wait_duration_in_millis);

        let batches = generate_batches(&schema, n_rows, chunk_length)?;

        // Source has 2 partitions
        let partitions = vec![
            Arc::new(TestStreamPartition {
                schema: Arc::clone(&schema),
                batches,
                idx: 0,
                state: PolingState::BatchReturn,
                sleep_duration: per_batch_wait_duration,
                send_exit,
            }) as _;
            n_partition
        ];
        let source = Arc::new(StreamingTableExec::try_new(
            Arc::clone(&schema),
            partitions,
            None,
            orderings,
            is_infinite,
            None,
        )?) as _;
        Ok(source)
    }

    // Tests NTH_VALUE(negative index) with memoize feature
    // To be able to trigger memoize feature for NTH_VALUE we need to
    // - feed BoundedWindowAggExec with batch stream data.
    // - Window frame should contain UNBOUNDED PRECEDING.
    // It hard to ensure these conditions are met, from the sql query.
    #[tokio::test]
    async fn test_window_nth_value_bounded_memoize() -> Result<()> {
        let config = SessionConfig::new().with_target_partitions(1);
        let task_ctx = Arc::new(TaskContext::default().with_session_config(config));

        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
        // Create a new batch of data to insert into the table
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3]))],
        )?;

        let memory_exec = TestMemoryExec::try_new_exec(
            &[vec![batch.clone(), batch.clone(), batch.clone()]],
            Arc::clone(&schema),
            None,
        )?;
        let col_a = col("a", &schema)?;
        let nth_value_func1 = create_udwf_window_expr(
            &nth_value_udwf(),
            &[
                Arc::clone(&col_a),
                Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
            ],
            &schema,
            "nth_value(-1)".to_string(),
            false,
        )?
        .reverse_expr()
        .unwrap();
        let nth_value_func2 = create_udwf_window_expr(
            &nth_value_udwf(),
            &[
                Arc::clone(&col_a),
                Arc::new(Literal::new(ScalarValue::Int32(Some(2)))),
            ],
            &schema,
            "nth_value(-2)".to_string(),
            false,
        )?
        .reverse_expr()
        .unwrap();

        let last_value_func = create_udwf_window_expr(
            &last_value_udwf(),
            &[Arc::clone(&col_a)],
            &schema,
            "last".to_string(),
            false,
        )?;

        let window_exprs = vec![
            // LAST_VALUE(a)
            Arc::new(StandardWindowExpr::new(
                last_value_func,
                &[],
                &[],
                Arc::new(WindowFrame::new_bounds(
                    WindowFrameUnits::Rows,
                    WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                    WindowFrameBound::CurrentRow,
                )),
            )) as _,
            // NTH_VALUE(a, -1)
            Arc::new(StandardWindowExpr::new(
                nth_value_func1,
                &[],
                &[],
                Arc::new(WindowFrame::new_bounds(
                    WindowFrameUnits::Rows,
                    WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                    WindowFrameBound::CurrentRow,
                )),
            )) as _,
            // NTH_VALUE(a, -2)
            Arc::new(StandardWindowExpr::new(
                nth_value_func2,
                &[],
                &[],
                Arc::new(WindowFrame::new_bounds(
                    WindowFrameUnits::Rows,
                    WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                    WindowFrameBound::CurrentRow,
                )),
            )) as _,
        ];
        let physical_plan = BoundedWindowAggExec::try_new(
            window_exprs,
            memory_exec,
            InputOrderMode::Sorted,
            true,
        )
        .map(|e| Arc::new(e) as Arc<dyn ExecutionPlan>)?;

        let batches = collect(physical_plan.execute(0, task_ctx)?).await?;

        // Get string representation of the plan
        assert_snapshot!(displayable(physical_plan.as_ref()).indent(true), @r#"
        BoundedWindowAggExec: wdw=[last: Field { "last": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, nth_value(-1): Field { "nth_value(-1)": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, nth_value(-2): Field { "nth_value(-2)": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
          DataSourceExec: partitions=1, partition_sizes=[3]
        "#);

        assert_snapshot!(batches_to_string(&batches), @r"
        +---+------+---------------+---------------+
        | a | last | nth_value(-1) | nth_value(-2) |
        +---+------+---------------+---------------+
        | 1 | 1    | 1             |               |
        | 2 | 2    | 2             | 1             |
        | 3 | 3    | 3             | 2             |
        | 1 | 1    | 1             | 3             |
        | 2 | 2    | 2             | 1             |
        | 3 | 3    | 3             | 2             |
        | 1 | 1    | 1             | 3             |
        | 2 | 2    | 2             | 1             |
        | 3 | 3    | 3             | 2             |
        +---+------+---------------+---------------+
        ");
        Ok(())
    }

    // In `Linear` mode, a partition may receive no new rows for several
    // input batches while other partitions keep growing. Once all of a
    // partition's buffered rows have results, the evaluation sweep skips
    // it until it receives rows again, so this test drives a partition
    // through quiet batches and then resumes it: the results after the
    // gap must continue from the retained accumulator state. Both frames
    // are causal, so results finalize in the batch their row arrives in
    // and the quiet partition is fully calculated while it waits.
    #[tokio::test]
    async fn bounded_window_linear_quiet_partition_resume() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("pk", DataType::UInt64, false),
            Field::new("ts", DataType::UInt64, false),
        ]));
        let make_batch = |rows: &[(u64, u64)]| -> Result<RecordBatch> {
            let mut pk = UInt64Builder::with_capacity(rows.len());
            let mut ts = UInt64Builder::with_capacity(rows.len());
            for (p, t) in rows {
                pk.append_value(*p);
                ts.append_value(*t);
            }
            Ok(RecordBatch::try_new(
                Arc::clone(&schema),
                vec![Arc::new(pk.finish()), Arc::new(ts.finish())],
            )?)
        };
        // `ts` ascends globally; partition 0 is absent from the middle batches.
        let batches = vec![
            make_batch(&[(0, 0), (0, 1), (1, 2)])?,
            make_batch(&[(1, 3), (1, 4)])?,
            make_batch(&[(1, 5)])?,
            make_batch(&[(0, 6), (1, 7)])?,
        ];
        let memory_exec =
            TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?;

        let partition_by = vec![col("pk", &schema)?];
        let order_by = [PhysicalSortExpr {
            expr: col("ts", &schema)?,
            options: SortOptions::default(),
        }];
        // A running COUNT (plain aggregate) and a SUM over the previous and
        // current row (sliding aggregate).
        let count_expr = create_window_expr(
            &WindowFunctionDefinition::AggregateUDF(count_udaf()),
            "count".to_string(),
            &[col("ts", &schema)?],
            &partition_by,
            &order_by,
            Arc::new(WindowFrame::new_bounds(
                WindowFrameUnits::Rows,
                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                WindowFrameBound::CurrentRow,
            )),
            Arc::clone(&schema),
            false,
            false,
            None,
        )?;
        let sum_expr = create_window_expr(
            &WindowFunctionDefinition::AggregateUDF(sum_udaf()),
            "sum".to_string(),
            &[col("ts", &schema)?],
            &partition_by,
            &order_by,
            Arc::new(WindowFrame::new_bounds(
                WindowFrameUnits::Rows,
                WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))),
                WindowFrameBound::CurrentRow,
            )),
            Arc::clone(&schema),
            false,
            false,
            None,
        )?;
        let physical_plan = BoundedWindowAggExec::try_new(
            vec![count_expr, sum_expr],
            memory_exec,
            InputOrderMode::Linear,
            true,
        )
        .map(|e| Arc::new(e) as Arc<dyn ExecutionPlan>)?;

        let batches = collect(physical_plan.execute(0, task_context())?).await?;

        assert_snapshot!(batches_to_string(&batches), @r"
        +----+----+-------+-----+
        | pk | ts | count | sum |
        +----+----+-------+-----+
        | 0  | 0  | 1     | 0   |
        | 0  | 1  | 2     | 1   |
        | 1  | 2  | 1     | 2   |
        | 1  | 3  | 2     | 5   |
        | 1  | 4  | 3     | 7   |
        | 1  | 5  | 4     | 9   |
        | 0  | 6  | 3     | 7   |
        | 1  | 7  | 5     | 12  |
        +----+----+-------+-----+
        ");
        Ok(())
    }

    // This test, tests whether most recent row guarantee by the input batch of the `BoundedWindowAggExec`
    // helps `BoundedWindowAggExec` to generate low latency result in the `Linear` mode.
    // Input data generated at the source is
    //       "+----+------+",
    //       "| sn | hash |",
    //       "+----+------+",
    //       "| 0  | 2    |",
    //       "| 1  | 2    |",
    //       "| 2  | 2    |",
    //       "| 3  | 2    |",
    //       "| 4  | 1    |",
    //       "| 5  | 1    |",
    //       "| 6  | 1    |",
    //       "| 7  | 1    |",
    //       "| 8  | 0    |",
    //       "| 9  | 0    |",
    //       "+----+------+",
    //
    // Effectively following query is run on this data
    //
    //   SELECT *, count(*) OVER(PARTITION BY duplicated_hash ORDER BY sn RANGE BETWEEN CURRENT ROW AND 1 FOLLOWING)
    //   FROM test;
    //
    // partition `duplicated_hash=2` receives following data from the input
    //
    //       "+----+------+",
    //       "| sn | hash |",
    //       "+----+------+",
    //       "| 0  | 2    |",
    //       "| 1  | 2    |",
    //       "| 2  | 2    |",
    //       "| 3  | 2    |",
    //       "+----+------+",
    // normally `BoundedWindowExec` can only generate following result from the input above
    //
    //       "+----+------+---------+",
    //       "| sn | hash |  count  |",
    //       "+----+------+---------+",
    //       "| 0  | 2    |  2      |",
    //       "| 1  | 2    |  2      |",
    //       "| 2  | 2    |<not yet>|",
    //       "| 3  | 2    |<not yet>|",
    //       "+----+------+---------+",
    // where result of last 2 row is missing. Since window frame end is not may change with future data
    // since window frame end is determined by 1 following (To generate result for row=3[where sn=2] we
    // need to received sn=4 to make sure window frame end bound won't change with future data).
    //
    // With the ability of different partitions to use global ordering at the input (where most up-to date
    //   row is
    //      "| 9  | 0    |",
    //   )
    //
    // `BoundedWindowExec` should be able to generate following result in the test
    //
    //       "+----+------+-------+",
    //       "| sn | hash | col_2 |",
    //       "+----+------+-------+",
    //       "| 0  | 2    | 2     |",
    //       "| 1  | 2    | 2     |",
    //       "| 2  | 2    | 2     |",
    //       "| 3  | 2    | 1     |",
    //       "| 4  | 1    | 2     |",
    //       "| 5  | 1    | 2     |",
    //       "| 6  | 1    | 2     |",
    //       "| 7  | 1    | 1     |",
    //       "+----+------+-------+",
    //
    // where result for all rows except last 2 is calculated (To calculate result for row 9 where sn=8
    //   we need to receive sn=10 value to calculate it result.).
    // In this test, out aim is to test for which portion of the input data `BoundedWindowExec` can generate
    // a result. To test this behaviour, we generated the data at the source infinitely (no `None` signal
    //    is sent to output from source). After, row:
    //
    //       "| 9  | 0    |",
    //
    // is sent. Source stops sending data to output. We collect, result emitted by the `BoundedWindowExec` at the
    // end of the pipeline with a timeout (Since no `None` is sent from source. Collection never ends otherwise).
    #[tokio::test]
    async fn bounded_window_exec_linear_mode_range_information() -> Result<()> {
        let n_rows = 10;
        let chunk_length = 2;
        let n_future_range = 1;

        let timeout_duration = Duration::from_millis(2000);

        let source =
            generate_never_ending_source(n_rows, chunk_length, 1, true, false, 5)?;

        let window =
            bounded_window_exec_pb_latent_range(source, n_future_range, "hash", "sn")?;

        let plan = projection_exec(window)?;

        // Get string representation of the plan
        assert_snapshot!(displayable(plan.as_ref()).indent(true), @r#"
        ProjectionExec: expr=[sn@0 as sn, hash@1 as hash, count([Column { name: "sn", index: 0 }]) PARTITION BY: [[Column { name: "hash", index: 1 }]], ORDER BY: [[PhysicalSortExpr { expr: Column { name: "sn", index: 0 }, options: SortOptions { descending: false, nulls_first: true } }]]@2 as col_2]
          BoundedWindowAggExec: wdw=[count([Column { name: "sn", index: 0 }]) PARTITION BY: [[Column { name: "hash", index: 1 }]], ORDER BY: [[PhysicalSortExpr { expr: Column { name: "sn", index: 0 }, options: SortOptions { descending: false, nulls_first: true } }]]: Field { "count([Column { name: \"sn\", index: 0 }]) PARTITION BY: [[Column { name: \"hash\", index: 1 }]], ORDER BY: [[PhysicalSortExpr { expr: Column { name: \"sn\", index: 0 }, options: SortOptions { descending: false, nulls_first: true } }]]": Int64 }, frame: RANGE BETWEEN CURRENT ROW AND 1 FOLLOWING], mode=[Linear]
            StreamingTableExec: partition_sizes=1, projection=[sn, hash], infinite_source=true, output_ordering=[sn@0 ASC NULLS LAST]
        "#);

        let task_ctx = task_context();
        let batches = collect_with_timeout(plan, task_ctx, timeout_duration).await?;

        assert_snapshot!(batches_to_string(&batches), @r"
        +----+------+-------+
        | sn | hash | col_2 |
        +----+------+-------+
        | 0  | 2    | 2     |
        | 1  | 2    | 2     |
        | 2  | 2    | 2     |
        | 3  | 2    | 1     |
        | 4  | 1    | 2     |
        | 5  | 1    | 2     |
        | 6  | 1    | 2     |
        | 7  | 1    | 1     |
        +----+------+-------+
        ");

        Ok(())
    }

    type Observation = (usize, PartitionKey, Vec<ScalarValue>);

    /// Test [`WindowStateObserver`] that records every callback into a shared
    /// `Vec` for later assertion.
    struct RecordingObserver {
        sink: Arc<std::sync::Mutex<Vec<Observation>>>,
    }

    impl WindowStateObserver for RecordingObserver {
        fn finalize_window_aggregate(
            &self,
            partition_idx: usize,
            _window_expr: &Arc<dyn WindowExpr>,
            partition_key: &PartitionKey,
            state: Vec<ScalarValue>,
        ) -> Result<()> {
            self.sink
                .lock()
                .unwrap()
                .push((partition_idx, partition_key.clone(), state));
            Ok(())
        }
    }

    /// Build a `BoundedWindowAggExec` for `count(sn) OVER (PARTITION BY hash
    /// ORDER BY sn <frame>)` over a fixed two-group source (hash=1 × 3,
    /// hash=2 × 3, sorted by (hash, sn)). Returns the plan pre-observer so
    /// callers can decide how to install it.
    fn build_partition_close_plan(frame: WindowFrame) -> Result<BoundedWindowAggExec> {
        let schema = test_schema();

        let mut sn_b = UInt64Builder::with_capacity(6);
        let mut hash_b = Int64Builder::with_capacity(6);
        for (sn, hash) in [(1u64, 1i64), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2)] {
            sn_b.append_value(sn);
            hash_b.append_value(hash);
        }
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())],
        )?;
        let ordering: LexOrdering = [
            PhysicalSortExpr {
                expr: col("hash", &schema)?,
                options: SortOptions::default(),
            },
            PhysicalSortExpr {
                expr: col("sn", &schema)?,
                options: SortOptions::default(),
            },
        ]
        .into();
        let source_raw =
            TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?
                .try_with_sort_information(vec![ordering])?;
        let source: Arc<dyn ExecutionPlan> =
            Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw)));

        let expr = create_window_expr(
            &WindowFunctionDefinition::AggregateUDF(count_udaf()),
            "cnt".to_string(),
            &[col("sn", &schema)?],
            &[col("hash", &schema)?],
            &[PhysicalSortExpr {
                expr: col("sn", &schema)?,
                options: SortOptions::default(),
            }],
            Arc::new(frame),
            source.schema(),
            false,
            false,
            None,
        )?;

        BoundedWindowAggExec::try_new(vec![expr], source, InputOrderMode::Sorted, false)
    }

    // Two PARTITION BY groups: hash=1 [sn=1,2,3] then hash=2 [sn=4,5,6].
    // Input is sorted by (hash, sn) so we can run in Sorted mode; in that
    // mode `mark_partition_end` closes the leading group mid-stream and
    // EOS closes the tail — both fire the observer for an ever-expanding
    // frame. Sliding frames are rejected at install time.

    #[tokio::test]
    async fn test_state_observer_rejects_sliding_frame() -> Result<()> {
        // `CURRENT ROW → UNBOUNDED FOLLOWING` is not ever-expanding, so this
        // maps to `SlidingAggregateWindowExpr` whose accumulator retracts as
        // rows leave the frame — at partition close the accumulator holds
        // only the last frame's rows, not the partition aggregate.
        // `with_state_observer` refuses this configuration.
        use std::sync::Mutex;

        let plan = build_partition_close_plan(WindowFrame::new_bounds(
            WindowFrameUnits::Rows,
            WindowFrameBound::CurrentRow,
            WindowFrameBound::Following(ScalarValue::UInt64(None)),
        ))?;
        let observer: Arc<dyn WindowStateObserver> = Arc::new(RecordingObserver {
            sink: Arc::new(Mutex::new(vec![])),
        });
        let err = plan.with_state_observer(Some(observer)).unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("sliding aggregate window frame"),
            "expected sliding-frame rejection, got: {msg}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_finalized_state_observer_fires_on_causal_frame() -> Result<()> {
        // `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` — ever-expanding,
        // `PlainAggregateWindowExpr` under the hood. At partition close the
        // accumulator holds the partition aggregate. Both mid-stream close
        // (hash=1 as hash=2 rows arrive) and EOS (hash=2 at drain) fire.
        use std::sync::Mutex;

        let task_ctx = Arc::new(TaskContext::default());
        let plan = build_partition_close_plan(WindowFrame::new_bounds(
            WindowFrameUnits::Rows,
            WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
            WindowFrameBound::CurrentRow,
        ))?;

        let observations: Arc<Mutex<Vec<Observation>>> = Arc::new(Mutex::new(vec![]));
        let observer: Arc<dyn WindowStateObserver> = Arc::new(RecordingObserver {
            sink: Arc::clone(&observations),
        });
        let plan = plan.with_state_observer(Some(observer))?;

        let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?;

        // count(sn) over each of hash=1 (3 rows) and hash=2 (3 rows), in
        // close order — hash=1 first (mid-stream close), hash=2 second (EOS).
        let observed: Vec<(usize, i64, Vec<ScalarValue>)> = observations
            .lock()
            .unwrap()
            .iter()
            .map(|(idx, key, state)| {
                let hash = match &key[0] {
                    ScalarValue::Int64(Some(v)) => *v,
                    other => panic!("unexpected partition-key element: {other:?}"),
                };
                (*idx, hash, state.clone())
            })
            .collect();
        assert_eq!(
            observed,
            vec![
                (0, 1, vec![ScalarValue::Int64(Some(3))]),
                (0, 2, vec![ScalarValue::Int64(Some(3))]),
            ]
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_finalized_state_observer_fires_exactly_once_across_batches()
    -> Result<()> {
        // Regression guard for the exactly-once observer contract when
        // partition close and pruning happen on different `compute_aggregates`
        // calls.
        //
        // The observer fires from `publish_finalized_states`, called at the
        // top of every `compute_aggregates`. Entries are only cleared by
        // `prune_state`, which runs only when `calculate_out_columns` returns
        // `Some`. Nothing in the type system ties the two together, so a
        // group whose state was published on batch N must not be re-published
        // on batch N+1 or at EOS.
        //
        // Layout: three PARTITION BY groups streamed across two input
        // batches, so each group closes on a distinct `compute_aggregates`
        // call:
        //   batch 1 = [hash=1 × 2]                — no close (single group).
        //   batch 2 = [hash=2 × 2, hash=3 × 2]    — `mark_partition_end`
        //                                            closes hash=1 and hash=2.
        //   EOS                                    — closes hash=3.
        //
        // Assertion: each key appears exactly once across all observations.
        use std::sync::Mutex;

        let task_ctx = Arc::new(TaskContext::default());
        let schema = test_schema();

        // Two batches, same output partition.
        let make_batch = |rows: &[(u64, i64)]| -> Result<RecordBatch> {
            let mut sn_b = UInt64Builder::with_capacity(rows.len());
            let mut hash_b = Int64Builder::with_capacity(rows.len());
            for &(sn, hash) in rows {
                sn_b.append_value(sn);
                hash_b.append_value(hash);
            }
            Ok(RecordBatch::try_new(
                Arc::clone(&schema),
                vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())],
            )?)
        };
        let batch1 = make_batch(&[(1, 1), (2, 1)])?;
        let batch2 = make_batch(&[(3, 2), (4, 2), (5, 3), (6, 3)])?;

        let ordering: LexOrdering = [
            PhysicalSortExpr {
                expr: col("hash", &schema)?,
                options: SortOptions::default(),
            },
            PhysicalSortExpr {
                expr: col("sn", &schema)?,
                options: SortOptions::default(),
            },
        ]
        .into();
        let source_raw =
            TestMemoryExec::try_new(&[vec![batch1, batch2]], Arc::clone(&schema), None)?
                .try_with_sort_information(vec![ordering])?;
        let source: Arc<dyn ExecutionPlan> =
            Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw)));

        let expr = create_window_expr(
            &WindowFunctionDefinition::AggregateUDF(count_udaf()),
            "cnt".to_string(),
            &[col("sn", &schema)?],
            &[col("hash", &schema)?],
            &[PhysicalSortExpr {
                expr: col("sn", &schema)?,
                options: SortOptions::default(),
            }],
            Arc::new(WindowFrame::new_bounds(
                WindowFrameUnits::Rows,
                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                WindowFrameBound::CurrentRow,
            )),
            source.schema(),
            false,
            false,
            None,
        )?;

        let observations: Arc<Mutex<Vec<Observation>>> = Arc::new(Mutex::new(vec![]));
        let observer: Arc<dyn WindowStateObserver> = Arc::new(RecordingObserver {
            sink: Arc::clone(&observations),
        });

        let plan = BoundedWindowAggExec::try_new(
            vec![expr],
            source,
            InputOrderMode::Sorted,
            false,
        )?
        .with_state_observer(Some(observer))?;

        let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?;

        let fired: Vec<i64> = observations
            .lock()
            .unwrap()
            .iter()
            .map(|(_, key, _)| match &key[0] {
                ScalarValue::Int64(Some(v)) => *v,
                other => panic!("unexpected partition-key element: {other:?}"),
            })
            .collect();
        // Each group closes on a distinct `compute_aggregates` call — hash=1
        // and hash=2 on batch 2's `mark_partition_end`, hash=3 at EOS — and
        // each appears exactly once, in close order.
        assert_eq!(fired, vec![1, 2, 3]);
        Ok(())
    }

    /// Run one task's local BWAG for `SUM(sn) OVER (ORDER BY sn ROWS
    /// UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, over
    /// `input` sorted ascending. Returns the per-row output values and the
    /// observed finalized state total (which the caller uses as a carry-in
    /// for the next task).
    async fn run_running_sum_task(
        input: &[u64],
        task_ctx: Arc<TaskContext>,
    ) -> Result<(Vec<u64>, u64)> {
        use arrow::array::UInt64Array;
        use datafusion_functions_aggregate::sum::sum_udaf;
        use std::sync::Mutex;

        /// Observer for `run_running_sum_task`: captures the single running
        /// SUM total published at EOS. Asserts exactly-one fire and rejects
        /// non-empty partition keys (this helper is no-PARTITION-BY only).
        struct RunningSumObserver {
            sink: Arc<Mutex<Option<u64>>>,
        }

        impl WindowStateObserver for RunningSumObserver {
            fn finalize_window_aggregate(
                &self,
                _partition_idx: usize,
                _window_expr: &Arc<dyn WindowExpr>,
                partition_key: &PartitionKey,
                state: Vec<ScalarValue>,
            ) -> Result<()> {
                assert!(
                    partition_key.is_empty(),
                    "empty PartitionKey for no-PARTITION-BY plan"
                );
                let total = match &state[0] {
                    ScalarValue::UInt64(Some(v)) => *v,
                    ScalarValue::Int64(Some(v)) => *v as u64,
                    other => panic!("unexpected sum state element: {other:?}"),
                };
                let prev = self.sink.lock().unwrap().replace(total);
                assert!(prev.is_none(), "observer must fire exactly once per task");
                Ok(())
            }
        }

        let schema = test_schema();
        let mut sn_b = UInt64Builder::with_capacity(input.len());
        let mut hash_b = Int64Builder::with_capacity(input.len());
        for &sn in input {
            sn_b.append_value(sn);
            hash_b.append_value(0);
        }
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())],
        )?;
        let ordering: LexOrdering = [PhysicalSortExpr {
            expr: col("sn", &schema)?,
            options: SortOptions::default(),
        }]
        .into();
        let source_raw =
            TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?
                .try_with_sort_information(vec![ordering])?;
        let source: Arc<dyn ExecutionPlan> =
            Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw)));

        let window_fn = WindowFunctionDefinition::AggregateUDF(sum_udaf());
        let args = vec![col("sn", &schema)?];
        let partition_by: Vec<Arc<dyn PhysicalExpr>> = vec![];
        let order_by = vec![PhysicalSortExpr {
            expr: col("sn", &schema)?,
            options: SortOptions::default(),
        }];
        let frame = WindowFrame::new_bounds(
            WindowFrameUnits::Rows,
            WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
            WindowFrameBound::CurrentRow,
        );
        let expr = create_window_expr(
            &window_fn,
            "running_sum".to_string(),
            &args,
            &partition_by,
            &order_by,
            Arc::new(frame),
            source.schema(),
            false,
            false,
            None,
        )?;

        let total_sink: Arc<Mutex<Option<u64>>> = Arc::new(Mutex::new(None));
        let observer: Arc<dyn WindowStateObserver> = Arc::new(RunningSumObserver {
            sink: Arc::clone(&total_sink),
        });

        let plan = BoundedWindowAggExec::try_new(
            vec![expr],
            source,
            InputOrderMode::Sorted,
            false,
        )?
        .with_state_observer(Some(observer))?;
        let batches = collect(Arc::new(plan).execute(0, task_ctx)?).await?;

        let mut out = Vec::with_capacity(input.len());
        for batch in &batches {
            let col = batch
                .column_by_name("running_sum")
                .expect("running_sum column present");
            let arr = col
                .as_any()
                .downcast_ref::<UInt64Array>()
                .expect("SUM(UInt64) → UInt64Array");
            for i in 0..arr.len() {
                out.push(arr.value(i));
            }
        }
        let total = total_sink
            .lock()
            .unwrap()
            .expect("observer must have fired at EOS");
        Ok((out, total))
    }

    /// Run one task's local BWAG for `approx_distinct(sn) OVER (ORDER BY sn
    /// ROWS UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, and
    /// return the single EOS-observed [`Accumulator::state`] Vec.
    async fn run_approx_distinct_task(
        input: &[u64],
        task_ctx: Arc<TaskContext>,
    ) -> Result<Vec<ScalarValue>> {
        use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf;
        use std::sync::Mutex;

        /// Observer for `run_approx_distinct_task`: capture the single EOS
        /// state. Asserts exactly-one fire and rejects non-empty partition
        /// keys (helper is no-PARTITION-BY only).
        struct ApproxDistinctObserver {
            sink: Arc<Mutex<Option<Vec<ScalarValue>>>>,
        }

        impl WindowStateObserver for ApproxDistinctObserver {
            fn finalize_window_aggregate(
                &self,
                _partition_idx: usize,
                _window_expr: &Arc<dyn WindowExpr>,
                partition_key: &PartitionKey,
                state: Vec<ScalarValue>,
            ) -> Result<()> {
                assert!(
                    partition_key.is_empty(),
                    "empty PartitionKey for no-PARTITION-BY plan"
                );
                let prev = self.sink.lock().unwrap().replace(state);
                assert!(prev.is_none(), "observer must fire exactly once per task");
                Ok(())
            }
        }

        let schema = test_schema();
        let mut sn_b = UInt64Builder::with_capacity(input.len());
        let mut hash_b = Int64Builder::with_capacity(input.len());
        for &sn in input {
            sn_b.append_value(sn);
            hash_b.append_value(0);
        }
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())],
        )?;
        let ordering: LexOrdering = [PhysicalSortExpr {
            expr: col("sn", &schema)?,
            options: SortOptions::default(),
        }]
        .into();
        let source_raw =
            TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?
                .try_with_sort_information(vec![ordering])?;
        let source: Arc<dyn ExecutionPlan> =
            Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw)));

        let expr = create_window_expr(
            &WindowFunctionDefinition::AggregateUDF(approx_distinct_udaf()),
            "approx_distinct_sn".to_string(),
            &[col("sn", &schema)?],
            &[],
            &[PhysicalSortExpr {
                expr: col("sn", &schema)?,
                options: SortOptions::default(),
            }],
            Arc::new(WindowFrame::new_bounds(
                WindowFrameUnits::Rows,
                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                WindowFrameBound::CurrentRow,
            )),
            source.schema(),
            false,
            false,
            None,
        )?;

        let state_sink: Arc<Mutex<Option<Vec<ScalarValue>>>> = Arc::new(Mutex::new(None));
        let observer: Arc<dyn WindowStateObserver> = Arc::new(ApproxDistinctObserver {
            sink: Arc::clone(&state_sink),
        });

        let plan = BoundedWindowAggExec::try_new(
            vec![expr],
            source,
            InputOrderMode::Sorted,
            false,
        )?
        .with_state_observer(Some(observer))?;
        let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?;

        state_sink
            .lock()
            .unwrap()
            .take()
            .ok_or_else(|| exec_datafusion_err!("observer never fired"))
    }

    #[tokio::test]
    async fn test_prefix_scan_across_tasks_matches_single_bwag() -> Result<()> {
        // Demonstrates the parallel-window shape reviewers asked about:
        // range-shuffle `SUM(sn) OVER (ORDER BY sn UNBOUNDED PRECEDING TO
        // CURRENT ROW)` across two tasks, then prefix-scan each task's
        // finalized state (from the observer) to carry-in the next task's
        // rows. Result must match a single BWAG over the concatenated input.
        let task_ctx = Arc::new(TaskContext::default());

        // Two tasks under range partition on sn:
        let (task1_out, task1_total) =
            run_running_sum_task(&[1, 1, 2, 2, 3, 3, 4, 4], Arc::clone(&task_ctx))
                .await?;
        let (task2_out, task2_total) =
            run_running_sum_task(&[5, 5, 6, 6, 7, 7, 8, 8], Arc::clone(&task_ctx))
                .await?;

        // Local (uncorrected) outputs and totals — first pass.
        assert_eq!(task1_out, vec![1, 2, 4, 6, 9, 12, 16, 20]);
        assert_eq!(task1_total, 20);
        assert_eq!(task2_out, vec![5, 10, 16, 22, 29, 36, 44, 52]);
        assert_eq!(task2_total, 52);

        // Prefix scan over per-task totals → carry-in for each task. Task 0's
        // carry-in is 0; task N's carry-in is the sum of tasks [0, N).
        let carry_ins = [0u64, task1_total];

        // Second pass: shift each task's local values by its carry-in.
        let task1_final: Vec<u64> = task1_out.iter().map(|v| v + carry_ins[0]).collect();
        let task2_final: Vec<u64> = task2_out.iter().map(|v| v + carry_ins[1]).collect();
        let parallel_result: Vec<u64> = task1_final
            .iter()
            .chain(task2_final.iter())
            .copied()
            .collect();

        // Oracle: single BWAG over the full concatenated input.
        let (single_result, single_total) = run_running_sum_task(
            &[1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8],
            task_ctx,
        )
        .await?;

        assert_eq!(
            parallel_result, single_result,
            "two-task prefix-scan must match single-BWAG oracle"
        );
        // And matches the sequence in the design discussion.
        assert_eq!(
            single_result,
            vec![1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, 42, 49, 56, 64, 72]
        );
        assert_eq!(single_total, 72);
        Ok(())
    }

    #[tokio::test]
    async fn test_prefix_merge_across_tasks_approx_distinct() -> Result<()> {
        // Load-bearing contract for the parallel-window use case: the state
        // exposed by `WindowStateObserver::finalize_window_aggregate` must be
        // compatible with `Accumulator::merge_batch` on a fresh accumulator
        // of the same UDAF. This is what allows non-decomposable aggregates
        // like `approx_distinct` (HLL sketch state) to be prefix-merged
        // across shard tasks — the reason we exposed accumulator state at
        // all. If this ever breaks, downstream parallel-window work has to
        // wait for a public API change.
        use arrow::array::{ArrayRef, BinaryArray};
        use arrow::datatypes::FieldRef;
        use datafusion_expr::function::AccumulatorArgs;
        use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf;

        let task_ctx = Arc::new(TaskContext::default());

        // Two tasks with overlapping inputs; concatenated distinct universe
        // is {1,2,3,4,5}.
        let state1 =
            run_approx_distinct_task(&[1, 1, 2, 3], Arc::clone(&task_ctx)).await?;
        let state2 = run_approx_distinct_task(&[3, 4, 5], Arc::clone(&task_ctx)).await?;
        let state_single =
            run_approx_distinct_task(&[1, 1, 2, 3, 3, 4, 5], Arc::clone(&task_ctx))
                .await?;

        // approx_distinct state is a single serialized-HLL Binary field.
        assert_eq!(state1.len(), 1, "single state field");
        assert_eq!(state2.len(), 1, "single state field");
        assert_eq!(state_single.len(), 1, "single state field");

        // Seed a fresh accumulator with the given serialized HLL states via
        // `merge_batch` and return its distinct-count evaluation.
        fn evaluate_merged(states: &[&ScalarValue]) -> Result<ScalarValue> {
            let udaf = approx_distinct_udaf();
            let input_schema =
                Arc::new(Schema::new(vec![Field::new("sn", DataType::UInt64, true)]));
            let return_field: FieldRef =
                Arc::new(Field::new("approx_distinct_sn", DataType::UInt64, true));
            let expr_field: FieldRef = Arc::new(Field::new("sn", DataType::UInt64, true));
            let physical_col: Arc<dyn PhysicalExpr> = col("sn", &input_schema)?;
            let args = AccumulatorArgs {
                return_field: Arc::clone(&return_field),
                schema: &input_schema,
                ignore_nulls: false,
                order_bys: &[],
                is_reversed: false,
                name: "approx_distinct",
                is_distinct: false,
                exprs: std::slice::from_ref(&physical_col),
                expr_fields: std::slice::from_ref(&expr_field),
            };
            let mut acc = udaf.accumulator(args)?;
            let byte_slices: Vec<&[u8]> = states
                .iter()
                .map(|s| match s {
                    ScalarValue::Binary(Some(v)) => v.as_slice(),
                    other => panic!("expected Binary state, got {other:?}"),
                })
                .collect();
            let bin: ArrayRef = Arc::new(BinaryArray::from_iter_values(byte_slices));
            acc.merge_batch(std::slice::from_ref(&bin))?;
            acc.evaluate()
        }

        let merged = evaluate_merged(&[&state1[0], &state2[0]])?;
        let oracle = evaluate_merged(&[&state_single[0]])?;

        assert_eq!(
            merged, oracle,
            "merged task states must match single-BWAG oracle — parallel prefix-merge contract"
        );
        // HLL is approximate but exact for a 5-element universe.
        assert_eq!(merged, ScalarValue::UInt64(Some(5)));
        Ok(())
    }

    #[test]
    fn test_bounded_window_agg_cardinality_effect() -> Result<()> {
        let schema = test_schema();
        let input: Arc<dyn ExecutionPlan> =
            Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?);
        let plan = bounded_window_exec_pb_latent_range(input, 1, "hash", "sn")?;
        let plan = plan
            .downcast_ref::<BoundedWindowAggExec>()
            .expect("expected BoundedWindowAggExec");

        assert!(matches!(
            plan.cardinality_effect(),
            CardinalityEffect::Equal
        ));
        Ok(())
    }

    /// Checks the per-partition batches that `LinearSearch` splits an input
    /// batch into: partitions appear in first-appearance order, rows within a
    /// partition keep their stream order, NULL keys form their own partition,
    /// and a single-partition batch is passed through without copying.
    #[test]
    fn test_linear_search_evaluate_partition_batches() -> Result<()> {
        use super::{LinearSearch, PartitionSearcher};
        use arrow::array::{Int32Array, Int64Array};

        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::Int32, true),
            Field::new("b", DataType::Int64, false),
        ]));
        let window_expr = create_window_expr(
            &WindowFunctionDefinition::AggregateUDF(count_udaf()),
            "count".to_string(),
            &[col("b", &schema)?],
            &[col("a", &schema)?],
            &[],
            Arc::new(WindowFrame::new(None)),
            Arc::clone(&schema),
            false,
            false,
            None,
        )?;
        let mut searcher = LinearSearch::new(vec![], Arc::clone(&schema));

        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![
                Arc::new(Int32Array::from(vec![
                    Some(1),
                    Some(2),
                    Some(1),
                    None,
                    Some(2),
                    Some(1),
                ])),
                Arc::new(Int64Array::from(vec![10, 20, 11, 30, 21, 12])),
            ],
        )?;
        let result =
            searcher.evaluate_partition_batches(&batch, &[Arc::clone(&window_expr)])?;
        assert_eq!(result.len(), 3);
        let expected = [
            (
                ScalarValue::Int32(Some(1)),
                vec![Some(1); 3],
                vec![10i64, 11, 12],
            ),
            (ScalarValue::Int32(Some(2)), vec![Some(2); 2], vec![20, 21]),
            (ScalarValue::Int32(None), vec![None], vec![30]),
        ];
        for ((key, partition_batch), (exp_key, exp_a, exp_b)) in
            result.iter().zip(expected)
        {
            assert_eq!(key, &vec![exp_key]);
            let exp_batch = RecordBatch::try_new(
                Arc::clone(&schema),
                vec![
                    Arc::new(Int32Array::from(exp_a)),
                    Arc::new(Int64Array::from(exp_b)),
                ],
            )?;
            assert_eq!(partition_batch, &exp_batch);
        }

        let single = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![
                Arc::new(Int32Array::from(vec![Some(7), Some(7)])),
                Arc::new(Int64Array::from(vec![70, 71])),
            ],
        )?;
        let result = searcher.evaluate_partition_batches(&single, &[window_expr])?;
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].0, vec![ScalarValue::Int32(Some(7))]);
        assert_eq!(result[0].1, single);
        // The whole batch belongs to one partition, so its columns are reused
        // rather than gathered into a new batch.
        assert!(Arc::ptr_eq(result[0].1.column(0), single.column(0)));
        Ok(())
    }
}