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: [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]
