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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
