sunchao commented on code in PR #4870: URL: https://github.com/apache/datafusion-comet/pull/4870#discussion_r3832322832
########## native/core/src/execution/operators/rank_limit.rs: ########## @@ -0,0 +1,609 @@ +// 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 `count > 0`. + rank: u64, + /// Total rows seen in the current partition (also 0-indexed cursor). + 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 { + fn process_batch(&mut self, batch: &RecordBatch) -> Result<RecordBatch> { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(batch.clone()); + } + + 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; + for i in 0..num_rows { + let same_partition = match &partition_rows { + Some(pr) => matches!(&self.prev_partition, Some(prev) if prev.row() == pr.row(i)), + // No PARTITION BY: state accumulates across every row, resetting + // only on the very first row of the stream. + None => self.count > 0, + }; + if !same_partition { + if let Some(pr) = &partition_rows { + 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); + self.count += 1; + continue; + } + + if order_rows.is_none() { + if let Some(k) = self.order_key.as_ref() { + order_rows = Some(k.encode(batch)?); + } + } + + // Whether this row's ORDER BY key ties with the previous emitted row. `false` + // on the first row of a partition and (vacuously) when there is no ORDER BY -- + // `prev_order` stays `None` across the whole partition in that case. + let ties_with_prev = matches!( + (&self.prev_order, &order_rows), + (Some(prev_o), Some(rows)) if prev_o.row() == rows.row(i) + ); + + let this_rank: u64 = + if self.prev_order.is_none() && self.kind != WindowFnKind::RowNumber { + // First row of a partition ranks 0 under RANK/DENSE_RANK. + 0 + } else { + match self.kind { + WindowFnKind::RowNumber => self.count, + _ if ties_with_prev => self.rank, + WindowFnKind::DenseRank => self.rank + 1, + WindowFnKind::Rank => self.count, + } + }; + + let keep = this_rank < self.limit; + mask_builder.append(keep); + if keep { + kept += 1; + } else { + // `this_rank` is monotonically nondecreasing within a partition for all three + // kinds (ROW_NUMBER: strictly, RANK / DENSE_RANK: nondecreasing), so once + // `keep` flips false every remaining row of this partition is dropped. + self.partition_exhausted = true; + } + + self.rank = this_rank; + // Only clone into an `OwnedRow` when the key actually changed. Under + // RANK/DENSE_RANK the common tail of a partition is a run of ties, + // so this avoids O(rows) heap allocations there. + if let Some(rows) = &order_rows { + if !ties_with_prev { + self.prev_order = Some(rows.row(i).owned()); + } + } + self.count += 1; + } + + if kept == num_rows { + return Ok(batch.clone()); + } + if kept == 0 { + return Ok(RecordBatch::new_empty(Arc::clone(&self.schema))); + } + let mask = BooleanArray::new(mask_builder.finish(), None); + Ok(filter_record_batch(batch, &mask)?) + } +} + +impl Stream for RankLimitStream { + type Item = Result<RecordBatch>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + loop { + match self.input.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + let processed = { + let _timer = self.baseline_metrics.elapsed_compute().timer(); + self.process_batch(&batch) Review Comment: **[P1] Release the metrics borrow before calling `process_batch`** Could we clone `elapsed_compute()` into a local before creating this timer? `ScopedTimerGuard` holds an immutable borrow of `self.baseline_metrics` until the end of this block, while `self.process_batch(&batch)` requires a mutable borrow of `self`. Compiling this file at `7ef79eb0` produces `E0502: cannot borrow self as mutable because it is also borrowed as immutable` on the `process_batch` call. This prevents native builds even when WindowGroupLimit is disabled. The preceding `add9758e` version compiles against the same dependencies. This small change compiles while preserving the timing: ```rust let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let _timer = elapsed_compute.timer(); self.process_batch(&batch) ``` ########## spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala: ########## @@ -0,0 +1,180 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} +import org.apache.spark.sql.catalyst.plans.physical.Partitioning +import org.apache.spark.sql.execution.SparkPlan + +import com.google.common.base.Objects + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass} +import org.apache.comet.serde.OperatorOuterClass.{Operator, RankLikeFunction} +import org.apache.comet.serde.QueryPlanSerde.exprToProto +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** + * Serde for Spark's `WindowGroupLimitExec` (Spark 3.5+, SPARK-37099). Handles ROW_NUMBER, RANK, + * and DENSE_RANK natively. ROW_NUMBER without PARTITION BY collapses to a `LocalLimitExec` over + * the Spark-sorted child. Every other combination (ROW_NUMBER partitioned, RANK/DENSE_RANK with + * or without PARTITION BY) maps onto Comet's streaming `PartitionedRankLimitExec`. + * + * The Scala type parameter is `SparkPlan` (not `WindowGroupLimitExec`) so this file stays + * compilable against Spark 3.4, where the exec class does not exist. Field extraction is + * delegated to the per-Spark-minor `ShimCometWindowGroupLimit`. + */ +object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { + + /** + * Fields extracted from a Spark `WindowGroupLimitExec` (Spark 3.5+). `mode` is a Spark-agnostic + * string ("Partial" or "Final") because Spark's `WindowGroupLimitMode` type does not exist on + * Spark 3.4, and the enclosing file must compile on that profile. + */ + case class Fields( + partitionSpec: Seq[Expression], + orderSpec: Seq[SortOrder], + rankLikeFunction: RankLikeFunction, + limit: Int, + mode: String) + + override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( + CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED) + + override def convert( + op: SparkPlan, + builder: Operator.Builder, + childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { + // Shim returns `None` for a Spark 3.4 plan (WGL does not exist) or if a future Spark + // introduces a rank-like function this shim does not know about. In both cases fall back + // to Spark rather than throwing, so a working query stays working across Spark upgrades. + val fields = ShimCometWindowGroupLimit.extract(op) match { + case Some(f) => f + case None => + withFallbackReason(op, "WindowGroupLimit: unsupported rank-like function") + return None + } + + if (fields.limit <= 0) { + // Spark's optimizer collapses limit <= 0 to an empty LocalRelation, but guard anyway. + withFallbackReason(op, s"WindowGroupLimit: non-positive limit ${fields.limit}") + return None + } + + val childOutput = op.children.head.output + val partitionProtos = fields.partitionSpec.map(e => e -> exprToProto(e, childOutput)) + val orderProtos = fields.orderSpec.map(e => e -> exprToProto(e, childOutput)) Review Comment: **[P2] Reject collated keys before native window-group pruning** Could we reject non-default string collations on these keys before lowering the operator? For an ordinary, uncollated Parquet table containing `(grp, s) = (1, 'A'), (1, 'a'), (1, 'b')`, this window expression filtered to `rk <= 1` must retain both `A` and `a`: ```sql RANK() OVER ( PARTITION BY grp ORDER BY CAST(s AS STRING COLLATE UTF8_LCASE) ) AS rk ``` The collated key is serialized as ordinary UTF8, and the new limiter compares its row-encoded bytes, so it drops `a`. `DENSE_RANK` has the same problem. The collated-scan fallback does not protect this case because the stored column is plain `STRING`. The cast stays in a native projection through the default codegen dispatcher, and the required two-key `[grp, cast-result]` sort passes `supportedSortType`, which only checks collations in its single-key branch. This particular `A, a, b` sequence is correctly sorted under both binary and case-insensitive ordering, so the lost row is not explained by an existing sort-order mismatch. The conversion path also permits this with `spark.comet.exec.window.enabled=false`, and a later Spark Window cannot restore a row already discarded by native Partial WGL. I verified that Spark 4.0.1 returns both peers and that the isolated native Rank and DenseRank operators each retain only `A`. The native probes used a diagnostic copy with only the timer-borrow compilation issue corrected, not an end-to-end run of the unchanged head. Please gate collated key types, including nested collated types, until peer equality preserves Spark's collation semantics. -- 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]
