[PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove opened a new pull request, #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279

   ## Which issue does this PR close?
   
   Closes #6253.
   
   ## Rationale for this change
   
   `PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at `UNBOUNDED FOLLOWING` run in DataFusion's `WindowAggExec`. That includes Spark's default frame for `agg(x) OVER (PARTITION BY k)`. `WindowAggExec` buffers the task's entire input, concatenates it, and reserves nothing. A large or skewed task therefore grows native memory that neither Comet's pool nor Spark can see.
   
   Reserving those batches is not enough on its own. `WindowAggExec` holds the whole task input, not one window partition, so a reservation would fail any task whose input is bigger than its pool share, even when every window partition is small. A build that only added the reservation failed `SUM(v) OVER (PARTITION BY k)` over 200,000 rows in 5-row window partitions under a 4 MB pool.
   
   ## What changes are included in this PR?
   
   - The planner now uses a new `CometWindowAggExec` where it used `WindowAggExec`. It evaluates the same `WindowExpr`s one window partition at a time, as `WindowAggExec` does, with two differences:
     - **It emits window partitions as they complete.** The input is sorted by the `PARTITION BY` keys, so once a batch starts a new window partition, every earlier one is complete. Those are evaluated and emitted at once, and only the window partition that may continue is kept. Without `PARTITION BY`, it still buffers the whole input.
     - **It reserves what it buffers.** The buffered batches are counted with `RecordBatchMemoryCounter`, so buffers shared between batches are counted once. The copy made to concatenate them is reserved before it is made, as apache/datafusion#25496 does for the hash join build side. Both are released once the output is emitted.
   - A refused reservation now says that the operator cannot spill and which settings to change.
   - `BoundedWindowAggExec` is unchanged. It keeps only the rows its frames still need, so its memory is bounded by the frame rather than the window partition. It still reserves nothing.
   - The operator tuning guide gets a short section on window functions.
   
   A window partition that does not fit now fails the task with a memory error instead of growing untracked. Spark would spill in that case. Spilling for `WindowAggExec` is still apache/datafusion#22946. In on-heap mode, which the Spark SQL tests use, the pool is unbounded, so nothing is ever refused there.
   
   Measured in release mode against `WindowAggExec` with a standalone bench over 2M rows in 8192-row batches:
   
   - **Speed.** Unchanged for 10-row window partitions, and 3–18% slower elsewhere. That is at most 0.6 ms per 2M rows, spent on a second partition-boundary pass per batch.
   - **Memory.** With 20M rows (480 MB) in 1,000-row window partitions:
   
   | | Peak RSS | Peak reserved | Largest output batch |
   | --- | --- | --- | --- |
   | `WindowAggExec` | 1,254 MB | 0 | 20M rows |
   | `CometWindowAggExec` | 480 MB (the input itself) | 609 KB | 9,000 rows |
   
   Without `PARTITION BY`, both peak at 1,090 MB, but `CometWindowAggExec` reserves 960 MB of it.
   
   ## How are these changes tested?
   
   - **Rust unit tests in `window_agg.rs`.**
     - A differential test against `WindowAggExec` over 20 seeds. It uses random window partition sizes, random batch boundaries including empty batches, null keys and two key columns, with and without `PARTITION BY`.
     - Tests for emission order and for empty input.
     - Memory tests:
       - many small window partitions pass under a pool 50 times smaller than the input;
       - a window partition that does not fit is refused, with the new message;
       - the concatenated copy is reserved.
     - Each of four mutations fails at least one of these tests: no batch reservation, no copy reservation, no release after emitting, and never continuing a window partition across batches.
   - **`CometWindowExecSuite`.** The new tests:
     - Whole-partition functions over 7-row batches.
     - Many small window partitions under a 4 MB pool pass and match Spark. This test fails with a build that only adds the reservation.
     - A window partition that does not fit fails the task with `Additional allocation failed for CometWindowAggExec`.
   
     The suite passes on Spark 3.4, 3.5 and 4.1.
   - **Other suites.**
     - The 20 window SQL file tests pass.
     - TPC-DS q12, q20, q47, q53, q57, q63, q89 and q98, and their v2.7 variants, match the golden files at SF1 under all three join configurations.
     - Spark 4.1.3's `DataFrameWindowFunctionsSuite` and `DataFrameWindowFramesSuite` ran from Comet's classpath with Comet enabled on-heap. They gave the same result on main and on this branch: 74 pass, and the same 4 fail. Those 4 are the tests the Spark diff marks `IgnoreComet` or adjusts for `CometWindowExec`.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#discussion_r4347402006


##########
native/core/src/execution/operators/window_agg.rs:
##########
@@ -0,0 +1,895 @@
+// 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.
+
+//! Window functions that need their whole window partition before they produce a value.
+//!
+//! `PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at
+//! `UNBOUNDED FOLLOWING` cannot run in DataFusion's `BoundedWindowAggExec`. DataFusion's
+//! `WindowAggExec` runs them by buffering its whole input, which is the task's entire Spark
+//! partition, and concatenating it before it evaluates anything. It reserves memory for
+//! neither, and it cannot spill (apache/datafusion#22946).
+//!
+//! [`CometWindowAggExec`] evaluates the same window expressions, one window partition at a
+//! time as `WindowAggExec` does, with two differences:
+//!
+//! - The input is sorted by the `PARTITION BY` keys, so once a batch starts a new window
+//!   partition, every earlier one is complete. Those are evaluated and emitted right away, and
+//!   only the window partition that may continue is kept. Without `PARTITION BY`, the whole
+//!   input is one window partition.
+//! - The buffered batches, and the copy made when concatenating them for evaluation, are
+//!   reserved against the task's memory pool. The operator still cannot spill, so a window
+//!   partition that does not fit fails the task with a memory error instead of growing
+//!   untracked. A window partition spread over several batches needs up to twice its size,
+//!   because the copy is reserved while the batches it was made from still are.
+//!
+//! Once DataFusion's `WindowAggExec` does both, this operator can be removed.
+
+use std::fmt::Formatter;
+use std::ops::Range;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::array::{Array, ArrayRef, RecordBatch, UInt32Array};
+use arrow::compute::{concat, concat_batches, take, SortColumn};
+use arrow::datatypes::SchemaRef;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::utils::memory::RecordBatchMemoryCounter;
+use datafusion::common::utils::{evaluate_partition_ranges, transpose};
+use datafusion::common::{internal_err, DataFusionError, Result, Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::window::WindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr, PhysicalSortExpr};
+use datafusion::physical_plan::execution_plan::{
+    CardinalityEffect, ChildrenPropertiesMode, EmissionType, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::windows::WindowAggExec;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan,
+    InputDistributionRequirements, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
+};
+use futures::{ready, Stream, StreamExt};
+
+#[derive(Debug)]
+pub(crate) struct CometWindowAggExec {
+    /// DataFusion's operator for the same window expressions. It provides the output schema,
+    /// the plan properties and the ordered `PARTITION BY` keys, and is never executed.
+    window: WindowAggExec,
+    cache: Arc<PlanProperties>,
+    metrics: ExecutionPlanMetricsSet,
+}
+
+impl CometWindowAggExec {
+    pub(crate) fn try_new(
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+        input: Arc<dyn ExecutionPlan>,
+        can_repartition: bool,
+    ) -> Result<Self> {
+        Ok(Self::from_datafusion(WindowAggExec::try_new(
+            window_expr,
+            input,
+            can_repartition,
+        )?))
+    }
+
+    fn from_datafusion(window: WindowAggExec) -> Self {
+        // Unlike `WindowAggExec`, the output is emitted as window partitions complete.
+        let emission_type = if window.window_expr()[0].partition_by().is_empty() {
+            EmissionType::Final
+        } else {
+            EmissionType::Incremental
+        };
+        let cache = Arc::new(
+            window
+                .properties()
+                .as_ref()
+                .clone()
+                .with_emission_type(emission_type),
+        );
+        Self {
+            window,
+            cache,
+            metrics: ExecutionPlanMetricsSet::new(),
+        }
+    }
+}
+
+impl DisplayAs for CometWindowAggExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        // `WindowAggExec` starts these formats with its own name.
+        if matches!(t, DisplayFormatType::Default | DisplayFormatType::Verbose) {
+            write!(f, "Comet")?;
+        }
+        self.window.fmt_as(t, f)
+    }
+}
+
+impl ExecutionPlan for CometWindowAggExec {
+    fn name(&self) -> &str {
+        "CometWindowAggExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.cache
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![self.window.input()]
+    }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        self.window.apply_expressions(f)
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        self.window.maintains_input_order()
+    }
+
+    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+        self.window.required_input_ordering()
+    }
+
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        self.window.input_distribution_requirements()
+    }
+
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let replaced = Arc::new(self.window.clone()).replace_children(children, options)?;
+        let Some(window) = replaced.downcast_ref::<WindowAggExec>() else {
+            return internal_err!("WindowAggExec child replacement changed its plan type");
+        };
+        Ok(Arc::new(Self::from_datafusion(window.clone())))
+    }
+
+    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 execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let partition_by_sort_keys = self.window.partition_by_sort_keys()?;
+        // The same check as `WindowAggStream::new`.
+        if partition_by_sort_keys.len() != self.window.window_expr()[0].partition_by().len() {
+            return internal_err!("All partition by columns should have an ordering");
+        }
+        let input = self
+            .window
+            .input()
+            .execute(partition, Arc::clone(&context))?;
+        Ok(Box::pin(CometWindowAggStream {
+            schema: self.window.schema(),
+            input,
+            window_expr: self.window.window_expr().to_vec(),
+            partition_by_sort_keys,
+            buffered: vec![],
+            buffered_memory: RecordBatchMemoryCounter::new(),
+            last_keys: vec![],
+            pool: Arc::clone(context.memory_pool()),
+            consumer_name: format!("{}[{partition}]", self.name()),
+            reservation: None,
+            baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+            finished: false,
+        }))
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
+    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+        self.window.child_stats_requests(partition)
+    }
+
+    fn statistics_from_inputs(
+        &self,
+        input_stats: &[Arc<Statistics>],
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
+        self.window.statistics_from_inputs(input_stats, args)
+    }
+
+    fn cardinality_effect(&self) -> CardinalityEffect {
+        self.window.cardinality_effect()
+    }
+}
+
+struct CometWindowAggStream {
+    schema: SchemaRef,
+    input: SendableRecordBatchStream,
+    window_expr: Vec<Arc<dyn WindowExpr>>,
+    partition_by_sort_keys: Vec<PhysicalSortExpr>,
+    /// The rows of the last window partition seen so far, which may continue in the next
+    /// batch. Without `PARTITION BY`, every row seen so far. Either way, one window partition.
+    buffered: Vec<RecordBatch>,
+    /// Counts each buffer that `buffered` retains once, however many batches share it.
+    buffered_memory: RecordBatchMemoryCounter,
+    /// The partition keys of the last buffered row, one row per key, which the next batch is
+    /// compared against.
+    last_keys: Vec<ArrayRef>,
+    pool: Arc<dyn MemoryPool>,
+    consumer_name: String,
+    /// Registered when the first batch is buffered rather than when the stream is created. The
+    /// pool divides its memory among the registered consumers, and a sort below this operator
+    /// has read all of its input by the time its first batch arrives. Registering earlier would
+    /// shrink the sort's share while it reads.
+    reservation: Option<MemoryReservation>,
+    baseline_metrics: BaselineMetrics,
+    finished: bool,
+}
+
+impl CometWindowAggStream {
+    /// Reserves `bytes` more, registering the consumer with the pool on first use.
+    fn grow(&mut self, bytes: usize) -> Result<()> {
+        self.reservation().try_grow(bytes).map_err(with_oom_context)
+    }
+
+    fn reservation(&mut self) -> &mut MemoryReservation {
+        self.reservation.get_or_insert_with(|| {
+            MemoryConsumer::new(self.consumer_name.clone()).register(&self.pool)
+        })
+    }
+
+    fn evaluate_partition_keys(&self, batch: &RecordBatch) -> Result<Vec<SortColumn>> {
+        self.partition_by_sort_keys
+            .iter()
+            .map(|key| key.evaluate_to_sort_column(batch))
+            .collect()
+    }
+
+    /// Buffers `batch`, then evaluates the window partitions that it completes.
+    fn push_batch(&mut self, batch: RecordBatch) -> Result<Option<RecordBatch>> {
+        let num_rows = batch.num_rows();
+        if num_rows == 0 {
+            return Ok(None);
+        }
+        let added = self.buffered_memory.count_batch(&batch);
+        self.grow(added)?;
+        if self.partition_by_sort_keys.is_empty() {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        let keys = self.evaluate_partition_keys(&batch)?;
+        let mut ranges = evaluate_partition_ranges(num_rows, &keys)?;
+        // The batch's last window partition may continue in the next batch.
+        let open_start = ranges.pop().map_or(0, |open| open.start);
+        let continues = self.continues_buffered_partition(&keys)?;
+        // Whatever happens to the rest, this batch's last row is the last one buffered.
+        let last_row = UInt32Array::from(vec![(num_rows - 1) as u32]);
+        self.last_keys = keys
+            .iter()
+            .map(|key| Ok(take(key.values.as_ref(), &last_row, None)?))
+            .collect::<Result<_>>()?;
+        if ranges.is_empty() && continues {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        // The window partitions that have ended are the buffered one, followed by those of the
+        // batch's rows before `open_start`. If the batch continues the buffered window
+        // partition, its first range extends it.
+        let offset = self
+            .buffered
+            .iter()
+            .map(RecordBatch::num_rows)
+            .sum::<usize>();
+        let mut ranges: Vec<_> = ranges
+            .into_iter()
+            .map(|range| range.start + offset..range.end + offset)
+            .collect();
+        if continues {
+            ranges[0].start = 0;
+        } else if offset > 0 {
+            ranges.insert(0, 0..offset);
+        }
+        let mut complete = std::mem::take(&mut self.buffered);
+        if open_start > 0 {
+            complete.push(batch.slice(0, open_start));
+        }
+        self.buffered
+            .push(batch.slice(open_start, num_rows - open_start));
+        self.evaluate(complete, ranges)
+    }
+
+    /// Whether the first row of the batch whose partition keys are `keys` belongs to the window
+    /// partition of the last buffered row.
+    fn continues_buffered_partition(&self, keys: &[SortColumn]) -> Result<bool> {
+        if self.buffered.is_empty() {
+            return Ok(false);
+        }
+        let pairs = self
+            .last_keys
+            .iter()
+            .zip(keys)
+            .map(|(last, key)| {
+                Ok(SortColumn {
+                    values: concat(&[last.as_ref(), key.values.slice(0, 1).as_ref()])?,
+                    options: key.options,
+                })
+            })
+            .collect::<Result<Vec<_>>>()?;
+        Ok(evaluate_partition_ranges(2, &pairs)?.len() == 1)
+    }
+
+    /// Evaluates the window expressions over `batches`, whose rows make up the window
+    /// partitions `ranges`, and returns the rows with the window columns appended.
+    fn evaluate(
+        &mut self,
+        batches: Vec<RecordBatch>,
+        ranges: impl IntoIterator<Item = Range<usize>>,
+    ) -> Result<Option<RecordBatch>> {
+        if batches.is_empty() {
+            return Ok(None);
+        }
+
+        let batch = concat_batches(&self.input.schema(), &batches)?;
+        // `buffered_memory` has counted every buffer of `batches`, so this counts only the
+        // buffers of the copy that it does not share with them, and none for a single batch,
+        // which is not copied. That is exact, where estimating the copy from each batch would
+        // charge every batch sliced from one list array for all of its values.
+        let copy = self.buffered_memory.count_batch(&batch);
+        self.grow(copy)?;
+        drop(batches);
+
+        let mut partition_results = vec![];
+        for range in ranges {
+            let partition = batch.slice(range.start, range.len());
+            partition_results.push(
+                self.window_expr
+                    .iter()
+                    .map(|expr| expr.evaluate(&partition))
+                    .collect::<Result<Vec<_>>>()?,
+            );
+        }
+        let mut columns = batch.columns().to_vec();
+        for results in transpose(partition_results) {
+            columns.push(concat(
+                &results
+                    .iter()
+                    .map(|array| array.as_ref())
+                    .collect::<Vec<_>>(),
+            )?);
+        }
+        let output = RecordBatch::try_new(Arc::clone(&self.schema), columns)?;
+
+        // The output is the consumer's to reserve, so only the rows still buffered stay
+        // reserved.
+        self.buffered_memory = RecordBatchMemoryCounter::new();
+        for batch in &self.buffered {
+            self.buffered_memory.count_batch(batch);
+        }
+        let buffered = self.buffered_memory.memory_usage();
+        self.reservation().resize(buffered);
+        Ok(Some(output))
+    }
+
+    fn poll_next_inner(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+        // Times the work on each batch, including the batches that only extend a window
+        // partition, but not the wait for the input.
+        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+        loop {
+            if self.finished {
+                return Poll::Ready(None);
+            }
+            let output = match ready!(self.input.poll_next_unpin(cx)) {
+                Some(Ok(batch)) => {
+                    let _timer = elapsed_compute.timer();
+                    self.push_batch(batch)
+                }
+                Some(Err(e)) => Err(e),
+                None => {
+                    self.finished = true;
+                    // Release the input pipeline's resources before evaluating the last window
+                    // partition.
+                    self.input = Box::pin(EmptyRecordBatchStream::new(self.input.schema()));
+                    let _timer = elapsed_compute.timer();
+                    let buffered = std::mem::take(&mut self.buffered);
+                    let num_rows = buffered.iter().map(RecordBatch::num_rows).sum();
+                    self.evaluate(buffered, std::iter::once(0..num_rows))

Review Comment:
   Done in 4158fcf0d1. The reservation is dropped right after the last window partition is evaluated, and the registration test now checks that the consumer is gone once the stream drains.



##########
native/core/src/execution/planner/window.rs:
##########
@@ -140,7 +142,7 @@ impl OperatorBuilder for WindowBuilder {
                 !partition_exprs.is_empty(),
             )?)
         } else {
-            Arc::new(WindowAggExec::try_new(
+            Arc::new(CometWindowAggExec::try_new(

Review Comment:
   Switched the example to `CometWindowAggExec` in 4158fcf0d1.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#discussion_r4303087935


##########
docs/source/user-guide/latest/tuning/operators.md:
##########
@@ -59,6 +59,17 @@ boundaries. A shuffled hash join can still filter probe batches after shuffle, b
 its filter back to an earlier scan stage. Compare the [runtime-filter and scan metrics](../metrics.md#hash-joins)
 with the setting disabled to distinguish reduced hash-probe work from reader I/O savings.
 
+## Window Functions
+
+`PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at `UNBOUNDED FOLLOWING`, which includes an
+aggregate over `PARTITION BY` with no `ORDER BY` such as `sum(x) OVER (PARTITION BY k)`, need a whole window partition
+before they return anything. Comet buffers one window partition at a time for them and reserves it from its native
+memory pool, but it does not yet provide spill-to-disk for them. A window partition that does not fit in its share of
+the pool, for example because of a heavily skewed key or a window without `PARTITION BY`, fails the task with a memory

Review Comment:
   Good call. I added a "Window Functions" entry to the 1.2.0 section of the upgrade guide, in the style of the 1.1.0 "Memory Pool Limits" entry. I didn't add a `spark.comet.legacy.*` key. The old behavior was memory that nobody tracked, so the only way back is to turn native windows off, and the entry says that.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["rich7420 (via GitHub)" <gi...@apache.org>]

rich7420 commented on code in PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#discussion_r4347102484


##########
native/core/src/execution/planner/window.rs:
##########
@@ -140,7 +142,7 @@ impl OperatorBuilder for WindowBuilder {
                 !partition_exprs.is_empty(),
             )?)
         } else {
-            Arc::new(WindowAggExec::try_new(
+            Arc::new(CometWindowAggExec::try_new(

Review Comment:
   nit: the window example in `docs/source/contributor-guide/debugging.md` still builds `WindowAggExec`. Could we switch it to `CometWindowAggExec` too?



##########
native/core/src/execution/operators/window_agg.rs:
##########
@@ -0,0 +1,895 @@
+// 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.
+
+//! Window functions that need their whole window partition before they produce a value.
+//!
+//! `PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at
+//! `UNBOUNDED FOLLOWING` cannot run in DataFusion's `BoundedWindowAggExec`. DataFusion's
+//! `WindowAggExec` runs them by buffering its whole input, which is the task's entire Spark
+//! partition, and concatenating it before it evaluates anything. It reserves memory for
+//! neither, and it cannot spill (apache/datafusion#22946).
+//!
+//! [`CometWindowAggExec`] evaluates the same window expressions, one window partition at a
+//! time as `WindowAggExec` does, with two differences:
+//!
+//! - The input is sorted by the `PARTITION BY` keys, so once a batch starts a new window
+//!   partition, every earlier one is complete. Those are evaluated and emitted right away, and
+//!   only the window partition that may continue is kept. Without `PARTITION BY`, the whole
+//!   input is one window partition.
+//! - The buffered batches, and the copy made when concatenating them for evaluation, are
+//!   reserved against the task's memory pool. The operator still cannot spill, so a window
+//!   partition that does not fit fails the task with a memory error instead of growing
+//!   untracked. A window partition spread over several batches needs up to twice its size,
+//!   because the copy is reserved while the batches it was made from still are.
+//!
+//! Once DataFusion's `WindowAggExec` does both, this operator can be removed.
+
+use std::fmt::Formatter;
+use std::ops::Range;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::array::{Array, ArrayRef, RecordBatch, UInt32Array};
+use arrow::compute::{concat, concat_batches, take, SortColumn};
+use arrow::datatypes::SchemaRef;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::utils::memory::RecordBatchMemoryCounter;
+use datafusion::common::utils::{evaluate_partition_ranges, transpose};
+use datafusion::common::{internal_err, DataFusionError, Result, Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::window::WindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr, PhysicalSortExpr};
+use datafusion::physical_plan::execution_plan::{
+    CardinalityEffect, ChildrenPropertiesMode, EmissionType, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::windows::WindowAggExec;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan,
+    InputDistributionRequirements, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
+};
+use futures::{ready, Stream, StreamExt};
+
+#[derive(Debug)]
+pub(crate) struct CometWindowAggExec {
+    /// DataFusion's operator for the same window expressions. It provides the output schema,
+    /// the plan properties and the ordered `PARTITION BY` keys, and is never executed.
+    window: WindowAggExec,
+    cache: Arc<PlanProperties>,
+    metrics: ExecutionPlanMetricsSet,
+}
+
+impl CometWindowAggExec {
+    pub(crate) fn try_new(
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+        input: Arc<dyn ExecutionPlan>,
+        can_repartition: bool,
+    ) -> Result<Self> {
+        Ok(Self::from_datafusion(WindowAggExec::try_new(
+            window_expr,
+            input,
+            can_repartition,
+        )?))
+    }
+
+    fn from_datafusion(window: WindowAggExec) -> Self {
+        // Unlike `WindowAggExec`, the output is emitted as window partitions complete.
+        let emission_type = if window.window_expr()[0].partition_by().is_empty() {
+            EmissionType::Final
+        } else {
+            EmissionType::Incremental
+        };
+        let cache = Arc::new(
+            window
+                .properties()
+                .as_ref()
+                .clone()
+                .with_emission_type(emission_type),
+        );
+        Self {
+            window,
+            cache,
+            metrics: ExecutionPlanMetricsSet::new(),
+        }
+    }
+}
+
+impl DisplayAs for CometWindowAggExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        // `WindowAggExec` starts these formats with its own name.
+        if matches!(t, DisplayFormatType::Default | DisplayFormatType::Verbose) {
+            write!(f, "Comet")?;
+        }
+        self.window.fmt_as(t, f)
+    }
+}
+
+impl ExecutionPlan for CometWindowAggExec {
+    fn name(&self) -> &str {
+        "CometWindowAggExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.cache
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![self.window.input()]
+    }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        self.window.apply_expressions(f)
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        self.window.maintains_input_order()
+    }
+
+    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+        self.window.required_input_ordering()
+    }
+
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        self.window.input_distribution_requirements()
+    }
+
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let replaced = Arc::new(self.window.clone()).replace_children(children, options)?;
+        let Some(window) = replaced.downcast_ref::<WindowAggExec>() else {
+            return internal_err!("WindowAggExec child replacement changed its plan type");
+        };
+        Ok(Arc::new(Self::from_datafusion(window.clone())))
+    }
+
+    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 execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let partition_by_sort_keys = self.window.partition_by_sort_keys()?;
+        // The same check as `WindowAggStream::new`.
+        if partition_by_sort_keys.len() != self.window.window_expr()[0].partition_by().len() {
+            return internal_err!("All partition by columns should have an ordering");
+        }
+        let input = self
+            .window
+            .input()
+            .execute(partition, Arc::clone(&context))?;
+        Ok(Box::pin(CometWindowAggStream {
+            schema: self.window.schema(),
+            input,
+            window_expr: self.window.window_expr().to_vec(),
+            partition_by_sort_keys,
+            buffered: vec![],
+            buffered_memory: RecordBatchMemoryCounter::new(),
+            last_keys: vec![],
+            pool: Arc::clone(context.memory_pool()),
+            consumer_name: format!("{}[{partition}]", self.name()),
+            reservation: None,
+            baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+            finished: false,
+        }))
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
+    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+        self.window.child_stats_requests(partition)
+    }
+
+    fn statistics_from_inputs(
+        &self,
+        input_stats: &[Arc<Statistics>],
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
+        self.window.statistics_from_inputs(input_stats, args)
+    }
+
+    fn cardinality_effect(&self) -> CardinalityEffect {
+        self.window.cardinality_effect()
+    }
+}
+
+struct CometWindowAggStream {
+    schema: SchemaRef,
+    input: SendableRecordBatchStream,
+    window_expr: Vec<Arc<dyn WindowExpr>>,
+    partition_by_sort_keys: Vec<PhysicalSortExpr>,
+    /// The rows of the last window partition seen so far, which may continue in the next
+    /// batch. Without `PARTITION BY`, every row seen so far. Either way, one window partition.
+    buffered: Vec<RecordBatch>,
+    /// Counts each buffer that `buffered` retains once, however many batches share it.
+    buffered_memory: RecordBatchMemoryCounter,
+    /// The partition keys of the last buffered row, one row per key, which the next batch is
+    /// compared against.
+    last_keys: Vec<ArrayRef>,
+    pool: Arc<dyn MemoryPool>,
+    consumer_name: String,
+    /// Registered when the first batch is buffered rather than when the stream is created. The
+    /// pool divides its memory among the registered consumers, and a sort below this operator
+    /// has read all of its input by the time its first batch arrives. Registering earlier would
+    /// shrink the sort's share while it reads.
+    reservation: Option<MemoryReservation>,
+    baseline_metrics: BaselineMetrics,
+    finished: bool,
+}
+
+impl CometWindowAggStream {
+    /// Reserves `bytes` more, registering the consumer with the pool on first use.
+    fn grow(&mut self, bytes: usize) -> Result<()> {
+        self.reservation().try_grow(bytes).map_err(with_oom_context)
+    }
+
+    fn reservation(&mut self) -> &mut MemoryReservation {
+        self.reservation.get_or_insert_with(|| {
+            MemoryConsumer::new(self.consumer_name.clone()).register(&self.pool)
+        })
+    }
+
+    fn evaluate_partition_keys(&self, batch: &RecordBatch) -> Result<Vec<SortColumn>> {
+        self.partition_by_sort_keys
+            .iter()
+            .map(|key| key.evaluate_to_sort_column(batch))
+            .collect()
+    }
+
+    /// Buffers `batch`, then evaluates the window partitions that it completes.
+    fn push_batch(&mut self, batch: RecordBatch) -> Result<Option<RecordBatch>> {
+        let num_rows = batch.num_rows();
+        if num_rows == 0 {
+            return Ok(None);
+        }
+        let added = self.buffered_memory.count_batch(&batch);
+        self.grow(added)?;
+        if self.partition_by_sort_keys.is_empty() {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        let keys = self.evaluate_partition_keys(&batch)?;
+        let mut ranges = evaluate_partition_ranges(num_rows, &keys)?;
+        // The batch's last window partition may continue in the next batch.
+        let open_start = ranges.pop().map_or(0, |open| open.start);
+        let continues = self.continues_buffered_partition(&keys)?;
+        // Whatever happens to the rest, this batch's last row is the last one buffered.
+        let last_row = UInt32Array::from(vec![(num_rows - 1) as u32]);
+        self.last_keys = keys
+            .iter()
+            .map(|key| Ok(take(key.values.as_ref(), &last_row, None)?))
+            .collect::<Result<_>>()?;
+        if ranges.is_empty() && continues {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        // The window partitions that have ended are the buffered one, followed by those of the
+        // batch's rows before `open_start`. If the batch continues the buffered window
+        // partition, its first range extends it.
+        let offset = self
+            .buffered
+            .iter()
+            .map(RecordBatch::num_rows)
+            .sum::<usize>();
+        let mut ranges: Vec<_> = ranges
+            .into_iter()
+            .map(|range| range.start + offset..range.end + offset)
+            .collect();
+        if continues {
+            ranges[0].start = 0;
+        } else if offset > 0 {
+            ranges.insert(0, 0..offset);
+        }
+        let mut complete = std::mem::take(&mut self.buffered);
+        if open_start > 0 {
+            complete.push(batch.slice(0, open_start));
+        }
+        self.buffered
+            .push(batch.slice(open_start, num_rows - open_start));
+        self.evaluate(complete, ranges)
+    }
+
+    /// Whether the first row of the batch whose partition keys are `keys` belongs to the window
+    /// partition of the last buffered row.
+    fn continues_buffered_partition(&self, keys: &[SortColumn]) -> Result<bool> {
+        if self.buffered.is_empty() {
+            return Ok(false);
+        }
+        let pairs = self
+            .last_keys
+            .iter()
+            .zip(keys)
+            .map(|(last, key)| {
+                Ok(SortColumn {
+                    values: concat(&[last.as_ref(), key.values.slice(0, 1).as_ref()])?,
+                    options: key.options,
+                })
+            })
+            .collect::<Result<Vec<_>>>()?;
+        Ok(evaluate_partition_ranges(2, &pairs)?.len() == 1)
+    }
+
+    /// Evaluates the window expressions over `batches`, whose rows make up the window
+    /// partitions `ranges`, and returns the rows with the window columns appended.
+    fn evaluate(
+        &mut self,
+        batches: Vec<RecordBatch>,
+        ranges: impl IntoIterator<Item = Range<usize>>,
+    ) -> Result<Option<RecordBatch>> {
+        if batches.is_empty() {
+            return Ok(None);
+        }
+
+        let batch = concat_batches(&self.input.schema(), &batches)?;
+        // `buffered_memory` has counted every buffer of `batches`, so this counts only the
+        // buffers of the copy that it does not share with them, and none for a single batch,
+        // which is not copied. That is exact, where estimating the copy from each batch would
+        // charge every batch sliced from one list array for all of its values.
+        let copy = self.buffered_memory.count_batch(&batch);
+        self.grow(copy)?;
+        drop(batches);
+
+        let mut partition_results = vec![];
+        for range in ranges {
+            let partition = batch.slice(range.start, range.len());
+            partition_results.push(
+                self.window_expr
+                    .iter()
+                    .map(|expr| expr.evaluate(&partition))
+                    .collect::<Result<Vec<_>>>()?,
+            );
+        }
+        let mut columns = batch.columns().to_vec();
+        for results in transpose(partition_results) {
+            columns.push(concat(
+                &results
+                    .iter()
+                    .map(|array| array.as_ref())
+                    .collect::<Vec<_>>(),
+            )?);
+        }
+        let output = RecordBatch::try_new(Arc::clone(&self.schema), columns)?;
+
+        // The output is the consumer's to reserve, so only the rows still buffered stay
+        // reserved.
+        self.buffered_memory = RecordBatchMemoryCounter::new();
+        for batch in &self.buffered {
+            self.buffered_memory.count_batch(batch);
+        }
+        let buffered = self.buffered_memory.memory_usage();
+        self.reservation().resize(buffered);
+        Ok(Some(output))
+    }
+
+    fn poll_next_inner(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+        // Times the work on each batch, including the batches that only extend a window
+        // partition, but not the wait for the input.
+        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+        loop {
+            if self.finished {
+                return Poll::Ready(None);
+            }
+            let output = match ready!(self.input.poll_next_unpin(cx)) {
+                Some(Ok(batch)) => {
+                    let _timer = elapsed_compute.timer();
+                    self.push_batch(batch)
+                }
+                Some(Err(e)) => Err(e),
+                None => {
+                    self.finished = true;
+                    // Release the input pipeline's resources before evaluating the last window
+                    // partition.
+                    self.input = Box::pin(EmptyRecordBatchStream::new(self.input.schema()));
+                    let _timer = elapsed_compute.timer();
+                    let buffered = std::mem::take(&mut self.buffered);
+                    let num_rows = buffered.iter().map(RecordBatch::num_rows).sum();
+                    self.evaluate(buffered, std::iter::once(0..num_rows))

Review Comment:
   Could we drop the reservation after evaluating the last partition? Otherwise, the finished stream stays registered, leaving less memory for the downstream sort while it inserts this batch.
   
   ```suggestion
                       let output = self.evaluate(buffered, std::iter::once(0..num_rows));
                       self.reservation = None;
                       output
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-6235269446

   Thanks, these were all useful. Pushed in 766e0dea10, which also merges main to clear the conflicts.
   
   The consumer now registers when the first batch is buffered, so it no longer takes a share from the sort while the sort is still reading its input. A new test checks that it registers then and not for empty input. The error hint and the tuning doc now say that a window partition needs up to twice its size within the operator's share, and both suggest `greedy_unified` before `spark.comet.exec.window.enabled=false`, with a note that the latter also moves the bounded window functions back to Spark. I left the pool itself alone (#5465).
   
   `evaluate` now counts the copy with the counter that already holds the buffered batches, so the second counter and the `batches.len() > 1` branch are gone. `push_batch` and the final evaluation are timed in `poll_next_inner` and the inner timer is removed, so batches that only extend the open partition record time too. The last row's keys are kept from the previous batch, as a copy so they don't pin the batch's buffers, instead of being re-evaluated from the last buffered batch every time.
   
   For the tests, the first Scala test is now `windows/window_whole_partition_batches.sql` without the `OVER ()` query. The failure test asserts the `CometWindowAggExec cannot spill` hint, and it passes, so the hint does reach Spark. In the Rust tests I dropped the interleaved empty-batch block and the `peak` assertions.
   
   This is now `Part of #6253`, and the description says why `BoundedWindowAggExec` is still open. The follow-up for `BoundedWindowAggExec` is #7092. The window is added to the pool list in `tuning/memory.md`, to the imported-batches paragraph in `memory_management.md` (an `OVER ()` window has no sort below it), and to the cannot-spill item in its checklist.
   
   Not yet on offering the stream upstream. I'll leave the removal condition as it is until there is upstream work to point at.
   
   `run-spark-4.1-tests` is on the PR. I haven't run the TPC-DS comparison under `fair_unified` yet.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-6264145733

   Thanks again for the reviews. I've pushed `35eac17a70`, which merges `main` and is ready for another look.
   
   Your earlier approvals were on `bde4aa0f76`, and the code has changed since then in three commits:
   
   - `2a81483e91` and `63f51fa22c` address @comphead's review.
     - The consumer now registers when the first batch is buffered, so it no longer takes a share from a sort that is still reading its input.
     - `evaluate` counts the concatenated copy with the counter that already holds the buffered batches. The second counter and the `batches.len() > 1` branch are gone.
     - `push_batch` and the final evaluation are timed in `poll_next_inner`, and the inner timer is removed.
     - The last row's keys are kept from the previous batch instead of being re-evaluated each time.
     - The failure test asserts the `CometWindowAggExec cannot spill` hint.
     - The first Scala test is now `windows/window_whole_partition_batches.sql`.
     - The docs and the upgrade guide are updated.
   - `766e0dea10` is a small simplification that shares the grow-and-wrap helper and the test pool lookup.
   - `35eac17a70` is the latest merge of `main`. #7014 moved window planning into `planner/window.rs`, so the swap to `CometWindowAggExec` now lives there. `planner.rs` is back to `main`'s version. That swap is the only change to existing code in the PR.
   
   Locally, the 7 `window_agg` tests pass, and `cargo fmt` and `cargo clippy -D warnings` are clean. The Spark 4.1 jobs passed on `766e0dea10` with `run-spark-4.1-tests`, and I'll check them again on the new head.
   
   One thing is still open. @comphead asked for a TPC-DS window query under `fair_unified`, with a pool small enough to make the sort spill, compared against `main`. I haven't run that yet. The standalone bench doesn't cover it, so I can't say yet whether the share dilution or the per-batch pool calls add cost there. Do you want that run before this merges?
   
   `BoundedWindowAggExec` is still untracked. That is #7092, and this PR is `Part of #6253`.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-6345773265

   btw I filed bunch of PRs/issues for window execution in DF 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["sunchao (via GitHub)" <gi...@apache.org>]

sunchao commented on code in PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#discussion_r4116302589


##########
native/core/src/execution/operators/window_agg.rs:
##########
@@ -0,0 +1,800 @@
+// 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.
+
+//! Window functions that need their whole window partition before they produce a value.
+//!
+//! `PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at
+//! `UNBOUNDED FOLLOWING` cannot run in DataFusion's `BoundedWindowAggExec`. DataFusion's
+//! `WindowAggExec` runs them by buffering its whole input, which is the task's entire Spark
+//! partition, and concatenating it before it evaluates anything. It reserves memory for
+//! neither, and it cannot spill (apache/datafusion#22946).
+//!
+//! [`CometWindowAggExec`] evaluates the same window expressions, one window partition at a
+//! time as `WindowAggExec` does, with two differences:
+//!
+//! - The input is sorted by the `PARTITION BY` keys, so once a batch starts a new window
+//!   partition, every earlier one is complete. Those are evaluated and emitted right away, and
+//!   only the window partition that may continue is kept. Without `PARTITION BY`, the whole
+//!   input is one window partition.
+//! - The buffered batches, and the copy made when concatenating them for evaluation, are
+//!   reserved against the task's memory pool. The operator still cannot spill, so a window
+//!   partition that does not fit fails the task with a memory error instead of growing
+//!   untracked.
+
+use std::fmt::Formatter;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::array::{Array, RecordBatch};
+use arrow::compute::{concat, concat_batches, SortColumn};
+use arrow::datatypes::SchemaRef;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::utils::memory::RecordBatchMemoryCounter;
+use datafusion::common::utils::{evaluate_partition_ranges, transpose};
+use datafusion::common::{internal_err, DataFusionError, Result, Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::window::WindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr, PhysicalSortExpr};
+use datafusion::physical_plan::execution_plan::{
+    CardinalityEffect, ChildrenPropertiesMode, EmissionType, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::windows::{get_ordered_partition_by_indices, WindowAggExec};
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan,
+    InputDistributionRequirements, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
+};
+use futures::{ready, Stream, StreamExt};
+
+#[derive(Debug)]
+pub(crate) struct CometWindowAggExec {
+    /// DataFusion's operator for the same window expressions. It provides the output schema,
+    /// the plan properties and the ordered `PARTITION BY` keys, and is never executed.
+    window: WindowAggExec,
+    /// How the `PARTITION BY` expressions map onto the input ordering, as `WindowAggExec`
+    /// computes it.
+    ordered_partition_by_indices: Vec<usize>,
+    cache: Arc<PlanProperties>,
+    metrics: ExecutionPlanMetricsSet,
+}
+
+impl CometWindowAggExec {
+    pub(crate) fn try_new(
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+        input: Arc<dyn ExecutionPlan>,
+        can_repartition: bool,
+    ) -> Result<Self> {
+        Self::from_datafusion(WindowAggExec::try_new(window_expr, input, can_repartition)?)
+    }
+
+    fn from_datafusion(window: WindowAggExec) -> Result<Self> {
+        let partition_by = window.window_expr()[0].partition_by();
+        let ordered_partition_by_indices =
+            get_ordered_partition_by_indices(partition_by, window.input())?;
+        // Unlike `WindowAggExec`, the output is emitted as window partitions complete.
+        let emission_type = if partition_by.is_empty() {
+            EmissionType::Final
+        } else {
+            EmissionType::Incremental
+        };
+        let cache = Arc::new(
+            window
+                .properties()
+                .as_ref()
+                .clone()
+                .with_emission_type(emission_type),
+        );
+        Ok(Self {
+            window,
+            ordered_partition_by_indices,
+            cache,
+            metrics: ExecutionPlanMetricsSet::new(),
+        })
+    }
+}
+
+impl DisplayAs for CometWindowAggExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        // `WindowAggExec` starts these formats with its own name.
+        if matches!(t, DisplayFormatType::Default | DisplayFormatType::Verbose) {
+            write!(f, "Comet")?;
+        }
+        self.window.fmt_as(t, f)
+    }
+}
+
+impl ExecutionPlan for CometWindowAggExec {
+    fn name(&self) -> &str {
+        "CometWindowAggExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.cache
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![self.window.input()]
+    }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        self.window.apply_expressions(f)
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        self.window.maintains_input_order()
+    }
+
+    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+        self.window.required_input_ordering()
+    }
+
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        self.window.input_distribution_requirements()
+    }
+
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return internal_err!("CometWindowAggExec requires one child");
+        }
+        let replaced = Arc::new(self.window.clone()).replace_children(children, options)?;
+        let Some(window) = replaced.downcast_ref::<WindowAggExec>() else {
+            return internal_err!("WindowAggExec child replacement changed its plan type");
+        };
+        Ok(Arc::new(Self::from_datafusion(window.clone())?))
+    }
+
+    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 execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let partition_by_sort_keys = self.window.partition_by_sort_keys()?;
+        let input = self
+            .window
+            .input()
+            .execute(partition, Arc::clone(&context))?;
+        let reservation = MemoryConsumer::new(format!("{}[{partition}]", self.name()))
+            .register(context.memory_pool());
+        Ok(Box::pin(CometWindowAggStream::try_new(
+            self.window.schema(),
+            self.window.window_expr().to_vec(),
+            input,
+            &partition_by_sort_keys,
+            &self.ordered_partition_by_indices,
+            reservation,
+            BaselineMetrics::new(&self.metrics, partition),
+        )?))
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
+    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+        self.window.child_stats_requests(partition)
+    }
+
+    fn statistics_from_inputs(
+        &self,
+        input_stats: &[Arc<Statistics>],
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
+        self.window.statistics_from_inputs(input_stats, args)
+    }
+
+    fn cardinality_effect(&self) -> CardinalityEffect {
+        self.window.cardinality_effect()
+    }
+}
+
+struct CometWindowAggStream {
+    schema: SchemaRef,
+    input_schema: SchemaRef,
+    input: SendableRecordBatchStream,
+    window_expr: Vec<Arc<dyn WindowExpr>>,
+    /// The `PARTITION BY` keys, in the order `WindowAggExec` evaluates them.
+    partition_by_sort_keys: Vec<PhysicalSortExpr>,
+    /// The rows of the last window partition seen so far, which may continue in the next
+    /// batch. Without `PARTITION BY`, every row seen so far.
+    buffered: Vec<RecordBatch>,
+    /// Counts each buffer that `buffered` retains once, however many batches share it.
+    buffered_memory: RecordBatchMemoryCounter,
+    reservation: MemoryReservation,
+    baseline_metrics: BaselineMetrics,
+    finished: bool,
+}
+
+impl CometWindowAggStream {
+    fn try_new(
+        schema: SchemaRef,
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+        input: SendableRecordBatchStream,
+        partition_by_sort_keys: &[PhysicalSortExpr],
+        ordered_partition_by_indices: &[usize],
+        reservation: MemoryReservation,
+        baseline_metrics: BaselineMetrics,
+    ) -> Result<Self> {
+        // The same check as `WindowAggStream::new`, which also makes the indexing below safe.
+        if window_expr[0].partition_by().len() != ordered_partition_by_indices.len() {
+            return internal_err!("All partition by columns should have an ordering");
+        }
+        let partition_by_sort_keys = ordered_partition_by_indices
+            .iter()
+            .map(|idx| partition_by_sort_keys[*idx].clone())
+            .collect();
+        Ok(Self {
+            schema,
+            input_schema: input.schema(),
+            input,
+            window_expr,
+            partition_by_sort_keys,
+            buffered: vec![],
+            buffered_memory: RecordBatchMemoryCounter::new(),
+            reservation,
+            baseline_metrics,
+            finished: false,
+        })
+    }
+
+    fn evaluate_partition_keys(&self, batch: &RecordBatch) -> Result<Vec<SortColumn>> {
+        self.partition_by_sort_keys
+            .iter()
+            .map(|key| key.evaluate_to_sort_column(batch))
+            .collect()
+    }
+
+    /// Buffers `batch`, then evaluates the window partitions that it completes.
+    fn push_batch(&mut self, batch: RecordBatch) -> Result<Option<RecordBatch>> {
+        let num_rows = batch.num_rows();
+        if num_rows == 0 {
+            return Ok(None);
+        }
+        let added = self.buffered_memory.count_batch(&batch);
+        self.reservation.try_grow(added).map_err(with_oom_context)?;
+        if self.partition_by_sort_keys.is_empty() {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        // Where the batch's last window partition starts. It may continue in the next batch.
+        let keys = self.evaluate_partition_keys(&batch)?;
+        let open_start = evaluate_partition_ranges(num_rows, &keys)?
+            .last()
+            .map_or(0, |range| range.start);
+        if open_start == 0 && self.continues_buffered_partition(&batch)? {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        // Every row before `open_start`, and every row buffered before this batch, belongs to a
+        // window partition that has ended.
+        let mut complete = std::mem::take(&mut self.buffered);
+        if open_start > 0 {
+            complete.push(batch.slice(0, open_start));
+        }
+        self.buffered
+            .push(batch.slice(open_start, num_rows - open_start));
+        self.evaluate(complete)
+    }
+
+    /// Whether the first row of `batch` belongs to the window partition of the last buffered
+    /// row.
+    fn continues_buffered_partition(&self, batch: &RecordBatch) -> Result<bool> {
+        let Some(last) = self.buffered.last() else {
+            return Ok(false);
+        };
+        let rows = concat_batches(
+            &self.input_schema,
+            [&last.slice(last.num_rows() - 1, 1), &batch.slice(0, 1)],
+        )?;
+        let keys = self.evaluate_partition_keys(&rows)?;
+        Ok(evaluate_partition_ranges(2, &keys)?.len() == 1)
+    }
+
+    /// Evaluates the window expressions over `batches`, which hold whole window partitions,
+    /// and returns their rows with the window columns appended.
+    fn evaluate(&mut self, batches: Vec<RecordBatch>) -> Result<Option<RecordBatch>> {
+        if batches.is_empty() {
+            return Ok(None);
+        }
+        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+        let _timer = elapsed_compute.timer();
+
+        // Concatenating a single batch is zero-copy. Otherwise reserve the copy before making
+        // it, while `batches` are still reserved.
+        if batches.len() > 1 {
+            let mut copy_size = 0;
+            for batch in &batches {
+                for column in batch.columns() {
+                    copy_size += column.to_data().get_slice_memory_size()?;

Review Comment:
   [P2] Could the copy estimate account for nested slice offsets before summing `get_slice_memory_size()`? For a sliced `ListArray`, this method recursively counts the entire child array, including values outside the slice. Multiple batches sharing that child therefore charge it repeatedly, although concatenation copies only the referenced values. A reproduced `COUNT(k) OVER ()` with a list payload fails under a 16 MiB pool even though its retained input and actual concatenated copy fit together. This introduces task failures based on batch representation rather than required memory. Please use offset-aware nested sizing and add a shared-list-slice regression.
   
   Evidence: Executed an isolated harness against the unchanged exact-head operator. Constructed 10,000 rows with `k = 0..9999` and an `a: List<Int64>` containing 0..63 per row, then supplied ten `original.slice(i * 1000, 1000)` batches to `COUNT(k) OVER ()` with `GreedyMemoryPool::new(16 * 1024 * 1024)`. Unique retained input measured 8,508,612 bytes and the actual concatenated batch measured 5,265,536 bytes, totaling less than 16 MiB. The new estimate requested 51,320,040 bytes and failed with `Resources exhausted: Failed to allocate additional 48.9 MB for CometWindowAggExec[0]`. The previous `WindowAggExec` succeeded. Materializing the same slices independently with `take_record_batch` also made the new operator succeed and return identical results.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-5874119437

   Thanks for this change. Emitting each window partition as soon as the sorted input moves past it matches Spark's partition-at-a-time `WindowExec` (minus spilling), and it is similar to how StarRocks' analytic operator processes clustered input. The differential test against `WindowAggExec` and the mutation-checked memory tests are convincing. A few things I'd like to see addressed.
   
   **Memory pool interplay**
   
   1. Under the default `fair_unified` pool, each consumer is limited to the pool divided by the number of registered consumers (`native/core/src/execution/memory_pools/fair_pool.rs:165-171`). A sort then window task has three (`ExternalSorter`, `ExternalSorterMerge` and the window). The window reserves a partition and its concatenated copy together, so by my reading the largest partition that fits is about a sixth of the pool. That share is the binding limit when a task has more consumers than there are running tasks, such as a skewed straggler, which is exactly the case this PR is for. In that case, registering in `execute()` (`native/core/src/execution/operators/window_agg.rs:191`) also cuts the sort's share from pool/2 to pool/3 while it reads its input, even when every window partition is tiny. I have not measured this. A few suggestions:
      - Register the consumer when the first batch is buffered. The sort has consumed its input by then. #5465 would be the real fix in the pool.
      - Have the error hint (`window_agg.rs:428-431`) and the tuning doc say that a partition needs roughly twice its size within its share.
      - Suggest `spark.comet.exec.memoryPool=greedy_unified` before disabling native windows, as `docs/source/user-guide/latest/migration-guide.md:88-90` already does. `spark.comet.exec.window.enabled=false` also moves bounded windows back to Spark, which may be worth saying.
   
   **Code**
   
   2. `evaluate` counts its inputs a second time (`window_agg.rs:332-344`). `self.buffered_memory` has already counted every buffer in `batches`, and `concat` only reuses input buffers, so `self.buffered_memory.count_batch(&batch)` should return the same number (and 0 for the zero-copy single-batch case). That would remove the second counter, the extra pass over every buffered batch, and the `batches.len() > 1` branch with its comment.
   3. Key evaluation, `evaluate_partition_ranges`, the continuation check and the reservation in `push_batch` run outside `elapsed_compute` (`window_agg.rs:329`). Batches that only extend the open partition record no time, while `WindowAggExec` times its range pass. Could we time `push_batch` and the final `evaluate` in `poll_next_inner` and drop the inner timer, since nested timers on one `Time` count twice? Also, `continues_buffered_partition` slices every column of the last buffered batch and re-evaluates its keys on every batch (`window_agg.rs:305`). Keeping the last row's key arrays from line 263 would avoid that and still use the same `partition` kernel.
   
   **Tests**
   
   4. A few suggestions:
      - The first new Scala test (`CometWindowExecSuite.scala:1517`) looks like a good fit for a SQL file test under `spark/src/test/resources/sql-tests/windows/`, with `-- Config: spark.comet.batchSize=7` and `INSERT ... SELECT ... FROM range(2000)`. Its second query (no `PARTITION BY`) could go. That path does not depend on the function, and it is already covered by "aggregate window function for all types" (`OVER()` at batch size 128) and by the Rust fuzz with `partitioned = false`.
      - Tests 2 and 3 need to stay in Scala, since `expect_error` requires Spark to fail as well. Test 3 (`:1596`) asserts DataFusion's `TrackConsumersPool` wording. Asserting on the new `CometWindowAggExec cannot spill` hint instead would also prove the hint reaches Spark. The `Context` text does survive into `CometNativeException`.
      - In the Rust tests, the interleaved empty-batch block in `empty_input_and_empty_batches` (`window_agg.rs:727-737`) repeats the fuzz, because `random_slices` already produces empty batches. The zero-row cases at 724-725 are the part the fuzz cannot produce. The `peak` assertions (750, 835) repeat what the neighboring refusals prove, so `run_with_limit` could return just the batches and keep its `reserved == 0` check.
   
   **Tracking and docs**
   
   5. The removal condition at `window_agg.rs:38` has no open upstream work behind it. apache/datafusion#22947 and apache/datafusion#23207 (reservation only, still buffering the whole input) were both closed unmerged. Would you consider offering this stream upstream as the no-spill first step of apache/datafusion#22946 and linking it there? Otherwise Comet carries its own copy of `WindowAggStream` and about ten delegated methods through every DataFusion upgrade.
   6. #6253 covers both window operators, and `BoundedWindowAggExec` still reserves nothing. Its `RANGE` frames keep a whole group of peer rows, so a skewed `ORDER BY` value can still grow untracked. Could this be `Part of #6253`, or could we file a follow-up for `BoundedWindowAggExec`?
   7. A few docs now leave the window out:
      - `docs/source/user-guide/latest/tuning/memory.md:44-46` lists what the pool tracks.
      - `docs/source/contributor-guide/memory_management.md:556-558` names only `ShuffledHashJoin` as an operator that cannot spill.
      - `docs/source/contributor-guide/memory_management.md:366-369` names only sort and hash join as reserving imported batches. `OVER ()` has no sort below it, so the window can buffer scan or shuffle batches directly.
   
   **CI and benchmarks**
   
   - The Spark SQL suites were skipped in PR CI. Since this changes the planner and adds a native operator, could we apply `run-spark-4.1-tests` before it goes to the merge queue?
   - If the standalone bench ran on a DataFusion pool, two costs are not in its numbers: extra sort spills from the share dilution above, and three JNI-backed pool calls for every batch that completes a partition. A TPC-DS window query under `fair_unified`, with a pool small enough to make the sort spill and compared against `main`, would settle both.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-6447591824

   > @comphead do you have any more feedback on this PR?
   
   actually no, but I'm thinking how we going to align the overlapping code after DF56


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#discussion_r4116382466


##########
native/core/src/execution/operators/window_agg.rs:
##########
@@ -0,0 +1,800 @@
+// 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.
+
+//! Window functions that need their whole window partition before they produce a value.
+//!
+//! `PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at
+//! `UNBOUNDED FOLLOWING` cannot run in DataFusion's `BoundedWindowAggExec`. DataFusion's
+//! `WindowAggExec` runs them by buffering its whole input, which is the task's entire Spark
+//! partition, and concatenating it before it evaluates anything. It reserves memory for
+//! neither, and it cannot spill (apache/datafusion#22946).
+//!
+//! [`CometWindowAggExec`] evaluates the same window expressions, one window partition at a
+//! time as `WindowAggExec` does, with two differences:
+//!
+//! - The input is sorted by the `PARTITION BY` keys, so once a batch starts a new window
+//!   partition, every earlier one is complete. Those are evaluated and emitted right away, and
+//!   only the window partition that may continue is kept. Without `PARTITION BY`, the whole
+//!   input is one window partition.
+//! - The buffered batches, and the copy made when concatenating them for evaluation, are
+//!   reserved against the task's memory pool. The operator still cannot spill, so a window
+//!   partition that does not fit fails the task with a memory error instead of growing
+//!   untracked.
+
+use std::fmt::Formatter;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::array::{Array, RecordBatch};
+use arrow::compute::{concat, concat_batches, SortColumn};
+use arrow::datatypes::SchemaRef;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::utils::memory::RecordBatchMemoryCounter;
+use datafusion::common::utils::{evaluate_partition_ranges, transpose};
+use datafusion::common::{internal_err, DataFusionError, Result, Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::window::WindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr, PhysicalSortExpr};
+use datafusion::physical_plan::execution_plan::{
+    CardinalityEffect, ChildrenPropertiesMode, EmissionType, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::windows::{get_ordered_partition_by_indices, WindowAggExec};
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan,
+    InputDistributionRequirements, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
+};
+use futures::{ready, Stream, StreamExt};
+
+#[derive(Debug)]
+pub(crate) struct CometWindowAggExec {
+    /// DataFusion's operator for the same window expressions. It provides the output schema,
+    /// the plan properties and the ordered `PARTITION BY` keys, and is never executed.
+    window: WindowAggExec,
+    /// How the `PARTITION BY` expressions map onto the input ordering, as `WindowAggExec`
+    /// computes it.
+    ordered_partition_by_indices: Vec<usize>,
+    cache: Arc<PlanProperties>,
+    metrics: ExecutionPlanMetricsSet,
+}
+
+impl CometWindowAggExec {
+    pub(crate) fn try_new(
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+        input: Arc<dyn ExecutionPlan>,
+        can_repartition: bool,
+    ) -> Result<Self> {
+        Self::from_datafusion(WindowAggExec::try_new(window_expr, input, can_repartition)?)
+    }
+
+    fn from_datafusion(window: WindowAggExec) -> Result<Self> {
+        let partition_by = window.window_expr()[0].partition_by();
+        let ordered_partition_by_indices =
+            get_ordered_partition_by_indices(partition_by, window.input())?;
+        // Unlike `WindowAggExec`, the output is emitted as window partitions complete.
+        let emission_type = if partition_by.is_empty() {
+            EmissionType::Final
+        } else {
+            EmissionType::Incremental
+        };
+        let cache = Arc::new(
+            window
+                .properties()
+                .as_ref()
+                .clone()
+                .with_emission_type(emission_type),
+        );
+        Ok(Self {
+            window,
+            ordered_partition_by_indices,
+            cache,
+            metrics: ExecutionPlanMetricsSet::new(),
+        })
+    }
+}
+
+impl DisplayAs for CometWindowAggExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        // `WindowAggExec` starts these formats with its own name.
+        if matches!(t, DisplayFormatType::Default | DisplayFormatType::Verbose) {
+            write!(f, "Comet")?;
+        }
+        self.window.fmt_as(t, f)
+    }
+}
+
+impl ExecutionPlan for CometWindowAggExec {
+    fn name(&self) -> &str {
+        "CometWindowAggExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.cache
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![self.window.input()]
+    }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        self.window.apply_expressions(f)
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        self.window.maintains_input_order()
+    }
+
+    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+        self.window.required_input_ordering()
+    }
+
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        self.window.input_distribution_requirements()
+    }
+
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return internal_err!("CometWindowAggExec requires one child");
+        }
+        let replaced = Arc::new(self.window.clone()).replace_children(children, options)?;
+        let Some(window) = replaced.downcast_ref::<WindowAggExec>() else {
+            return internal_err!("WindowAggExec child replacement changed its plan type");
+        };
+        Ok(Arc::new(Self::from_datafusion(window.clone())?))
+    }
+
+    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 execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let partition_by_sort_keys = self.window.partition_by_sort_keys()?;
+        let input = self
+            .window
+            .input()
+            .execute(partition, Arc::clone(&context))?;
+        let reservation = MemoryConsumer::new(format!("{}[{partition}]", self.name()))
+            .register(context.memory_pool());
+        Ok(Box::pin(CometWindowAggStream::try_new(
+            self.window.schema(),
+            self.window.window_expr().to_vec(),
+            input,
+            &partition_by_sort_keys,
+            &self.ordered_partition_by_indices,
+            reservation,
+            BaselineMetrics::new(&self.metrics, partition),
+        )?))
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
+    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+        self.window.child_stats_requests(partition)
+    }
+
+    fn statistics_from_inputs(
+        &self,
+        input_stats: &[Arc<Statistics>],
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
+        self.window.statistics_from_inputs(input_stats, args)
+    }
+
+    fn cardinality_effect(&self) -> CardinalityEffect {
+        self.window.cardinality_effect()
+    }
+}
+
+struct CometWindowAggStream {
+    schema: SchemaRef,
+    input_schema: SchemaRef,
+    input: SendableRecordBatchStream,
+    window_expr: Vec<Arc<dyn WindowExpr>>,
+    /// The `PARTITION BY` keys, in the order `WindowAggExec` evaluates them.
+    partition_by_sort_keys: Vec<PhysicalSortExpr>,
+    /// The rows of the last window partition seen so far, which may continue in the next
+    /// batch. Without `PARTITION BY`, every row seen so far.
+    buffered: Vec<RecordBatch>,
+    /// Counts each buffer that `buffered` retains once, however many batches share it.
+    buffered_memory: RecordBatchMemoryCounter,
+    reservation: MemoryReservation,
+    baseline_metrics: BaselineMetrics,
+    finished: bool,
+}
+
+impl CometWindowAggStream {
+    fn try_new(
+        schema: SchemaRef,
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+        input: SendableRecordBatchStream,
+        partition_by_sort_keys: &[PhysicalSortExpr],
+        ordered_partition_by_indices: &[usize],
+        reservation: MemoryReservation,
+        baseline_metrics: BaselineMetrics,
+    ) -> Result<Self> {
+        // The same check as `WindowAggStream::new`, which also makes the indexing below safe.
+        if window_expr[0].partition_by().len() != ordered_partition_by_indices.len() {
+            return internal_err!("All partition by columns should have an ordering");
+        }
+        let partition_by_sort_keys = ordered_partition_by_indices
+            .iter()
+            .map(|idx| partition_by_sort_keys[*idx].clone())
+            .collect();
+        Ok(Self {
+            schema,
+            input_schema: input.schema(),
+            input,
+            window_expr,
+            partition_by_sort_keys,
+            buffered: vec![],
+            buffered_memory: RecordBatchMemoryCounter::new(),
+            reservation,
+            baseline_metrics,
+            finished: false,
+        })
+    }
+
+    fn evaluate_partition_keys(&self, batch: &RecordBatch) -> Result<Vec<SortColumn>> {
+        self.partition_by_sort_keys
+            .iter()
+            .map(|key| key.evaluate_to_sort_column(batch))
+            .collect()
+    }
+
+    /// Buffers `batch`, then evaluates the window partitions that it completes.
+    fn push_batch(&mut self, batch: RecordBatch) -> Result<Option<RecordBatch>> {
+        let num_rows = batch.num_rows();
+        if num_rows == 0 {
+            return Ok(None);
+        }
+        let added = self.buffered_memory.count_batch(&batch);
+        self.reservation.try_grow(added).map_err(with_oom_context)?;
+        if self.partition_by_sort_keys.is_empty() {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        // Where the batch's last window partition starts. It may continue in the next batch.
+        let keys = self.evaluate_partition_keys(&batch)?;
+        let open_start = evaluate_partition_ranges(num_rows, &keys)?
+            .last()
+            .map_or(0, |range| range.start);
+        if open_start == 0 && self.continues_buffered_partition(&batch)? {
+            self.buffered.push(batch);
+            return Ok(None);
+        }
+
+        // Every row before `open_start`, and every row buffered before this batch, belongs to a
+        // window partition that has ended.
+        let mut complete = std::mem::take(&mut self.buffered);
+        if open_start > 0 {
+            complete.push(batch.slice(0, open_start));
+        }
+        self.buffered
+            .push(batch.slice(open_start, num_rows - open_start));
+        self.evaluate(complete)
+    }
+
+    /// Whether the first row of `batch` belongs to the window partition of the last buffered
+    /// row.
+    fn continues_buffered_partition(&self, batch: &RecordBatch) -> Result<bool> {
+        let Some(last) = self.buffered.last() else {
+            return Ok(false);
+        };
+        let rows = concat_batches(
+            &self.input_schema,
+            [&last.slice(last.num_rows() - 1, 1), &batch.slice(0, 1)],
+        )?;
+        let keys = self.evaluate_partition_keys(&rows)?;
+        Ok(evaluate_partition_ranges(2, &keys)?.len() == 1)
+    }
+
+    /// Evaluates the window expressions over `batches`, which hold whole window partitions,
+    /// and returns their rows with the window columns appended.
+    fn evaluate(&mut self, batches: Vec<RecordBatch>) -> Result<Option<RecordBatch>> {
+        if batches.is_empty() {
+            return Ok(None);
+        }
+        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+        let _timer = elapsed_compute.timer();
+
+        // Concatenating a single batch is zero-copy. Otherwise reserve the copy before making
+        // it, while `batches` are still reserved.
+        if batches.len() > 1 {
+            let mut copy_size = 0;
+            for batch in &batches {
+                for column in batch.columns() {
+                    copy_size += column.to_data().get_slice_memory_size()?;

Review Comment:
   Thanks, good catch. The copy is now reserved right after concatenating, counting only the buffers it doesn't share with the input batches, which is exact for nested types instead of estimated per slice. I added `batches_sliced_from_one_list_array_are_charged_once` with your repro's shape; it's refused with the old estimate and passes now.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-6345846784

   Here are the DataFusion items I filed for window execution:
   
   Spilling and memory
   - apache/datafusion#25637 is the spilling coverage epic. It tracks `WindowAggExec` and `BoundedWindowAggExec` as non-spilling next to Spark and DuckDB. The upstream spill work for `WindowAggExec` is apache/datafusion#22946.
   
   `BoundedWindowAggExec` and the value functions, which relate to #7092
   - apache/datafusion#26537 proposes memoizing `LAST_VALUE` and `NTH_VALUE` with `IGNORE NULLS` so `BoundedWindowAggExec` can prune its buffer.
   - apache/datafusion#26539: `LAST_VALUE` and `NTH_VALUE` with `IGNORE NULLS` keep the whole partition buffered.
   - apache/datafusion#26541: `BoundedWindowAggExec` has a high fixed cost per partition when partitions are small.
   
   Performance
   - apache/datafusion#26538: `NTH_VALUE(expr, -n) IGNORE NULLS` is quadratic in the partition size.
   - apache/datafusion#26540: `FIRST_VALUE`, `LAST_VALUE` and `NTH_VALUE` build a `ScalarValue` for every row.
   - apache/datafusion#26532 and apache/datafusion#26533 are the issue and the fix that made `LAST_VALUE ... IGNORE NULLS` linear. #26533 merged on 2026-10-21.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on code in PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#discussion_r4146848657


##########
docs/source/user-guide/latest/tuning/operators.md:
##########
@@ -59,6 +59,17 @@ boundaries. A shuffled hash join can still filter probe batches after shuffle, b
 its filter back to an earlier scan stage. Compare the [runtime-filter and scan metrics](../metrics.md#hash-joins)
 with the setting disabled to distinguish reduced hash-probe work from reader I/O savings.
 
+## Window Functions
+
+`PERCENT_RANK`, `CUME_DIST`, `NTILE`, and aggregates whose frame ends at `UNBOUNDED FOLLOWING`, which includes an
+aggregate over `PARTITION BY` with no `ORDER BY` such as `sum(x) OVER (PARTITION BY k)`, need a whole window partition
+before they return anything. Comet buffers one window partition at a time for them and reserves it from its native
+memory pool, but it does not yet provide spill-to-disk for them. A window partition that does not fit in its share of
+the pool, for example because of a heavily skewed key or a window without `PARTITION BY`, fails the task with a memory

Review Comment:
   This section describes a change in outcome. A window partition that does not fit its share now fails the task, where before it ran on untracked memory. Would it make sense to add an entry for it to the upgrade guide? `docs/source/contributor-guide/config_conventions.md` (Changing the Behavior of an Existing Config) asks for one, plus a `spark.comet.legacy.*` key. I'm not sure a key is expected for a memory accounting fix, so a note like the 1.1.0 "Memory Pool Limits" entry may be enough.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-5955550875

   I would like to do a deeper review of this PR before it is merged.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: native windows reserve the rows they buffer and buffer one window partition at a time [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #6279:
URL: https://github.com/apache/datafusion-comet/pull/6279#issuecomment-6439980243

   @comphead do you have any more feedback on this PR?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org