sunchao commented on code in PR #4870: URL: https://github.com/apache/datafusion-comet/pull/4870#discussion_r3855658566
########## native/core/src/execution/operators/rank_limit.rs: ########## @@ -0,0 +1,637 @@ +// 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. + +//! Streaming top-K per partition operator for Spark's `WindowGroupLimitExec`. +//! +//! The child stream must be sorted by `[partition_keys..., order_keys...]`. +//! Spark's `WindowGroupLimitExec.requiredChildOrdering` guarantees this via +//! `EnsureRequirements`; the operator relies on the injected sort so a single +//! streaming pass decides emit-or-drop per row. Tie behavior matches Spark's +//! `RankLimitIterator` / `SimpleLimitIterator` exactly. +//! +//! ROW_NUMBER without PARTITION BY is served by a plain `LocalLimitExec` in the +//! planner and never reaches this operator. + +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::{ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch}; +use arrow::compute::filter_record_batch; +use arrow::datatypes::SchemaRef; +use arrow::row::{OwnedRow, RowConverter, Rows, SortField}; +use datafusion::common::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{ + LexOrdering, OrderingRequirements, PhysicalExpr, PhysicalSortExpr, +}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, +}; +use futures::{Stream, StreamExt}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WindowFnKind { + RowNumber, + Rank, + DenseRank, +} + +#[derive(Debug)] +pub struct PartitionedRankLimitExec { + input: Arc<dyn ExecutionPlan>, + /// PARTITION BY expressions. Empty means "no PARTITION BY" (global top-K + /// within each input DataFusion partition). + partition_keys: Vec<PhysicalSortExpr>, + /// ORDER BY expressions. Empty means "no ORDER BY" and every row within a + /// partition ties. + order_keys: Vec<PhysicalSortExpr>, + fetch: usize, + kind: WindowFnKind, + cache: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl PartitionedRankLimitExec { + pub fn try_new( + input: Arc<dyn ExecutionPlan>, + partition_keys: Vec<PhysicalSortExpr>, + order_keys: Vec<PhysicalSortExpr>, + fetch: usize, + kind: WindowFnKind, + ) -> Result<Self> { + let cache = Arc::new(Self::compute_properties( + &input, + &partition_keys, + &order_keys, + )?); + Ok(Self { + input, + partition_keys, + order_keys, + fetch, + kind, + cache, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + partition_keys: &[PhysicalSortExpr], + order_keys: &[PhysicalSortExpr], + ) -> Result<PlanProperties> { + let mut eq_properties = input.equivalence_properties().clone(); + if let Some(ordering) = full_ordering(partition_keys, order_keys) { + eq_properties.reorder(ordering)?; + } + Ok(PlanProperties::new( + eq_properties, + input.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } +} + +/// `[partition_keys..., order_keys...]` as a single `LexOrdering`, or `None` when both lists +/// are empty. Dedup by `LexOrdering::new` is fine here because this ordering is only used to +/// declare equivalence properties and the input-ordering requirement; the streaming operator +/// itself operates on the un-deduped `partition_keys` / `order_keys` slices so a duplicate +/// (e.g. `PARTITION BY a, a`) never turns into an internal error. +fn full_ordering( + partition_keys: &[PhysicalSortExpr], + order_keys: &[PhysicalSortExpr], +) -> Option<LexOrdering> { + let sort_exprs: Vec<PhysicalSortExpr> = partition_keys + .iter() + .chain(order_keys.iter()) + .cloned() + .collect(); + LexOrdering::new(sort_exprs) +} + +impl DisplayAs for PartitionedRankLimitExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let partition = self + .partition_keys + .iter() + .map(|e| e.to_string()) + .collect::<Vec<_>>() + .join(", "); + let order = self + .order_keys + .iter() + .map(|e| e.to_string()) + .collect::<Vec<_>>() + .join(", "); + write!( + f, + "CometPartitionedRankLimitExec: kind={:?}, fetch={}, partition_by=[{}], order_by=[{}]", + self.kind, self.fetch, partition, order + ) + } + DisplayFormatType::TreeRender => unimplemented!(), + } + } +} + +impl ExecutionPlan for PartitionedRankLimitExec { + fn name(&self) -> &str { + "CometPartitionedRankLimitExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + assert_eq!(children.len(), 1); + Ok(Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&children[0]), + self.partition_keys.clone(), + self.order_keys.clone(), + self.fetch, + self.kind, + )?)) + } + + // The operator's correctness depends on the input being sorted by + // `[partition_keys..., order_keys...]`. Spark's Catalyst injects the required sort above + // `WindowGroupLimitExec`, and Comet executes the deserialized plan directly without + // running any DataFusion physical optimizer pass, so this method is informational: it + // documents the ordering contract and shows up in `DisplayableExecutionPlan` output. It + // is not a safety net -- if the sort is missing upstream, results are wrong. + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![full_ordering(&self.partition_keys, &self.order_keys).map(OrderingRequirements::from)] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, context)?; + let schema = input.schema(); + + let partition_key = build_key_encoder(&self.partition_keys, &schema)?; + + // ROW_NUMBER's rank formula is just the running count, so it never reads the + // ORDER BY key. Skip building the converter and evaluating order columns. + // For RANK/DENSE_RANK, the encoder drives tie detection on the ORDER BY suffix. + // When the suffix is empty (query has no ORDER BY) `build_key_encoder` returns + // `None` and every row within a partition ties. + let order_key = if self.kind == WindowFnKind::RowNumber { + None + } else { + build_key_encoder(&self.order_keys, &schema)? + }; + + Ok(Box::pin(RankLimitStream { + input, + schema, + partition_key, + order_key, + limit: self.fetch as u64, + kind: self.kind, + baseline_metrics: BaselineMetrics::new(&self.metrics, partition), + prev_partition: None, + prev_order: None, + rank: 0, + count: 0, + partition_exhausted: false, + })) + } +} + +/// Row-encoded key for either PARTITION BY or ORDER BY columns. Only constructed +/// when the corresponding expression list is non-empty. +struct KeyEncoder { + converter: RowConverter, + exprs: Vec<Arc<dyn PhysicalExpr>>, +} + +impl KeyEncoder { + fn encode(&self, batch: &RecordBatch) -> Result<Rows> { + let num_rows = batch.num_rows(); + let cols: Vec<ArrayRef> = self + .exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::<Result<_>>()?; + Ok(self.converter.convert_columns(&cols)?) + } +} + +fn build_key_encoder(exprs: &[PhysicalSortExpr], schema: &SchemaRef) -> Result<Option<KeyEncoder>> { + if exprs.is_empty() { + return Ok(None); + } + let sort_fields = build_sort_fields(exprs, schema)?; + let converter = RowConverter::new(sort_fields)?; + let exprs = exprs.iter().map(|e| Arc::clone(&e.expr)).collect(); + Ok(Some(KeyEncoder { converter, exprs })) +} + +fn build_sort_fields(ordering: &[PhysicalSortExpr], schema: &SchemaRef) -> Result<Vec<SortField>> { + ordering + .iter() + .map(|e| { + Ok(SortField::new_with_options( + e.expr.data_type(schema)?, + e.options, + )) + }) + .collect() +} + +struct RankLimitStream { + input: SendableRecordBatchStream, + schema: SchemaRef, + /// `None` when there is no PARTITION BY (global top-K per DF input partition). + partition_key: Option<KeyEncoder>, + /// `None` when there is no ORDER BY (every row within a partition ties), and + /// always `None` for ROW_NUMBER (rank formula never reads order keys). + order_key: Option<KeyEncoder>, + limit: u64, + kind: WindowFnKind, + baseline_metrics: BaselineMetrics, + + // Per-partition streaming state, persisted across batches. + prev_partition: Option<OwnedRow>, + prev_order: Option<OwnedRow>, + /// Rank of the most recently seen row (0-indexed). Only meaningful when + /// `prev_order.is_some()` -- the two reads below both sit past the point where + /// `prev_order` was set for the current partition. + rank: u64, + /// 0-indexed cursor into the current partition for rank arithmetic. Advances only on + /// non-exhausted rows -- once `partition_exhausted` fires, subsequent rows in the same + /// partition skip the increment. `count` is thus NOT rows-seen; it freezes at + /// `first_dropped_at` for the tail of an exhausted partition. + count: u64, + /// Set once `this_rank >= limit` inside the current partition and cleared when a new + /// partition starts. Mirrors Spark's `GroupedLimitIterator.skipRemainingRows`: for a + /// partition already past the limit we skip order-key encoding, tie detection, and + /// rank arithmetic on the remaining rows. + partition_exhausted: bool, +} + +impl RankLimitStream { + /// Filter a batch to the rows this operator keeps. `Ok(None)` means the batch produced no + /// output; the caller must not surface it downstream. Passing an empty batch also returns + /// `Ok(None)` (nothing to emit) rather than an empty pass-through, so the caller never + /// has to strip zero-row batches. + fn process_batch(&mut self, batch: &RecordBatch) -> Result<Option<RecordBatch>> { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(None); + } + + let partition_rows = self + .partition_key + .as_ref() + .map(|k| k.encode(batch)) + .transpose()?; + // Lazily encoded: skipped entirely for a batch that is wholly inside an already- + // exhausted partition, so a giant skewed partition after the limit costs O(rows) + // partition-key checks instead of O(rows) full row encodings. + let mut order_rows: Option<Rows> = None; + + let mut mask_builder = BooleanBufferBuilder::new(num_rows); + let mut kept: usize = 0; + // Position of the first dropped row in this batch. When `kept == first_dropped_at`, + // every kept row lies at positions `0..kept`, so the output is `batch.slice(0, kept)` + // -- one `Arc` reslice per column, no bitmap scan or per-column filter kernel. + let mut first_dropped_at: Option<usize> = None; + for i in 0..num_rows { + // Only a PARTITION-BY-shaped stream has partition boundaries. With no + // PARTITION BY the whole stream is one partition, so state accumulates + // across every row and no reset is needed. + if let Some(pr) = &partition_rows { + let same_partition = + matches!(&self.prev_partition, Some(prev) if prev.row() == pr.row(i)); + if !same_partition { + self.prev_partition = Some(pr.row(i).owned()); + self.prev_order = None; + self.rank = 0; + self.count = 0; + self.partition_exhausted = false; + } + } + + if self.partition_exhausted { + mask_builder.append(false); + continue; Review Comment: **[P1] Record skipped rows before taking the prefix shortcut** Could we set `first_dropped_at` when this branch drops a row? If a batch starts inside a partition exhausted by the previous batch, its leading drops are not recorded. With `fetch = 1` and sorted `(partition, order)` rows: ```text Batch 0: (1,10), (1,20), (1,30), (1,40) Batch 1: (1,50), (2,10), (2,20), (3,10) Expected output: (1,10), (2,10), (3,10) Actual output: (1,10), (1,50), (2,10) ``` Batch 1 builds the correct mask `[false, true, false, true]`, but records `first_dropped_at = 2` and `kept = 2`. The shortcut at lines 419–420 therefore returns `batch.slice(0, 2)` instead of applying that mask. This permanently loses partition 3's top row, which a subsequent Final limiter or window/filter cannot recover. I reproduced this for `ROW_NUMBER`, `RANK`, and `DENSE_RANK` by compiling and executing the unchanged `rank_limit.rs` from `77702dba` in an isolated harness with DataFusion 54.1.0 and Arrow 58.4.0. The six existing tests pass and all three regression cases fail. Recording the first dropped position in this branch makes all nine pass. Could we include this mixed-partition batch case in the native tests? -- 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]
