From add9758e9f5a46a74c1469a6e9d954fcfcebb35c Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 21 Aug 2026 09:17:27 -0700 Subject: [PATCH 1/6] feat: support `WindowGroupLimit` --- dev/diffs/3.5.9.diff | 24 + dev/diffs/4.0.4.diff | 24 + dev/diffs/4.1.3.diff | 24 + native/core/src/execution/jni_api.rs | 1 + native/core/src/execution/operators/mod.rs | 2 + .../src/execution/operators/rank_limit.rs | 378 ++++++++++++ native/core/src/execution/planner.rs | 103 +++- .../execution/planner/operator_registry.rs | 1 + native/proto/src/proto/operator.proto | 17 + .../scala/org/apache/comet/CometConf.scala | 2 + .../apache/comet/rules/CometExecRule.scala | 8 +- .../sql/comet/CometWindowGroupLimitExec.scala | 153 +++++ .../shims/ShimCometWindowGroupLimit.scala | 40 ++ .../shims/ShimCometWindowGroupLimit.scala | 51 ++ .../shims/ShimCometWindowGroupLimit.scala | 52 ++ .../window/window_group_limit_datatypes.sql | 136 +++++ .../window/window_group_limit_edge.sql | 107 ++++ .../window/window_group_limit_rank.sql | 394 ++++++++++++ .../window_group_limit_rank_dense_rank.sql | 106 ++++ .../window/window_group_limit_row_number.sql | 168 ++++++ .../window_group_limit_scalar_subquery.sql | 70 +++ .../q44/extended.txt | 115 ++-- .../q67/extended.txt | 77 ++- .../q70/extended.txt | 99 ++- .../q44/extended.txt | 84 ++- .../q67a/extended.txt | 571 +++++++++--------- .../q70a/extended.txt | 287 +++++---- 27 files changed, 2462 insertions(+), 632 deletions(-) create mode 100644 native/core/src/execution/operators/rank_limit.rs create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala create mode 100644 spark/src/main/spark-3.4/org/apache/comet/shims/ShimCometWindowGroupLimit.scala create mode 100644 spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala create mode 100644 spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index 64c587039c9..6ff471ae4bc 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -1446,6 +1446,30 @@ index 005e764cc30..92ec088efab 100644 } private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala +index 46ed8fdfd21..4585dbab5b8 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution + + import org.apache.spark.sql.{DataFrame, QueryTest} ++import org.apache.spark.sql.comet.CometWindowGroupLimitExec + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.window.WindowGroupLimitExec + import org.apache.spark.sql.functions.lit +@@ -30,7 +31,10 @@ abstract class RemoveRedundantWindowGroupLimitsSuiteBase + + private def checkNumWindowGroupLimits(df: DataFrame, count: Int): Unit = { + val plan = df.queryExecution.executedPlan +- assert(collectWithSubqueries(plan) { case exec: WindowGroupLimitExec => exec }.length == count) ++ assert(collectWithSubqueries(plan) { ++ case exec: WindowGroupLimitExec => exec ++ case exec: CometWindowGroupLimitExec => exec ++ }.length == count) + } + + private def checkWindowGroupLimits(query: String, count: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala index 47679ed7865..9ffbaecb98e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index 4f2667fbc61..fe3bbf23e55 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -1862,6 +1862,30 @@ index 005e764cc30..92ec088efab 100644 } private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala +index 46ed8fdfd21..4585dbab5b8 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution + + import org.apache.spark.sql.{DataFrame, QueryTest} ++import org.apache.spark.sql.comet.CometWindowGroupLimitExec + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.window.WindowGroupLimitExec + import org.apache.spark.sql.functions.lit +@@ -30,7 +31,10 @@ abstract class RemoveRedundantWindowGroupLimitsSuiteBase + + private def checkNumWindowGroupLimits(df: DataFrame, count: Int): Unit = { + val plan = df.queryExecution.executedPlan +- assert(collectWithSubqueries(plan) { case exec: WindowGroupLimitExec => exec }.length == count) ++ assert(collectWithSubqueries(plan) { ++ case exec: WindowGroupLimitExec => exec ++ case exec: CometWindowGroupLimitExec => exec ++ }.length == count) + } + + private def checkWindowGroupLimits(query: String, count: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala index 47679ed7865..9ffbaecb98e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index 21eaa1266e5..ad34e8eaf53 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -1985,6 +1985,30 @@ index 005e764cc30..92ec088efab 100644 } private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala +index 46ed8fdfd21..4585dbab5b8 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution + + import org.apache.spark.sql.{DataFrame, QueryTest} ++import org.apache.spark.sql.comet.CometWindowGroupLimitExec + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.window.WindowGroupLimitExec + import org.apache.spark.sql.functions.lit +@@ -30,7 +31,10 @@ abstract class RemoveRedundantWindowGroupLimitsSuiteBase + + private def checkNumWindowGroupLimits(df: DataFrame, count: Int): Unit = { + val plan = df.queryExecution.executedPlan +- assert(collectWithSubqueries(plan) { case exec: WindowGroupLimitExec => exec }.length == count) ++ assert(collectWithSubqueries(plan) { ++ case exec: WindowGroupLimitExec => exec ++ case exec: CometWindowGroupLimitExec => exec ++ }.length == count) + } + + private def checkWindowGroupLimits(query: String, count: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala index 47679ed7865..9ffbaecb98e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d80754736b7..82bd6b3fc11 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -276,6 +276,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin", OpStruct::Sample(_) => "Sample", OpStruct::ContribScan(_) => "ContribScan", + OpStruct::WindowGroupLimit(_) => "WindowGroupLimit", } } diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index b9b2b0fbd73..81615643bc3 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -35,6 +35,8 @@ mod csv_scan; pub mod projection; mod sample; pub use sample::SampleExec; +mod rank_limit; +pub use rank_limit::{PartitionedRankLimitExec, WindowFnKind}; mod scan; mod shuffle_scan; pub use csv_scan::init_csv_datasource_exec; diff --git a/native/core/src/execution/operators/rank_limit.rs b/native/core/src/execution/operators/rank_limit.rs new file mode 100644 index 00000000000..8d8c0ddce8d --- /dev/null +++ b/native/core/src/execution/operators/rank_limit.rs @@ -0,0 +1,378 @@ +// 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::{DataFusionError, 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::{ + 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, + /// Full sort expression `[partition_keys..., order_keys...]`. + expr: LexOrdering, + /// Leading count of expressions in `expr` that form the partition key. + /// Zero means "no PARTITION BY" (global top-K within each input partition). + /// Can equal `expr.len()` when `LexOrdering::new` dedup collapses the ORDER BY suffix. + partition_prefix_len: usize, + fetch: usize, + kind: WindowFnKind, + cache: Arc, +} + +impl PartitionedRankLimitExec { + pub fn try_new( + input: Arc, + expr: LexOrdering, + partition_prefix_len: usize, + fetch: usize, + kind: WindowFnKind, + ) -> Result { + // Guard against `LexOrdering::new` dedup dropping a partition key. + if partition_prefix_len > expr.len() { + return Err(DataFusionError::Internal(format!( + "PartitionedRankLimitExec: partition prefix ({partition_prefix_len}) exceeds \ + ordering length ({})", + expr.len() + ))); + } + let cache = Arc::new(Self::compute_properties(&input, &expr)?); + Ok(Self { + input, + expr, + partition_prefix_len, + fetch, + kind, + cache, + }) + } + + fn compute_properties( + input: &Arc, + sort_exprs: &LexOrdering, + ) -> Result { + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.reorder(sort_exprs.clone())?; + Ok(PlanProperties::new( + eq_properties, + input.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } +} + +impl DisplayAs for PartitionedRankLimitExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "CometPartitionedRankLimitExec: kind={:?}, fetch={}, partition_prefix_len={}, \ + expr=[{}]", + self.kind, self.fetch, self.partition_prefix_len, self.expr + ), + DisplayFormatType::TreeRender => unimplemented!(), + } + } +} + +impl ExecutionPlan for PartitionedRankLimitExec { + fn name(&self) -> &str { + "CometPartitionedRankLimitExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&children[0]), + self.expr.clone(), + self.partition_prefix_len, + self.fetch, + self.kind, + )?)) + } + + // The operator's correctness depends on the input being sorted by + // `[partition_keys..., order_keys...]`. Declaring this lets DataFusion's + // `EnforceSorting` insert a `SortExec` if the sort somehow got dropped + // during plan construction, though in practice Spark's Catalyst already + // injects the sort above `WindowGroupLimitExec`. + fn required_input_ordering(&self) -> Vec> { + vec![Some(OrderingRequirements::from(self.expr.clone()))] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.input.execute(partition, context)?; + let schema = input.schema(); + + let partition_key = build_key_encoder(&self.expr[..self.partition_prefix_len], &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, or every ORDER BY + // column was deduplicated with a PARTITION BY column by `LexOrdering::new`), + // `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.expr[self.partition_prefix_len..], &schema)? + }; + + Ok(Box::pin(RankLimitStream { + input, + schema, + partition_key, + order_key, + limit: self.fetch as u64, + kind: self.kind, + prev_partition: None, + prev_order: None, + rank: 0, + count: 0, + })) + } +} + +/// 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>, +} + +impl KeyEncoder { + fn encode(&self, batch: &RecordBatch) -> Result { + let num_rows = batch.num_rows(); + let cols: Vec = self + .exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + Ok(self.converter.convert_columns(&cols)?) + } +} + +fn build_key_encoder(exprs: &[PhysicalSortExpr], schema: &SchemaRef) -> Result> { + 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> { + 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, + /// `None` when the ORDER BY suffix is empty (fully covered by PARTITION BY + /// or absent entirely), and always `None` for ROW_NUMBER (rank formula + /// never reads order keys). + order_key: Option, + limit: u64, + kind: WindowFnKind, + + // Per-partition streaming state, persisted across batches. + prev_partition: Option, + prev_order: Option, + /// 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, +} + +impl RankLimitStream { + fn process_batch(&mut self, batch: &RecordBatch) -> Result { + 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()?; + let order_rows = self + .order_key + .as_ref() + .map(|k| k.encode(batch)) + .transpose()?; + + 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; + } + + // 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 suffix — `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; + } + + 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; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match self.input.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => match self.process_batch(&batch) { + // Skip fully-filtered batches so downstream never sees + // spurious empty batches between real ones. + Ok(out) if out.num_rows() == 0 => continue, + Ok(out) => return Poll::Ready(Some(Ok(out))), + Err(e) => return Poll::Ready(Some(Err(e))), + }, + Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), + Poll::Ready(None) => return Poll::Ready(None), + Poll::Pending => return Poll::Pending, + } + } + } +} + +impl RecordBatchStream for RankLimitStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d109627e825..8f9a193ab09 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -33,6 +33,7 @@ mod delta_scan; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::AlignedArrowStreamReader; use crate::execution::operators::IcebergScanExec; +use crate::execution::operators::{PartitionedRankLimitExec, WindowFnKind}; use crate::execution::{ expressions::list_empty_to_null::ListEmptyToNullExpr, expressions::list_positions::ListPositionsExpr, @@ -135,7 +136,8 @@ use datafusion_comet_proto::{ spark_operator::{ self, lower_window_frame_bound::LowerFrameBoundStruct, operator::OpStruct, upper_window_frame_bound::UpperFrameBoundStruct, AggregateMode as ProtoAggregateMode, - BuildSide, CompressionCodec as SparkCompressionCodec, JoinType, Operator, WindowFrameType, + BuildSide, CompressionCodec as SparkCompressionCodec, JoinType, Operator, RankLikeFunction, + WindowFrameType, }, spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, }; @@ -2377,6 +2379,105 @@ impl PhysicalPlanner { Arc::new(SparkPlan::new(spark_plan.plan_id, final_plan, vec![child])), )) } + OpStruct::WindowGroupLimit(wgl) => { + // Route: + // * ROW_NUMBER + empty PARTITION BY -> `LocalLimitExec`. Spark's WGL requires + // the child to be sorted by ORDER BY, so `first K rows per input partition` + // IS the global RANK top-K (Partial phase); the Final phase runs on the + // single-partition post-shuffle stream and stays correct. + // * Everything else (ROW_NUMBER partitioned, RANK / DENSE_RANK with any + // PARTITION BY) -> Comet's streaming `PartitionedRankLimitExec`. Relies on + // the Spark-injected sort of `[partition_keys, order_keys]`, which lets it + // preserve input order for rows tied on the ORDER BY keys (matching Spark's + // `SimpleLimitIterator` for ROW_NUMBER and `RankLimitIterator` for RANK / + // DENSE_RANK). `partition_prefix_len == 0` degenerates to one big partition + // per DF input partition. + assert_eq!(children.len(), 1); + let (scans, shuffle_scans, child) = + self.create_plan(&children[0], inputs, partition_count)?; + let input_schema = child.schema(); + + let kind = + match RankLikeFunction::try_from(wgl.rank_like_function).map_err(|_| { + GeneralError(format!( + "Unsupported WindowGroupLimit rank function tag: {}", + wgl.rank_like_function + )) + })? { + RankLikeFunction::RowNumber => WindowFnKind::RowNumber, + RankLikeFunction::Rank => WindowFnKind::Rank, + RankLikeFunction::DenseRank => WindowFnKind::DenseRank, + }; + + let partition_prefix_len = wgl.partition_by_list.len(); + let fetch = wgl.limit as usize; + + // Fast path: ROW_NUMBER without PARTITION BY collapses to `first K rows per + // input DF partition`, which is exactly what `LocalLimitExec` does. No + // row-encoding overhead, no ORDER BY tie detection needed. + let topk: Arc = if partition_prefix_len == 0 + && kind == WindowFnKind::RowNumber + { + Arc::new(LocalLimitExec::new(Arc::clone(&child.native_plan), fetch)) + } else { + // Partition keys arrive as bare exprs (no direction). Spark's WGL + // requires the child to be sorted by partition-keys ASCENDING then by + // order-by keys, so materialize partition keys with SortOptions matching + // Spark's `SortOrder(_, Ascending)` (ascending, nulls_first) and + // concatenate with the ORDER BY exprs into a single LexOrdering + // `[partition_keys, order_keys]`. + let mut sort_exprs: Vec = + Vec::with_capacity(partition_prefix_len + wgl.order_by_list.len()); + for expr in &wgl.partition_by_list { + let phys = self.create_expr(expr, Arc::clone(&input_schema))?; + sort_exprs.push(PhysicalSortExpr { + expr: phys, + options: SortOptions::default(), + }); + } + for expr in &wgl.order_by_list { + sort_exprs.push(self.create_sort_expr(expr, Arc::clone(&input_schema))?); + } + + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + GeneralError("WindowGroupLimit produced empty LexOrdering".to_string()) + })?; + + // `LexOrdering::new` deduplicates by underlying `PhysicalExpr`, so if a + // PARTITION BY column also appears in ORDER BY (e.g. `PARTITION BY t2c + // ORDER BY t2c`, produced by SPARK-46526-shaped scalar-subquery + // rewrites) the ORDER BY suffix collapses away. The streaming operator + // handles that tie-everything degenerate case (every row shares the + // K-th ORDER BY value) the same way it handles any other tie: rank + // stays at 0 for RANK/DENSE_RANK, or increments per row for ROW_NUMBER. + + // Always route to the streaming `PartitionedRankLimitExec`. Spark's + // `WindowGroupLimitExec.requiredChildOrdering` guarantees the child is + // sorted by `[partition_keys..., order_keys...]`, so a single pass + // suffices. `PartitionedTopKExec` (heap-based) is faster asymptotically + // but reorders tied rows, which breaks Spark's `SimpleLimitIterator` / + // `RankLimitIterator` semantics: for `ROW_NUMBER()`, Spark assigns ranks + // in input order to rows tied on the ORDER BY keys, and any deviation + // shows up as a wrong-row diff in SPARK-37099-shaped tests. + Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + kind, + )?) + }; + + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new( + spark_plan.plan_id, + topk, + vec![Arc::clone(&child)], + )), + )) + } OpStruct::ShuffleScan(scan) => { let data_types = scan.fields.iter().map(to_arrow_datatype).collect_vec(); diff --git a/native/core/src/execution/planner/operator_registry.rs b/native/core/src/execution/planner/operator_registry.rs index 7e67dc181ca..c644741894c 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -159,5 +159,6 @@ fn get_operator_type(spark_operator: &Operator) -> Option { // match. No contrib-specific logic lives here -- we just signal "no OperatorType mapping" // so the supports-mixed-codegen check skips it. OpStruct::ContribScan(_) => None, + OpStruct::WindowGroupLimit(_) => None, // Not yet in OperatorType enum } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ec9448709cf..e96b545ef62 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -66,6 +66,7 @@ message Operator { ShuffleScan shuffle_scan = 116; BroadcastNestedLoopJoin broadcast_nested_loop_join = 117; Sample sample = 118; + WindowGroupLimit window_group_limit = 119; // Extension point for optional, out-of-tree contrib scans (Delta, Lance, ...). The concrete // scan message (e.g. `DeltaScan`) is packed into this envelope on the JVM side and dispatched // by `type_url` on the native side. Using a single permanent field -- rather than a new oneof @@ -860,3 +861,19 @@ message Window { repeated spark.spark_expression.Expr partition_by_list = 3; Operator child = 4; } + +// Top-K rows per partition group. Corresponds to Spark's WindowGroupLimitExec +// (Spark 3.5+). Rank-like function drives the semantics; the operator only +// accepts ROW_NUMBER, RANK, and DENSE_RANK. +message WindowGroupLimit { + repeated spark.spark_expression.Expr partition_by_list = 1; + repeated spark.spark_expression.Expr order_by_list = 2; + int32 limit = 3; + RankLikeFunction rank_like_function = 4; +} + +enum RankLikeFunction { + RowNumber = 0; + Rank = 1; + DenseRank = 2; +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 475b73ed692..d8fe5b69890 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -249,6 +249,8 @@ object CometConf extends ShimCometConf { createExecEnabledConfig("explode", defaultValue = true) val COMET_EXEC_WINDOW_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("window", defaultValue = true) + val COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED: ConfigEntry[Boolean] = + createExecEnabledConfig("windowGroupLimit", defaultValue = true) val COMET_EXEC_TAKE_ORDERED_AND_PROJECT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("takeOrderedAndProject", defaultValue = true) val COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED: ConfigEntry[Boolean] = diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 3c2668cfe92..c69602fc80b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -55,7 +55,7 @@ import org.apache.comet.CometSparkSessionExtensions._ import org.apache.comet.rules.CometExecRule.allExecs import org.apache.comet.serde._ import org.apache.comet.serde.operator._ -import org.apache.comet.shims.{ShimCometStreaming, ShimSubqueryBroadcast} +import org.apache.comet.shims.{ShimCometStreaming, ShimCometWindowGroupLimit, ShimSubqueryBroadcast} object CometExecRule { @@ -71,7 +71,7 @@ object CometExecRule { * Fully native operators. */ val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = - Map( + Map[Class[_ <: SparkPlan], CometOperatorSerde[_]]( classOf[ProjectExec] -> CometProjectExec, classOf[FilterExec] -> CometFilterExec, classOf[LocalLimitExec] -> CometLocalLimitExec, @@ -87,7 +87,9 @@ object CometExecRule { classOf[SortExec] -> CometSortExec, classOf[LocalTableScanExec] -> CometLocalTableScanExec, classOf[SampleExec] -> CometSampleExec, - classOf[WindowExec] -> CometWindowExec) + classOf[WindowExec] -> CometWindowExec) ++ + // WindowGroupLimitExec exists only on Spark 3.5+; the shim returns None on 3.4. + ShimCometWindowGroupLimit.windowGroupLimitClass.map(_ -> CometWindowGroupLimitExec) /** * Sinks that have a native plan of ScanExec. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala new file mode 100644 index 00000000000..d64d5950b25 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -0,0 +1,153 @@ +/* + * 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+). */ + case class Fields( + partitionSpec: Seq[Expression], + orderSpec: Seq[SortOrder], + rankLikeFunction: RankLikeFunction, + limit: Int) + + 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] = { + // `nativeExecs` only routes here for a real `WindowGroupLimitExec` on Spark 3.5+, so the + // shim always returns Some. + val fields = ShimCometWindowGroupLimit.extract(op).get + + 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)) + + val failing = (partitionProtos ++ orderProtos).collect { case (e, None) => e } + if (failing.nonEmpty) { + withFallbackReason( + op, + failing.map(_.sql).mkString("WindowGroupLimit: unsupported expressions: ", ", ", "")) + return None + } + + val wglBuilder = OperatorOuterClass.WindowGroupLimit + .newBuilder() + .setLimit(fields.limit) + .setRankLikeFunction(fields.rankLikeFunction) + wglBuilder.addAllPartitionByList(partitionProtos.map(_._2.get).asJava) + wglBuilder.addAllOrderByList(orderProtos.map(_._2.get).asJava) + Some(builder.setWindowGroupLimit(wglBuilder).build()) + } + + override def createExec(nativeOp: Operator, op: SparkPlan): CometNativeExec = { + val fields = ShimCometWindowGroupLimit + .extract(op) + .getOrElse( + throw new IllegalStateException( + "createExec called on a non-WindowGroupLimitExec operator: " + op.nodeName)) + CometWindowGroupLimitExec( + nativeOp, + op, + op.output, + fields.partitionSpec, + fields.orderSpec, + fields.limit, + op.children.head, + SerializedPlan(None)) + } +} + +/** + * Comet physical plan node for Spark `WindowGroupLimitExec`. The Spark Partial/Final split is + * preserved unchanged in the Spark plan tree (each side is planned as its own native subtree), so + * the case class doesn't carry a mode field. + */ +case class CometWindowGroupLimitExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + partitionSpec: Seq[Expression], + orderSpec: Seq[SortOrder], + limit: Int, + child: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometUnaryExec { + + override def nodeName: String = "CometWindowGroupLimitExec" + + override def outputOrdering: Seq[SortOrder] = child.outputOrdering + + override def outputPartitioning: Partitioning = child.outputPartitioning + + protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + this.copy(child = newChild) + + override def stringArgs: Iterator[Any] = + Iterator(output, partitionSpec, orderSpec, limit, child) + + override def equals(obj: Any): Boolean = obj match { + case other: CometWindowGroupLimitExec => + this.output == other.output && + this.partitionSpec == other.partitionSpec && + this.orderSpec == other.orderSpec && + this.limit == other.limit && + this.child == other.child && + this.serializedPlanOpt == other.serializedPlanOpt + case _ => false + } + + override def hashCode(): Int = + Objects.hashCode(output, partitionSpec, orderSpec, Integer.valueOf(limit), child) +} diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimCometWindowGroupLimit.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimCometWindowGroupLimit.scala new file mode 100644 index 00000000000..b338a32b05f --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -0,0 +1,40 @@ +/* + * 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.comet.shims + +import org.apache.spark.sql.comet.CometWindowGroupLimitExec.Fields +import org.apache.spark.sql.execution.SparkPlan + +/** + * Spark 3.4 does not have `WindowGroupLimitExec` (introduced by SPARK-37099 in Spark 3.5). This + * no-op shim keeps the shared Comet code compiling against Spark 3.4 while ensuring the operator + * is never registered for conversion on this profile. + */ +object ShimCometWindowGroupLimit { + + /** + * The `WindowGroupLimitExec` class object on Spark 3.5+, or `None` on Spark 3.4. + * `CometExecRule` uses this to conditionally register the serde in `nativeExecs`. + */ + def windowGroupLimitClass: Option[Class[_ <: SparkPlan]] = None + + /** Extract WGL fields when `op` is a `WindowGroupLimitExec` (Spark 3.5+). */ + def extract(op: SparkPlan): Option[Fields] = None +} diff --git a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala new file mode 100644 index 00000000000..4f17f4ec98e --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -0,0 +1,51 @@ +/* + * 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.comet.shims + +import org.apache.spark.sql.catalyst.expressions.{DenseRank, Rank, RowNumber} +import org.apache.spark.sql.comet.CometWindowGroupLimitExec.Fields +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.window.WindowGroupLimitExec + +import org.apache.comet.serde.OperatorOuterClass.RankLikeFunction + +/** + * Spark 3.5+ shim exposing `WindowGroupLimitExec` (SPARK-37099) to the shared Comet code without + * causing the 3.4 build to fail on a missing class reference. + */ +object ShimCometWindowGroupLimit { + + def windowGroupLimitClass: Option[Class[_ <: SparkPlan]] = Some(classOf[WindowGroupLimitExec]) + + def extract(op: SparkPlan): Option[Fields] = op match { + case w: WindowGroupLimitExec => + val fn = w.rankLikeFunction match { + case _: RowNumber => RankLikeFunction.RowNumber + case _: Rank => RankLikeFunction.Rank + case _: DenseRank => RankLikeFunction.DenseRank + case other => + throw new IllegalStateException( + s"Unexpected rank-like function in WindowGroupLimitExec: ${other.getClass.getName}") + } + Some(Fields(w.partitionSpec, w.orderSpec, fn, w.limit)) + case _ => + None + } +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala new file mode 100644 index 00000000000..3175d30f39f --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -0,0 +1,52 @@ +/* + * 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.comet.shims + +import org.apache.spark.sql.catalyst.expressions.{DenseRank, Rank, RowNumber} +import org.apache.spark.sql.comet.CometWindowGroupLimitExec.Fields +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.window.WindowGroupLimitExec + +import org.apache.comet.serde.OperatorOuterClass.RankLikeFunction + +/** + * Spark 4.x shim exposing `WindowGroupLimitExec` to the shared Comet code. Content mirrors the + * Spark 3.5 shim; the split exists because the per-Spark-minor `spark-3.5/` and `spark-4.x/` + * source dirs are activated by disjoint Maven profiles. + */ +object ShimCometWindowGroupLimit { + + def windowGroupLimitClass: Option[Class[_ <: SparkPlan]] = Some(classOf[WindowGroupLimitExec]) + + def extract(op: SparkPlan): Option[Fields] = op match { + case w: WindowGroupLimitExec => + val fn = w.rankLikeFunction match { + case _: RowNumber => RankLikeFunction.RowNumber + case _: Rank => RankLikeFunction.Rank + case _: DenseRank => RankLikeFunction.DenseRank + case other => + throw new IllegalStateException( + s"Unexpected rank-like function in WindowGroupLimitExec: ${other.getClass.getName}") + } + Some(Fields(w.partitionSpec, w.orderSpec, fn, w.limit)) + case _ => + None + } +} diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql new file mode 100644 index 00000000000..75ec86ce187 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql @@ -0,0 +1,136 @@ +-- 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. + +-- WindowGroupLimit data-type coverage for ORDER BY / PARTITION BY keys. +-- Exercises int, bigint (past Int32), double (NaN, +/-Inf, +/-0.0), decimal, date, +-- timestamp, string (incl. empty and NULL). Also covers decimal and date partition keys, +-- which upstream Spark's window.sql does not exercise. All rank-like functions +-- (ROW_NUMBER, RANK, DENSE_RANK) route to Comet's streaming `PartitionedRankLimitExec` +-- and run natively. +-- +-- Each ROW_NUMBER's ORDER BY carries `payload ASC` as a secondary key because `payload` is +-- unique per row; that removes tie-breaking non-determinism between engines even where +-- the streaming operator's input-order guarantee doesn't apply (e.g. plain ordering +-- comparisons in the outer query). + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 + +statement +CREATE TABLE test_wgl_types( + k_str string, + k_int int, + k_long bigint, + k_dbl double, + k_dec decimal(18,4), + k_date date, + k_ts timestamp, + payload int +) USING parquet + +statement +INSERT INTO test_wgl_types VALUES + ('a', 1, 2147483650, 1.5, cast('1.2345' as decimal(18,4)), date '2020-01-01', timestamp '2020-01-01 00:00:00', 1), + ('a', 2, 2147483651, double('NaN'), cast('99999999999999.9999' as decimal(18,4)), date '2020-01-02', timestamp '2020-01-01 00:00:01', 2), + ('a', -2147483648, -9223372036854775808, double('Infinity'),cast('-99999999999999.9999' as decimal(18,4)), date '1970-01-01', timestamp '1970-01-01 00:00:00', 3), + ('a', 2147483647, 9223372036854775807, double('-Infinity'),cast('0.0000' as decimal(18,4)), date '9999-12-31', timestamp '9999-12-31 23:59:59', 4), + ('a', 0, 0, 0.0, cast('0.0001' as decimal(18,4)), date '2020-06-15', timestamp '2020-06-15 12:00:00', 5), + ('a', 0, 0, -0.0, cast('-0.0001' as decimal(18,4)), date '2020-06-15', timestamp '2020-06-15 12:00:00', 6), + ('b', 10, 10, double('NaN'), cast('1.0' as decimal(18,4)), date '2020-01-01', timestamp '2020-01-01 00:00:00', 7), + ('b', NULL, NULL, NULL, NULL, NULL, NULL, 8), + ('', 42, 42, 42.0, cast('42' as decimal(18,4)), date '2000-02-29', timestamp '2000-02-29 03:14:15', 9), + (NULL, 1, 1, 1.0, cast('1' as decimal(18,4)), date '2000-01-01', timestamp '2000-01-01 00:00:00', 10) + +-- Order by int, top-1 per string partition. +query +SELECT k_str, k_int, payload FROM ( + SELECT k_str, k_int, payload, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_int DESC NULLS LAST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn = 1 ORDER BY k_str NULLS LAST + +-- Order by bigint past Int32 range. +query +SELECT k_str, k_long, payload FROM ( + SELECT k_str, k_long, payload, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_long ASC NULLS LAST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn <= 2 ORDER BY k_str NULLS LAST, rn + +-- Order by double (NaN, +/-Inf, +/-0.0). Spark treats NaN as max, +0.0 == -0.0. +query +SELECT k_str, k_dbl, payload FROM ( + SELECT k_str, k_dbl, payload, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_dbl DESC NULLS LAST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn <= 3 ORDER BY k_str NULLS LAST, rn + +-- Order by decimal. +query +SELECT k_str, k_dec, payload FROM ( + SELECT k_str, k_dec, payload, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_dec DESC NULLS LAST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn <= 2 ORDER BY k_str NULLS LAST, rn + +-- Order by date. +query +SELECT k_str, k_date, payload FROM ( + SELECT k_str, k_date, payload, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_date ASC NULLS FIRST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn <= 2 ORDER BY k_str NULLS LAST, rn + +-- Order by timestamp. +query +SELECT k_str, k_ts, payload FROM ( + SELECT k_str, k_ts, payload, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_ts DESC NULLS LAST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn <= 2 ORDER BY k_str NULLS LAST, rn + +-- Order by string (incl. empty string and NULL); partition by a derived int. +query +SELECT k_str, payload FROM ( + SELECT k_str, payload, + ROW_NUMBER() OVER (PARTITION BY payload % 3 ORDER BY k_str ASC NULLS FIRST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn <= 2 ORDER BY payload % 3, rn + +-- Decimal partition key with RANK. +query +SELECT k_dec, k_int FROM ( + SELECT k_dec, k_int, + RANK() OVER (PARTITION BY k_dec ORDER BY k_int DESC NULLS LAST) AS rk + FROM test_wgl_types +) t WHERE rk = 1 ORDER BY k_dec NULLS LAST, k_int DESC NULLS LAST + +-- Date partition key with DENSE_RANK. +query +SELECT k_date, k_int FROM ( + SELECT k_date, k_int, + DENSE_RANK() OVER (PARTITION BY k_date ORDER BY k_int ASC NULLS LAST) AS dr + FROM test_wgl_types +) t WHERE dr <= 2 ORDER BY k_date NULLS LAST, dr, k_int NULLS LAST + +-- Partition column IS the order column. +query +SELECT k_str, k_int FROM ( + SELECT k_str, k_int, + ROW_NUMBER() OVER (PARTITION BY k_str ORDER BY k_str, k_int ASC NULLS LAST, payload ASC) AS rn + FROM test_wgl_types +) t WHERE rn = 1 ORDER BY k_str NULLS LAST diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql new file mode 100644 index 00000000000..80df37b6c0c --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql @@ -0,0 +1,107 @@ +-- 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. + +-- WindowGroupLimit boundary and negative cases: empty table, single-row table, all-null +-- order column, all-tied rows, and filter shapes Spark's InferWindowGroupLimit rule does +-- NOT push (rn > k, disjunction with a non-rank predicate, SizeBasedWindowFunction sibling +-- like percent_rank per SPARK-46941). The negative cases must still return correct rows. + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 + +statement +CREATE TABLE test_wgl_empty(a int, b int) USING parquet + +statement +CREATE TABLE test_wgl_one(a int, b int) USING parquet + +statement +INSERT INTO test_wgl_one VALUES (1, 100) + +statement +CREATE TABLE test_wgl_all_null(a int, b int) USING parquet + +statement +INSERT INTO test_wgl_all_null VALUES (1, NULL), (1, NULL), (1, NULL) + +statement +CREATE TABLE test_wgl_all_tie(a int, b int) USING parquet + +statement +INSERT INTO test_wgl_all_tie VALUES (1, 5), (1, 5), (1, 5), (1, 5) + +-- Empty table. +query +SELECT a, b FROM ( + SELECT a, b, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b DESC) AS rn FROM test_wgl_empty +) t WHERE rn <= 3 + +-- Single-row table. +query +SELECT a, b, rn FROM ( + SELECT a, b, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b DESC) AS rn FROM test_wgl_one +) t WHERE rn <= 3 + +-- All-null order column with RANK: every row shares rank 1. +query +SELECT a, b FROM ( + SELECT a, b, RANK() OVER (PARTITION BY a ORDER BY b DESC NULLS LAST) AS rk + FROM test_wgl_all_null +) t WHERE rk = 1 + +-- All-null order column with DENSE_RANK: every row shares rank 1. +query +SELECT a, b FROM ( + SELECT a, b, DENSE_RANK() OVER (PARTITION BY a ORDER BY b DESC NULLS LAST) AS dr + FROM test_wgl_all_null +) t WHERE dr = 1 + +-- Every row tied — row_number=1 returns 1 row. +query +SELECT a, b FROM ( + SELECT a, b, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) AS rn + FROM test_wgl_all_tie +) t WHERE rn = 1 + +-- Every row tied — rank=1 returns all rows. +query +SELECT count(*) AS cnt FROM ( + SELECT a, b, RANK() OVER (PARTITION BY a ORDER BY b) AS rk + FROM test_wgl_all_tie +) t WHERE rk = 1 + +-- Every row tied — dense_rank=1 returns all rows. +query +SELECT count(*) AS cnt FROM ( + SELECT a, b, DENSE_RANK() OVER (PARTITION BY a ORDER BY b) AS dr + FROM test_wgl_all_tie +) t WHERE dr = 1 + +-- Unsupported filter form `rn > k` — InferWindowGroupLimit does NOT push, so no fallback +-- attributable to WindowGroupLimit. Result must still be correct. +query +SELECT a, b FROM ( + SELECT a, b, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b DESC) AS rn + FROM test_wgl_all_tie +) t WHERE rn > 2 ORDER BY rn + +-- Unsupported filter form: disjunction with a non-rank predicate — no WGL pushdown. +query +SELECT a, b FROM ( + SELECT a, b, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b DESC) AS rn + FROM test_wgl_one +) t WHERE rn = 1 OR b > 50 diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql new file mode 100644 index 00000000000..8530874138b --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql @@ -0,0 +1,394 @@ +-- 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. + +-- WindowGroupLimit RANK() native pushdown tests. Comet routes RANK with a non-empty +-- PARTITION BY through the streaming `PartitionedRankLimitExec`, which relies on the +-- Spark-injected `[partition_keys, order_keys]` sort and retains all rows tied at the +-- K-th ORDER BY value -- matching Spark's `RankLimitIterator`. +-- +-- Every case here uses `query` mode: results are asserted against Spark AND every +-- pushdown-eligible plan must be Comet-native. ConfigMatrix threshold=-1 exercises +-- the plain Window+Filter shape (no WGL); threshold=1000 exercises the WGL pushdown +-- through the new native operator. Both must pass. +-- +-- Tie-collapsing note: since these tests project only the partition + order columns +-- (no per-row disambiguator), any row that shares (part, rk, ord) with another is +-- indistinguishable in the output, so the ORDER BY does not need a further tiebreaker +-- for `checkAnswer` to compare row-wise. + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 +-- ConfigMatrix: parquet.enable.dictionary=false,true + +statement +CREATE TABLE test_rank(part string, ord int, extra int) USING parquet + +statement +INSERT INTO test_rank VALUES + ('a', 100, 1), + ('a', 100, 2), -- tied with row 1 at rank 1 + ('a', 90, 3), -- rank 3 under RANK (gap after tie) + ('a', 80, 4), + ('a', 80, 5), -- tied at rank 4 under RANK + ('a', 70, 6), + ('b', 50, 7), + ('b', 40, 8), + ('b', 40, 9), -- tied at rank 2 under RANK + ('b', NULL, 10), + ('c', 1, 11), + ('d', NULL, 12), + ('d', NULL, 13) -- all NULLs in one partition + + +-- ================================================================================ +-- Basic tie semantics: RANK <= 2 must retain rows tied at rank 1, then advance to +-- the next distinct ORDER BY value. Under RANK's gap semantics, `rk <= 2` really +-- means "rank 1 rows only" when two rows tied at rank 1 push rank 2 out. +-- ================================================================================ +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 2 ORDER BY part, rk, ord DESC NULLS LAST + +-- Same shape with `< k` -- Spark rewrites to `<= k-1`. +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk < 3 ORDER BY part, rk, ord DESC NULLS LAST + +-- Literal on the left of the comparison (`k >= rk`). +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE 3 >= rk ORDER BY part, rk, ord DESC NULLS LAST + +-- Equality: rk = 1 returns every row tied at the top of its partition, so partitions +-- with duplicate top values return multiple rows (partition 'a' has two 100s). +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk = 1 ORDER BY part, ord DESC NULLS LAST + + +-- ================================================================================ +-- K variations +-- ================================================================================ + +-- Larger K covering every partition. +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 100 ORDER BY part, rk, ord DESC NULLS LAST + +-- Outer LIMIT after the WGL pushdown. Comet keeps the pushdown and stacks the +-- outer LIMIT above. +query +SELECT part, ord, rk FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 2 ORDER BY part, rk, ord DESC NULLS LAST LIMIT 5 + +-- Extra AND-predicate on a non-rank column. The predicate does not affect WGL +-- pushdown (Spark keeps rk<=k separate) but the answer must still be correct. +-- Project `extra` here so it survives the outer select and disambiguates output rows. +query +SELECT part, ord, extra FROM ( + SELECT part, ord, extra, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 3 AND extra >= 3 ORDER BY part, rk, ord DESC NULLS LAST, extra + + +-- ================================================================================ +-- ORDER BY variations (ASC/DESC, NULLS FIRST/LAST, multi-column) +-- ================================================================================ + +-- ASC ordering with NULLS FIRST. +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord ASC NULLS FIRST) AS rk + FROM test_rank +) t WHERE rk <= 2 ORDER BY part, rk, ord ASC NULLS FIRST + +-- ASC ordering with NULLS LAST. +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (PARTITION BY part ORDER BY ord ASC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 2 ORDER BY part, rk, ord ASC NULLS LAST + +-- Two-column ORDER BY. The secondary key eliminates every intra-rank tie so +-- exactly the top-K distinct rows are returned per partition. +query +SELECT part, ord, extra FROM ( + SELECT part, ord, extra, + RANK() OVER (PARTITION BY part ORDER BY ord DESC NULLS LAST, extra ASC NULLS FIRST) AS rk + FROM test_rank +) t WHERE rk <= 2 ORDER BY part, rk, ord DESC NULLS LAST, extra + + +-- ================================================================================ +-- Multi-column PARTITION BY +-- ================================================================================ + +statement +CREATE TABLE test_rank_multipart(a string, b int, score int) USING parquet + +statement +INSERT INTO test_rank_multipart VALUES + ('x', 1, 10), ('x', 1, 10), ('x', 1, 9), -- (x,1): tie at top, then rank 3 + ('x', 2, 5), ('x', 2, 5), + ('y', 1, 100), ('y', 1, 100), ('y', 1, 100), -- (y,1): all tied + (NULL, 1, 7), (NULL, 1, 7) -- NULL in partition column + +query +SELECT a, b, score FROM ( + SELECT a, b, score, + RANK() OVER (PARTITION BY a, b ORDER BY score DESC NULLS LAST) AS rk + FROM test_rank_multipart +) t WHERE rk <= 2 ORDER BY a NULLS LAST, b, rk, score DESC NULLS LAST + + +-- ================================================================================ +-- FP edge cases in the ORDER BY column (NaN / +Inf / -Inf / 0 / NULL). NaN sorts +-- greater than any finite value in Spark, and Arrow's row-format total-ordering +-- encoding matches -- so NaN gets rank 1 under DESC. Note: -0.0 vs 0.0 is a known +-- Spark-vs-Arrow-row divergence (Spark treats them equal in ORDER BY; Arrow's +-- bitwise total_cmp treats them distinct), so this test avoids exercising that +-- boundary directly by keeping the cutoff (rk <= 3) above the 0-values. +-- ================================================================================ + +statement +CREATE TABLE test_rank_fp(part string, v double) USING parquet + +statement +INSERT INTO test_rank_fp VALUES + ('p', double('NaN')), + ('p', double('Infinity')), + ('p', double('Infinity')), -- tied with prior +Inf under DESC + ('p', 1.0), + ('p', 0.0), + ('p', double('-Infinity')), + ('p', NULL) + +-- DESC NULLS LAST -> NaN, +Inf, +Inf, 1.0, 0.0, -Inf, NULL. +-- RANK <= 3 keeps the NaN row and both +Inf rows (rank 2 tied, rank 3 empty). +query +SELECT part, v FROM ( + SELECT part, v, + RANK() OVER (PARTITION BY part ORDER BY v DESC NULLS LAST) AS rk + FROM test_rank_fp +) t WHERE rk <= 3 ORDER BY rk, v DESC NULLS LAST + + +-- ================================================================================ +-- Decimal / bigint / integer boundary values in ORDER BY column +-- ================================================================================ + +statement +CREATE TABLE test_rank_num(part string, i32 int, i64 bigint, dec38 decimal(38, 4)) +USING parquet + +statement +INSERT INTO test_rank_num VALUES + ('p', -2147483648, -9223372036854775808, cast('-9999999999999999999999999999999999.9999' AS decimal(38,4))), + ('p', -2147483648, -9223372036854775808, cast('-9999999999999999999999999999999999.9999' AS decimal(38,4))), + ('p', 2147483647, 9223372036854775807, cast(' 9999999999999999999999999999999999.9999' AS decimal(38,4))), + ('p', 0, 0, cast('0.0000' AS decimal(38,4))), + ('p', NULL, NULL, NULL) + +-- ORDER BY the max-boundary decimal column, then rank <= 2. Verifies the row-encoder +-- handles wide decimals across partition boundaries without truncation. +query +SELECT part, dec38 FROM ( + SELECT part, dec38, + RANK() OVER (PARTITION BY part ORDER BY dec38 DESC NULLS LAST) AS rk + FROM test_rank_num +) t WHERE rk <= 2 ORDER BY rk, dec38 DESC NULLS LAST + +-- ORDER BY the int32/int64 extremes. RANK <= 2 with duplicate min values. +query +SELECT part, i32, i64 FROM ( + SELECT part, i32, i64, + RANK() OVER (PARTITION BY part ORDER BY i64 ASC NULLS LAST) AS rk + FROM test_rank_num +) t WHERE rk <= 2 ORDER BY rk, i64 ASC NULLS LAST, i32 + + +-- ================================================================================ +-- TPC-DS q67-shaped: PARTITION BY category ORDER BY sum DESC, rank <= 3 +-- Distilled from tpc-ds/q67.sql (`rank() over (partition by i_category order by +-- sumsales desc) rk ... where rk <= 100`). +-- ================================================================================ + +statement +CREATE TABLE q67_sales(category string, item string, revenue double) USING parquet + +statement +INSERT INTO q67_sales VALUES + ('food', 'apple', 500.0), + ('food', 'banana', 500.0), -- tie with apple at rank 1 + ('food', 'cherry', 300.0), + ('food', 'donut', 200.0), + ('food', 'eclair', 100.0), + ('food', 'fig', 50.0), + ('elec', 'phone', 1000.0), + ('elec', 'laptop', 900.0), + ('elec', 'tablet', 900.0), -- tie at rank 2 + ('elec', 'monitor', 500.0), + ('elec', 'cable', 10.0), + ('elec', 'mouse', 10.0), + ('elec', 'kbd', 10.0), + ('home', 'sofa', 750.0), + ('home', 'lamp', NULL), -- NULL revenue + (NULL, 'unknown', 42.0) -- NULL category + +query +SELECT category, item, revenue FROM ( + SELECT category, item, revenue, + RANK() OVER (PARTITION BY category ORDER BY revenue DESC NULLS LAST) AS rk + FROM q67_sales +) t WHERE rk <= 3 ORDER BY category NULLS LAST, rk, revenue DESC NULLS LAST, item + + +-- ================================================================================ +-- TPC-DS q70-shaped: PARTITION BY state, rank <= 5, filtering to top-K states +-- Distilled from tpc-ds/q70.sql (`rank() over (partition by s_state order by +-- sum(ss_net_profit) desc) as ranking ... where ranking <= 5`). +-- ================================================================================ + +statement +CREATE TABLE q70_state_sales(state string, county string, net_profit int) USING parquet + +statement +INSERT INTO q70_state_sales VALUES + ('CA', 'LA', 100), ('CA', 'SF', 90), ('CA', 'SD', 80), + ('CA', 'SJ', 70), ('CA', 'OA', 60), ('CA', 'FR', 50), + ('TX', 'HTX', 200), ('TX', 'DAL', 150), ('TX', 'AUS', 150), -- tie at rank 2 + ('TX', 'SAT', 50), + ('NY', 'NYC', 300), ('NY', 'BUF', 100), + ('WA', 'SEA', 80), ('WA', 'TAC', 80), ('WA', 'SPO', 80), -- all tied + ('OR', NULL, 10), -- NULL county + (NULL, 'X', 42) + +-- Match q70's structure: rank each state's counties by profit, keep top-5. +query +SELECT state, county, net_profit FROM ( + SELECT state, county, net_profit, + RANK() OVER (PARTITION BY state ORDER BY net_profit DESC NULLS LAST) AS ranking + FROM q70_state_sales +) t WHERE ranking <= 5 +ORDER BY state NULLS LAST, ranking, net_profit DESC NULLS LAST, county NULLS LAST + +-- Same query with the outer LIMIT clause q70 uses. +query +SELECT state, county, net_profit FROM ( + SELECT state, county, net_profit, + RANK() OVER (PARTITION BY state ORDER BY net_profit DESC NULLS LAST) AS ranking + FROM q70_state_sales +) t WHERE ranking <= 3 +ORDER BY state NULLS LAST, ranking, net_profit DESC NULLS LAST, county NULLS LAST LIMIT 10 + + +-- ================================================================================ +-- Grouped aggregate then RANK -- exercises the q67 pattern where the input to WGL +-- comes from an aggregation whose output feeds the ranked window. +-- ================================================================================ + +query +SELECT category, item, total FROM ( + SELECT category, item, total, + RANK() OVER (PARTITION BY category ORDER BY total DESC NULLS LAST) AS rk + FROM ( + SELECT category, item, sum(revenue) AS total + FROM q67_sales + GROUP BY category, item + ) agg +) t WHERE rk <= 2 ORDER BY category NULLS LAST, rk, total DESC NULLS LAST, item + + +-- ================================================================================ +-- Global RANK top-K (empty PARTITION BY) -- routes through Comet's streaming +-- `PartitionedRankLimitExec` with `partition_prefix_len == 0`. State never resets +-- across the DF partition; Spark's WGL Partial phase runs per input partition and +-- the Final phase runs after a SinglePartition shuffle. Distilled from TPC-DS q44 +-- (`rank() over (order by rank_col asc) rnk ... where rnk < 11`). +-- ================================================================================ + +-- q44 shape: global RANK ASC, filter `rnk < 11`. +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (ORDER BY ord ASC NULLS LAST) AS rnk + FROM test_rank +) t WHERE rnk < 4 ORDER BY rnk, ord ASC NULLS LAST + +-- Global RANK DESC with ties at the top: `rk <= 1` returns EVERY row tied at +-- the max, `rk <= 2` also returns them (rank 2 is skipped because two ties push it +-- out). +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 1 ORDER BY rk, ord DESC NULLS LAST, part NULLS LAST + +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 2 ORDER BY rk, ord DESC NULLS LAST, part NULLS LAST + +-- Global RANK with `<` filter form. +query +SELECT part, ord FROM ( + SELECT part, ord, + RANK() OVER (ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk < 3 ORDER BY rk, ord DESC NULLS LAST, part NULLS LAST + +-- Global RANK with outer LIMIT (q44 uses `LIMIT 100` on the outer query). +query +SELECT part, ord, rk FROM ( + SELECT part, ord, + RANK() OVER (ORDER BY ord DESC NULLS LAST) AS rk + FROM test_rank +) t WHERE rk <= 3 ORDER BY rk, ord DESC NULLS LAST, part NULLS LAST LIMIT 5 + +-- Global ROW_NUMBER without PARTITION BY should route to LocalLimitExec (Spark +-- rewrites this pattern to plain Limit under low thresholds, but the WGL path is +-- exercised at threshold=1000). +query +SELECT part, ord FROM ( + SELECT part, ord, + ROW_NUMBER() OVER (ORDER BY ord DESC NULLS LAST, part) AS rn + FROM test_rank +) t WHERE rn <= 3 ORDER BY rn diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql new file mode 100644 index 00000000000..4a20ef4b787 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql @@ -0,0 +1,106 @@ +-- 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. + +-- WindowGroupLimit tie-semantics for RANK vs DENSE_RANK. RANK leaves gaps on ties +-- (`rk <= 2` skips rank 2 if two rows tied at 1), DENSE_RANK does not. Both rank-like +-- functions route to Comet's streaming `PartitionedRankLimitExec` on the native side, +-- which preserves Spark's tie ordering (via `RankLimitIterator`). Strict `query` mode +-- verifies both correctness and full native operator coverage. +-- The `rk = 0` case at the bottom collapses to an empty LocalRelation via Spark's +-- optimizer under threshold=1000, so it uses `spark_answer_only` — the resulting +-- LocalTableScan is not natively supported without `spark.comet.exec.localTableScan.enabled`. + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 +-- Config: spark.comet.exec.localTableScan.enabled=true + +statement +CREATE TABLE test_wgl_rank(grp string, score int) USING parquet + +statement +INSERT INTO test_wgl_rank VALUES + ('a', 100), ('a', 100), ('a', 90), ('a', 90), ('a', 80), ('a', 70), + ('b', 50), ('b', 40), ('b', 40), ('b', NULL), + ('c', 1), + (NULL, 42), (NULL, 42), + ('d', NULL), ('d', NULL) + +-- RANK <= 2: two tied #1 rows fill the "first two slots" and rank 2 is skipped. +query +SELECT grp, score FROM ( + SELECT grp, score, + RANK() OVER (PARTITION BY grp ORDER BY score DESC NULLS LAST) AS rk + FROM test_wgl_rank +) t WHERE rk <= 2 ORDER BY grp NULLS LAST, rk, score DESC NULLS LAST + +-- DENSE_RANK <= 2: top-2 distinct order values per partition (ties still counted once). +query +SELECT grp, score FROM ( + SELECT grp, score, + DENSE_RANK() OVER (PARTITION BY grp ORDER BY score DESC NULLS LAST) AS dr + FROM test_wgl_rank +) t WHERE dr <= 2 ORDER BY grp NULLS LAST, dr, score DESC NULLS LAST + +-- RANK = 1: all rows tied at the top. +query +SELECT grp, score FROM ( + SELECT grp, score, + RANK() OVER (PARTITION BY grp ORDER BY score DESC NULLS LAST) AS rk + FROM test_wgl_rank +) t WHERE rk = 1 ORDER BY grp NULLS LAST, score DESC NULLS LAST + +-- DENSE_RANK < 3. +query +SELECT grp, score FROM ( + SELECT grp, score, + DENSE_RANK() OVER (PARTITION BY grp ORDER BY score DESC NULLS LAST) AS dr + FROM test_wgl_rank +) t WHERE dr < 3 ORDER BY grp NULLS LAST, dr, score DESC NULLS LAST + +-- RANK with no PARTITION BY. Unlike ROW_NUMBER, Spark still inserts WGL here. +query +SELECT grp, score FROM ( + SELECT grp, score, + RANK() OVER (ORDER BY score DESC NULLS LAST) AS rk + FROM test_wgl_rank +) t WHERE rk <= 2 ORDER BY rk, grp NULLS LAST + +-- DENSE_RANK with no PARTITION BY. +query +SELECT grp, score FROM ( + SELECT grp, score, + DENSE_RANK() OVER (ORDER BY score DESC NULLS LAST) AS dr + FROM test_wgl_rank +) t WHERE dr <= 2 ORDER BY dr, grp NULLS LAST + +-- Two rank-like windows in one SELECT, both filtered. Spark picks the minimum implied +-- limit (row_number wins as cheaper), and both filters still apply to the result. +query +SELECT grp, score FROM ( + SELECT grp, score, + ROW_NUMBER() OVER (PARTITION BY grp ORDER BY score DESC NULLS LAST) AS rn, + RANK() OVER (PARTITION BY grp ORDER BY score DESC NULLS LAST) AS rk + FROM test_wgl_rank +) t WHERE rn <= 3 AND rk <= 2 ORDER BY grp NULLS LAST, rn + +-- Limit = 0 collapses to empty. +query +SELECT grp, score FROM ( + SELECT grp, score, + RANK() OVER (PARTITION BY grp ORDER BY score DESC) AS rk + FROM test_wgl_rank +) t WHERE rk = 0 diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql new file mode 100644 index 00000000000..bdd5af5a7d2 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql @@ -0,0 +1,168 @@ +-- 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. + +-- WindowGroupLimit tests with ROW_NUMBER. +-- Spark's InferWindowGroupLimit optimizer rule (SPARK-37099, since 3.5.0) inserts a +-- WindowGroupLimitExec below Window when a rank-like output is filtered by <=, <, or = +-- against an integer literal. Threshold is toggled via ConfigMatrix to exercise both the +-- optimized (WGL-inserted) and the naive (Window + Filter) plan shapes. +-- +-- Comet routes the ROW_NUMBER + non-empty PARTITION BY pushdown case to DataFusion's +-- PartitionedTopKExec via `CometWindowGroupLimitExec`. Every pushdown-eligible query below +-- runs natively; queries where Spark's optimizer converts to a plain Limit (global top-K) +-- or skips WGL pushdown still exercise the ORDER BY + Filter path without a fallback. +-- +-- Tie-breaking note: ROW_NUMBER is non-deterministic when ORDER BY has duplicate values. +-- Spark's stream-based SimpleLimitIterator preserves input order; DataFusion's heap-based +-- PartitionedTopKExec does not. Every ORDER BY below therefore includes `hire_yr` as a +-- secondary key so the tests exercise WGL without hitting tie-breaking non-determinism. + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 +-- ConfigMatrix: parquet.enable.dictionary=false,true +-- The `rn = 0` case at line 99 collapses to an empty LocalRelation under threshold=1000, +-- which lowers to a `LocalTableScan`; enable native execution for it so strict `query` +-- mode passes. +-- Config: spark.comet.exec.localTableScan.enabled=true + +statement +CREATE TABLE test_wgl_rn(dept string, salary int, hire_yr int) USING parquet + +statement +INSERT INTO test_wgl_rn VALUES + ('eng', 100, 2020), + ('eng', 200, 2019), + ('eng', 200, 2021), + ('eng', 50, 2022), + ('eng', NULL, 2018), + ('sales', 90, 2020), + ('sales', 90, 2021), + ('sales', 300, 2015), + ('hr', 75, 2023), + (NULL, 500, 2010), + (NULL, NULL, NULL) + +-- Top-1 per partition, DESC. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn = 1 ORDER BY dept NULLS LAST + +-- Top-3 per partition with `<= k`. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 3 ORDER BY dept NULLS LAST, rn + +-- `< k` form (Spark decrements to `<= k-1`). +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn < 3 ORDER BY dept NULLS LAST, rn + +-- Literal on the left of the comparison. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE 3 >= rn ORDER BY dept NULLS LAST, rn + +-- Limit larger than any partition. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 100 ORDER BY dept NULLS LAST, rn + +-- Limit = 0 collapses to an empty relation. Under threshold=1000, Spark's +-- InferWindowGroupLimit rewrites this to a `LocalTableScan` (Comet-eligible via the +-- fixture's `spark.comet.exec.localTableScan.enabled=true`); under threshold=-1 the +-- query runs via Window+Filter. Both plans return empty. +query +SELECT dept, salary FROM ( + SELECT dept, salary, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) AS rn + FROM test_wgl_rn +) t WHERE rn = 0 + +-- ASC with NULLS FIRST. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary ASC NULLS FIRST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 2 ORDER BY dept NULLS LAST, rn + +-- ASC with NULLS LAST. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary ASC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 2 ORDER BY dept NULLS LAST, rn + +-- Multi-column ORDER BY (asc, desc). Already deterministic via the two-key ordering. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER ( + PARTITION BY dept + ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST + ) AS rn + FROM test_wgl_rn +) t WHERE rn <= 2 ORDER BY dept NULLS LAST, rn + +-- Multi-column PARTITION BY. Each (dept, hire_yr) pair is unique, so no ties. +query +SELECT dept, hire_yr, salary FROM ( + SELECT dept, hire_yr, salary, + ROW_NUMBER() OVER (PARTITION BY dept, hire_yr ORDER BY salary DESC NULLS LAST) AS rn + FROM test_wgl_rn +) t WHERE rn = 1 ORDER BY dept NULLS LAST, hire_yr NULLS LAST + +-- No PARTITION BY (global top-K). Spark converts this to a plain Limit for ROW_NUMBER +-- when limit < topKSortFallbackThreshold, so no WGL should appear. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 3 ORDER BY rn + +-- Extra AND predicate on a non-rank column. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 2 AND hire_yr >= 2019 ORDER BY dept NULLS LAST, rn + +-- Outer LIMIT after WGL pushdown. +query +SELECT dept, salary, hire_yr FROM ( + SELECT dept, salary, hire_yr, + ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC NULLS LAST, hire_yr ASC NULLS FIRST) AS rn + FROM test_wgl_rn +) t WHERE rn <= 2 ORDER BY dept NULLS LAST, rn LIMIT 3 diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql new file mode 100644 index 00000000000..29e3e2c38db --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql @@ -0,0 +1,70 @@ +-- 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. + +-- Repro for SPARK-46526-shaped correlated scalar subquery with `ORDER BY colX +-- LIMIT 1`. Spark decorrelates by wrapping the subquery in a rank window; if +-- the ORDER BY column collapses into the correlation (e.g. ORDER BY t2c when +-- the correlation predicate is `WHERE t2c = t1c`), the resulting Window may +-- have an EMPTY orderSpec but non-empty partitionSpec. Comet's serde must not +-- hand DataFusion's `PartitionedTopKExec` an ordering with only partition +-- keys -- that panics at execute time. +-- +-- Spark 3.5 rejects this correlation form at analysis +-- (UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.ACCESSING_OUTER_QUERY_COLUMN_IS_NOT_ALLOWED); +-- the decorrelation that produces the WGL landed in Spark 4.0. + +-- MinSparkVersion: 4.0 + +statement +CREATE TABLE t1(t1a string, t1b smallint, t1c int, t1d int) USING parquet + +statement +INSERT INTO t1 VALUES + ('t1a-1', cast(1 AS smallint), 1, 10), + ('t1a-2', cast(2 AS smallint), 2, 5), + ('t1a-3', cast(3 AS smallint), 3, NULL), + ('t1a-4', cast(4 AS smallint), NULL, 20) + +statement +CREATE TABLE t2(t2a string, t2b smallint, t2c int, t2d int) USING parquet + +statement +INSERT INTO t2 VALUES + ('t2a-1', cast(1 AS smallint), 1, 100), + ('t2a-1b', cast(11 AS smallint), 1, 200), + ('t2a-2', cast(2 AS smallint), 2, 300), + ('t2a-3', cast(3 AS smallint), 3, 400) + +-- Original failing query. +query +SELECT t1a, t1b +FROM t1 +WHERE t1c = (SELECT t2c + FROM t2 + WHERE t2c = t1c + ORDER BY t2c LIMIT 1) +ORDER BY t1a + +-- Same shape but with a non-degenerate ORDER BY column (t2d instead of t2c). +query +SELECT t1a, t1b +FROM t1 +WHERE t1c = (SELECT t2c + FROM t2 + WHERE t2c = t1c + ORDER BY t2d LIMIT 1) +ORDER BY t1a diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q44/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q44/extended.txt index 0db84abb506..abf74def7ce 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q44/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q44/extended.txt @@ -1,64 +1,59 @@ -TakeOrderedAndProject -+- Project - +- BroadcastHashJoin - :- Project - : +- BroadcastHashJoin - : :- Project - : : +- SortMergeJoin - : : :- Sort - : : : +- Project - : : : +- Filter - : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : : +- CometColumnarToRow - : : : +- CometSort - : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : : +- CometColumnarToRow - : : : +- CometSort - : : : +- CometFilter - : : : : +- Subquery - : : : : +- CometColumnarToRow - : : : : +- CometHashAggregate - : : : : +- CometExchange - : : : : +- CometHashAggregate - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- CometHashAggregate - : : : +- CometExchange - : : : +- CometHashAggregate - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : +- Sort - : : +- Project - : : +- Filter - : : +- Window - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : +- CometColumnarToRow - : : +- CometSort - : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : +- CometColumnarToRow - : : +- CometSort - : : +- CometFilter - : : : +- ReusedSubquery - : : +- CometHashAggregate - : : +- CometExchange - : : +- CometHashAggregate - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales - : +- BroadcastExchange - : +- CometColumnarToRow - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - +- BroadcastExchange - +- CometColumnarToRow +CometColumnarToRow ++- CometTakeOrderedAndProject + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometSortMergeJoin + : : :- CometSort + : : : +- CometProject + : : : +- CometFilter + : : : +- CometWindowExec + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometExchange + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometFilter + : : : : +- Subquery + : : : : +- CometColumnarToRow + : : : : +- CometHashAggregate + : : : : +- CometExchange + : : : : +- CometHashAggregate + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometHashAggregate + : : : +- CometExchange + : : : +- CometHashAggregate + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : +- CometSort + : : +- CometProject + : : +- CometFilter + : : +- CometWindowExec + : : +- CometWindowGroupLimitExec + : : +- CometSort + : : +- CometExchange + : : +- CometWindowGroupLimitExec + : : +- CometSort + : : +- CometFilter + : : : +- ReusedSubquery + : : +- CometHashAggregate + : : +- CometExchange + : : +- CometHashAggregate + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + +- CometBroadcastExchange +- CometProject +- CometFilter +- CometNativeScan parquet spark_catalog.default.item -Comet accelerated 32 out of 54 eligible operators (59%). Final plan contains 7 transitions between Spark and Comet. Accelerated expressions: 14 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 53 out of 54 eligible operators (98%). Final plan contains 2 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q67/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q67/extended.txt index 83193e5e657..1fa2938247b 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q67/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q67/extended.txt @@ -1,41 +1,40 @@ -TakeOrderedAndProject -+- Filter - +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - +- CometColumnarToRow +CometColumnarToRow ++- CometTakeOrderedAndProject + +- CometFilter + +- CometWindowExec + +- CometWindowGroupLimitExec +- CometSort - +- CometColumnarExchange - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - +- CometColumnarToRow - +- CometSort - +- CometHashAggregate - +- CometExchange - +- CometHashAggregate - +- CometExpand - +- CometProject - +- CometBroadcastHashJoin - :- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- CometSubqueryBroadcast - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- CometBroadcastExchange - +- CometProject - +- CometFilter - +- CometNativeScan parquet spark_catalog.default.item + +- CometExchange + +- CometWindowGroupLimitExec + +- CometSort + +- CometHashAggregate + +- CometExchange + +- CometHashAggregate + +- CometExpand + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometSubqueryBroadcast + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.store + +- CometBroadcastExchange + +- CometProject + +- CometFilter + +- CometNativeScan parquet spark_catalog.default.item -Comet accelerated 32 out of 37 eligible operators (86%). Final plan contains 2 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 37 out of 37 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt index ed3e7d1b913..c8b2cadf0e8 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt @@ -5,55 +5,52 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange - +- HashAggregate - +- Expand - +- Project - +- BroadcastHashJoin - :- CometColumnarToRow - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : +- CometSubqueryBroadcast - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.date_dim - +- BroadcastExchange - +- Project - +- BroadcastHashJoin - :- CometColumnarToRow - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- BroadcastExchange - +- Project - +- Filter - +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - +- CometColumnarToRow - +- CometSort - +- CometHashAggregate - +- CometExchange - +- CometHashAggregate - +- CometProject - +- CometBroadcastHashJoin - :- CometProject - : +- CometBroadcastHashJoin - : :- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : +- ReusedSubquery - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- CometBroadcastExchange - +- CometProject - +- CometFilter - +- CometNativeScan parquet spark_catalog.default.date_dim + +- CometExchange + +- CometHashAggregate + +- CometExpand + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : +- CometSubqueryBroadcast + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.date_dim + +- CometBroadcastExchange + +- CometProject + +- CometBroadcastHashJoin + :- CometFilter + : +- CometNativeScan parquet spark_catalog.default.store + +- CometBroadcastExchange + +- CometProject + +- CometFilter + +- CometWindowExec + +- CometWindowGroupLimitExec + +- CometSort + +- CometHashAggregate + +- CometExchange + +- CometHashAggregate + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : +- ReusedSubquery + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.store + +- CometBroadcastExchange + +- CometProject + +- CometFilter + +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 40 out of 52 eligible operators (76%). Final plan contains 4 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 52 out of 52 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q44/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q44/extended.txt index e113f28d063..00cd5e33b10 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q44/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q44/extended.txt @@ -7,52 +7,48 @@ CometColumnarToRow : :- CometProject : : +- CometSortMergeJoin : : :- CometSort - : : : +- CometColumnarExchange - : : : +- Project - : : : +- Filter - : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : : +- CometColumnarToRow - : : : +- CometSort - : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : : +- CometColumnarToRow - : : : +- CometSort - : : : +- CometFilter - : : : : +- Subquery - : : : : +- CometColumnarToRow - : : : : +- CometHashAggregate - : : : : +- CometExchange - : : : : +- CometHashAggregate - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometWindowExec + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometExchange + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometFilter + : : : : +- Subquery + : : : : +- CometColumnarToRow + : : : : +- CometHashAggregate + : : : : +- CometExchange + : : : : +- CometHashAggregate + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometHashAggregate + : : : +- CometExchange : : : +- CometHashAggregate - : : : +- CometExchange - : : : +- CometHashAggregate - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales : : +- CometSort - : : +- CometColumnarExchange - : : +- Project - : : +- Filter - : : +- Window - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : +- CometColumnarToRow - : : +- CometSort - : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : : +- CometColumnarToRow - : : +- CometSort - : : +- CometFilter - : : : +- ReusedSubquery + : : +- CometExchange + : : +- CometProject + : : +- CometFilter + : : +- CometWindowExec + : : +- CometWindowGroupLimitExec + : : +- CometSort + : : +- CometExchange + : : +- CometWindowGroupLimitExec + : : +- CometSort + : : +- CometFilter + : : : +- ReusedSubquery + : : +- CometHashAggregate + : : +- CometExchange : : +- CometHashAggregate - : : +- CometExchange - : : +- CometHashAggregate - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales : +- CometBroadcastExchange : +- CometProject : +- CometFilter @@ -62,4 +58,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.item -Comet accelerated 45 out of 56 eligible operators (80%). Final plan contains 6 transitions between Spark and Comet. Accelerated expressions: 14 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 55 out of 56 eligible operators (98%). Final plan contains 2 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q67a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q67a/extended.txt index 6f95655293b..6af947205a3 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q67a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q67a/extended.txt @@ -1,289 +1,288 @@ -TakeOrderedAndProject -+- Filter - +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - +- CometColumnarToRow +CometColumnarToRow ++- CometTakeOrderedAndProject + +- CometFilter + +- CometWindowExec + +- CometWindowGroupLimitExec +- CometSort - +- CometColumnarExchange - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - +- CometColumnarToRow - +- CometSort - +- CometUnion - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - :- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : : +- CometSubqueryBroadcast - : : : : +- CometBroadcastExchange - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - +- CometHashAggregate - +- CometExchange + +- CometExchange + +- CometWindowGroupLimitExec + +- CometSort + +- CometUnion + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : : +- CometSubqueryBroadcast + : : : : +- CometBroadcastExchange + : : : : +- CometProject + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + +- CometHashAggregate + +- CometExchange + +- CometHashAggregate +- CometHashAggregate - +- CometHashAggregate - +- CometExchange - +- CometHashAggregate - +- CometProject - +- CometBroadcastHashJoin - :- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- CometSubqueryBroadcast - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- CometBroadcastExchange - +- CometProject - +- CometFilter - +- CometNativeScan parquet spark_catalog.default.item + +- CometExchange + +- CometHashAggregate + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometSubqueryBroadcast + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.store + +- CometBroadcastExchange + +- CometProject + +- CometFilter + +- CometNativeScan parquet spark_catalog.default.item -Comet accelerated 280 out of 285 eligible operators (98%). Final plan contains 2 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 285 out of 285 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt index 25cf0af8d17..8e5298ac981 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt @@ -9,160 +9,151 @@ CometColumnarToRow +- CometHashAggregate +- CometUnion :- CometHashAggregate - : +- CometColumnarExchange - : +- HashAggregate - : +- Project - : +- BroadcastHashJoin - : :- CometColumnarToRow - : : +- CometProject - : : +- CometBroadcastHashJoin - : : :- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- CometSubqueryBroadcast - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- BroadcastExchange - : +- Project - : +- BroadcastHashJoin - : :- CometColumnarToRow - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- BroadcastExchange - : +- Project - : +- Filter - : +- Window - : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : +- CometColumnarToRow - : +- CometSort - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- ReusedSubquery - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometSubqueryBroadcast + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometWindowExec + : +- CometWindowGroupLimitExec + : +- CometSort + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- ReusedSubquery + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.date_dim :- CometHashAggregate : +- CometExchange : +- CometHashAggregate : +- CometHashAggregate - : +- CometColumnarExchange - : +- HashAggregate - : +- Project - : +- BroadcastHashJoin - : :- CometColumnarToRow - : : +- CometProject - : : +- CometBroadcastHashJoin - : : :- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- CometSubqueryBroadcast - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- BroadcastExchange - : +- Project - : +- BroadcastHashJoin - : :- CometColumnarToRow - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- BroadcastExchange - : +- Project - : +- Filter - : +- Window - : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - : +- CometColumnarToRow - : +- CometSort - : +- CometHashAggregate - : +- CometExchange - : +- CometHashAggregate - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : : +- ReusedSubquery - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometSubqueryBroadcast + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometWindowExec + : +- CometWindowGroupLimitExec + : +- CometSort + : +- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- ReusedSubquery + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.date_dim +- CometHashAggregate +- CometExchange +- CometHashAggregate +- CometHashAggregate - +- CometColumnarExchange - +- HashAggregate - +- Project - +- BroadcastHashJoin - :- CometColumnarToRow - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : +- CometSubqueryBroadcast - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.date_dim - +- BroadcastExchange - +- Project - +- BroadcastHashJoin - :- CometColumnarToRow - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- BroadcastExchange - +- Project - +- Filter - +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] - +- CometColumnarToRow - +- CometSort - +- CometHashAggregate - +- CometExchange - +- CometHashAggregate - +- CometProject - +- CometBroadcastHashJoin - :- CometProject - : +- CometBroadcastHashJoin - : :- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales - : : +- ReusedSubquery - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- CometBroadcastExchange - +- CometProject - +- CometFilter - +- CometNativeScan parquet spark_catalog.default.date_dim + +- CometExchange + +- CometHashAggregate + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : +- CometSubqueryBroadcast + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.date_dim + +- CometBroadcastExchange + +- CometProject + +- CometBroadcastHashJoin + :- CometFilter + : +- CometNativeScan parquet spark_catalog.default.store + +- CometBroadcastExchange + +- CometProject + +- CometFilter + +- CometWindowExec + +- CometWindowGroupLimitExec + +- CometSort + +- CometHashAggregate + +- CometExchange + +- CometHashAggregate + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : +- ReusedSubquery + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.store + +- CometBroadcastExchange + +- CometProject + +- CometFilter + +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 120 out of 153 eligible operators (78%). Final plan contains 10 transitions between Spark and Comet. Accelerated expressions: 12 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 153 out of 153 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 0 codegen dispatch. \ No newline at end of file From 7ef79eb05b86bf826428934883da4acdb5aefcea Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 21 Aug 2026 09:40:11 -0700 Subject: [PATCH 2/6] feat: support `WindowGroupLimit` --- docs/source/contributor-guide/roadmap.md | 5 +- .../latest/compatibility/floating-point.md | 13 + .../latest/compatibility/operators.md | 4 +- docs/source/user-guide/latest/operators.md | 2 +- .../src/execution/operators/rank_limit.rs | 351 +++++++++++++++--- native/core/src/execution/planner.rs | 40 +- .../sql/comet/CometWindowGroupLimitExec.scala | 47 ++- .../shims/ShimCometWindowGroupLimit.scala | 10 +- .../shims/ShimCometWindowGroupLimit.scala | 10 +- .../window_group_limit_datatypes.sql | 0 .../window_group_limit_edge.sql | 0 .../window_group_limit_rank.sql | 29 +- .../window_group_limit_rank_dense_rank.sql | 0 .../window_group_limit_row_number.sql | 14 +- .../window_group_limit_scalar_subquery.sql | 4 +- 15 files changed, 406 insertions(+), 123 deletions(-) rename spark/src/test/resources/sql-tests/{expressions/window => windows}/window_group_limit_datatypes.sql (100%) rename spark/src/test/resources/sql-tests/{expressions/window => windows}/window_group_limit_edge.sql (100%) rename spark/src/test/resources/sql-tests/{expressions/window => windows}/window_group_limit_rank.sql (93%) rename spark/src/test/resources/sql-tests/{expressions/window => windows}/window_group_limit_rank_dense_rank.sql (100%) rename spark/src/test/resources/sql-tests/{expressions/window => windows}/window_group_limit_row_number.sql (92%) rename spark/src/test/resources/sql-tests/{expressions/window => windows}/window_group_limit_scalar_subquery.sql (95%) diff --git a/docs/source/contributor-guide/roadmap.md b/docs/source/contributor-guide/roadmap.md index aa7b643232e..4f0dfd8b6d5 100644 --- a/docs/source/contributor-guide/roadmap.md +++ b/docs/source/contributor-guide/roadmap.md @@ -29,8 +29,8 @@ Native window execution runs by default (`spark.comet.exec.window.enabled`). The `first_value`, `last_value`), and the `count`, `min`, `max`, `sum`, and `avg` aggregates are accelerated. Remaining work is to close the gaps that still fall back to Spark: statistical aggregates (`stddev`, `variance`, `corr`, `covar`) and `collect_list` / `collect_set` as window functions ([#4766]), `GROUPS` frames ([#4836]), `RANGE` frames with explicit date or -decimal offsets ([#4834]), `first_value` / `last_value` on `RANGE` frames with a literal offset ([#4835]), -non-literal `lag` / `lead` default values ([#4268]), and `WindowGroupLimitExec` ([#4837]). See the +decimal offsets ([#4834]), `first_value` / `last_value` on `RANGE` frames with a literal offset ([#4835]), and +non-literal `lag` / `lead` default values ([#4268]). See the [window function compatibility guide](../user-guide/latest/compatibility/operators.md) for the complete list of supported functions, frames, and fallback cases. @@ -39,7 +39,6 @@ supported functions, frames, and fallback cases. [#4834]: https://github.com/apache/datafusion-comet/issues/4834 [#4835]: https://github.com/apache/datafusion-comet/issues/4835 [#4836]: https://github.com/apache/datafusion-comet/issues/4836 -[#4837]: https://github.com/apache/datafusion-comet/issues/4837 ## Native Lambda Evaluation diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index ffab0550609..a068e12a500 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -27,3 +27,16 @@ So Comet adds additional normalization expression of NaN and zero for comparison to Spark in some cases, especially when the data contains both positive and negative zero. This is likely an edge case that is not of concern for many users. If it is a concern, setting `spark.comet.exec.strictFloatingPoint=true` will make relevant operations fall back to Spark. + +## Ordering: signed zero (`-0.0` vs `+0.0`) + +Spark's `ORDER BY`, `RANK`, `DENSE_RANK`, and window frame comparisons route through +`SQLOrderingUtil.compareDoubles` / `compareFloats`, which explicitly define `-0.0 == 0.0`. Comet's +native sort and `WindowGroupLimitExec` use the `arrow-row` row-format encoder for `ORDER BY` keys, +which applies Rust's total-ordering transform to the raw IEEE-754 bits. Under that encoding `-0.0` +sorts strictly less than `+0.0`, so a partition that mixes the two zeros can produce a rank +distribution that differs from Spark. For example, `RANK() OVER (ORDER BY v ASC)` over +`[-0.0, 0.0, 1.0]` filtered to `rk <= 1` returns two rows in Spark (both zeros tied at rank 1) but +one row in Comet (`-0.0` at rank 1, `+0.0` at rank 2). If your workload materially mixes `-0.0` +and `+0.0` in a ranked column, prefer Spark for that stage or normalize the column to `+0.0` +upstream. diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 13aa6d943ca..2c8eea4a5d3 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -71,8 +71,8 @@ incorrect result. When any single window expression in a `WindowExec` falls back window are not supported by Spark either. - Any `PARTITION BY` or `ORDER BY` expression that Comet cannot serialize. -`WindowGroupLimitExec` (window-based limit pushdown) is not yet supported and falls back to Spark -([#4837](https://github.com/apache/datafusion-comet/issues/4837)). +`WindowGroupLimitExec` (window-based limit pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`) +runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` (default: true). ## Round-Robin Partitioning diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 42235b1367e..20e52a9b1a7 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -103,7 +103,7 @@ omitted from the tables below and may be reconsidered based on demand: | Operator | Status | Notes | | ---------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `WindowExec` | ⚠️ | Runs natively and is enabled by default. A broad set of window functions is accelerated; unsupported shapes fall back to Spark. See [window function compatibility](compatibility/operators.md). | -| `WindowGroupLimitExec` | 🔜 | Window-based limit pushdown falls back today ([#4837](https://github.com/apache/datafusion-comet/issues/4837)). | +| `WindowGroupLimitExec` | ✅ | Streaming per-partition top-K pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`. | ## Generators and set operations diff --git a/native/core/src/execution/operators/rank_limit.rs b/native/core/src/execution/operators/rank_limit.rs index 8d8c0ddce8d..d6bd0e81e78 100644 --- a/native/core/src/execution/operators/rank_limit.rs +++ b/native/core/src/execution/operators/rank_limit.rs @@ -35,12 +35,13 @@ 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::{DataFusionError, Result}; +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, @@ -57,50 +58,51 @@ pub enum WindowFnKind { #[derive(Debug)] pub struct PartitionedRankLimitExec { input: Arc, - /// Full sort expression `[partition_keys..., order_keys...]`. - expr: LexOrdering, - /// Leading count of expressions in `expr` that form the partition key. - /// Zero means "no PARTITION BY" (global top-K within each input partition). - /// Can equal `expr.len()` when `LexOrdering::new` dedup collapses the ORDER BY suffix. - partition_prefix_len: usize, + /// PARTITION BY expressions. Empty means "no PARTITION BY" (global top-K + /// within each input DataFusion partition). + partition_keys: Vec, + /// ORDER BY expressions. Empty means "no ORDER BY" and every row within a + /// partition ties. + order_keys: Vec, fetch: usize, kind: WindowFnKind, cache: Arc, + metrics: ExecutionPlanMetricsSet, } impl PartitionedRankLimitExec { pub fn try_new( input: Arc, - expr: LexOrdering, - partition_prefix_len: usize, + partition_keys: Vec, + order_keys: Vec, fetch: usize, kind: WindowFnKind, ) -> Result { - // Guard against `LexOrdering::new` dedup dropping a partition key. - if partition_prefix_len > expr.len() { - return Err(DataFusionError::Internal(format!( - "PartitionedRankLimitExec: partition prefix ({partition_prefix_len}) exceeds \ - ordering length ({})", - expr.len() - ))); - } - let cache = Arc::new(Self::compute_properties(&input, &expr)?); + let cache = Arc::new(Self::compute_properties( + &input, + &partition_keys, + &order_keys, + )?); Ok(Self { input, - expr, - partition_prefix_len, + partition_keys, + order_keys, fetch, kind, cache, + metrics: ExecutionPlanMetricsSet::new(), }) } fn compute_properties( input: &Arc, - sort_exprs: &LexOrdering, + partition_keys: &[PhysicalSortExpr], + order_keys: &[PhysicalSortExpr], ) -> Result { let mut eq_properties = input.equivalence_properties().clone(); - eq_properties.reorder(sort_exprs.clone())?; + if let Some(ordering) = full_ordering(partition_keys, order_keys) { + eq_properties.reorder(ordering)?; + } Ok(PlanProperties::new( eq_properties, input.output_partitioning().clone(), @@ -110,15 +112,45 @@ impl PartitionedRankLimitExec { } } +/// `[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 { + let sort_exprs: Vec = 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 => write!( - f, - "CometPartitionedRankLimitExec: kind={:?}, fetch={}, partition_prefix_len={}, \ - expr=[{}]", - self.kind, self.fetch, self.partition_prefix_len, self.expr - ), + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let partition = self + .partition_keys + .iter() + .map(|e| e.to_string()) + .collect::>() + .join(", "); + let order = self + .order_keys + .iter() + .map(|e| e.to_string()) + .collect::>() + .join(", "); + write!( + f, + "CometPartitionedRankLimitExec: kind={:?}, fetch={}, partition_by=[{}], order_by=[{}]", + self.kind, self.fetch, partition, order + ) + } DisplayFormatType::TreeRender => unimplemented!(), } } @@ -144,26 +176,31 @@ impl ExecutionPlan for PartitionedRankLimitExec { assert_eq!(children.len(), 1); Ok(Arc::new(PartitionedRankLimitExec::try_new( Arc::clone(&children[0]), - self.expr.clone(), - self.partition_prefix_len, + 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...]`. Declaring this lets DataFusion's - // `EnforceSorting` insert a `SortExec` if the sort somehow got dropped - // during plan construction, though in practice Spark's Catalyst already - // injects the sort above `WindowGroupLimitExec`. + // `[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> { - vec![Some(OrderingRequirements::from(self.expr.clone()))] + vec![full_ordering(&self.partition_keys, &self.order_keys).map(OrderingRequirements::from)] } fn maintains_input_order(&self) -> Vec { vec![true] } + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + fn execute( &self, partition: usize, @@ -172,18 +209,17 @@ impl ExecutionPlan for PartitionedRankLimitExec { let input = self.input.execute(partition, context)?; let schema = input.schema(); - let partition_key = build_key_encoder(&self.expr[..self.partition_prefix_len], &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, or every ORDER BY - // column was deduplicated with a PARTITION BY column by `LexOrdering::new`), - // `build_key_encoder` returns `None` and every row within a partition ties. + // 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.expr[self.partition_prefix_len..], &schema)? + build_key_encoder(&self.order_keys, &schema)? }; Ok(Box::pin(RankLimitStream { @@ -193,10 +229,12 @@ impl ExecutionPlan for PartitionedRankLimitExec { 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, })) } } @@ -247,12 +285,12 @@ struct RankLimitStream { schema: SchemaRef, /// `None` when there is no PARTITION BY (global top-K per DF input partition). partition_key: Option, - /// `None` when the ORDER BY suffix is empty (fully covered by PARTITION BY - /// or absent entirely), and always `None` for ROW_NUMBER (rank formula - /// never reads order keys). + /// `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, limit: u64, kind: WindowFnKind, + baseline_metrics: BaselineMetrics, // Per-partition streaming state, persisted across batches. prev_partition: Option, @@ -261,6 +299,11 @@ struct RankLimitStream { 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 { @@ -275,11 +318,10 @@ impl RankLimitStream { .as_ref() .map(|k| k.encode(batch)) .transpose()?; - let order_rows = self - .order_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 = None; let mut mask_builder = BooleanBufferBuilder::new(num_rows); let mut kept: usize = 0; @@ -297,12 +339,24 @@ impl RankLimitStream { self.prev_order = None; self.rank = 0; self.count = 0; + self.partition_exhausted = false; } - // 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 suffix — `prev_order` stays `None` across - // the whole partition in that case. + 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) @@ -325,6 +379,11 @@ impl RankLimitStream { 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; @@ -356,13 +415,23 @@ impl Stream for RankLimitStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { match self.input.poll_next_unpin(cx) { - Poll::Ready(Some(Ok(batch))) => match self.process_batch(&batch) { - // Skip fully-filtered batches so downstream never sees - // spurious empty batches between real ones. - Ok(out) if out.num_rows() == 0 => continue, - Ok(out) => return Poll::Ready(Some(Ok(out))), - Err(e) => return Poll::Ready(Some(Err(e))), - }, + Poll::Ready(Some(Ok(batch))) => { + let processed = { + let _timer = self.baseline_metrics.elapsed_compute().timer(); + self.process_batch(&batch) + }; + match processed { + // Skip fully-filtered batches so downstream never sees + // spurious empty batches between real ones. + Ok(out) if out.num_rows() == 0 => continue, + Ok(out) => { + return self + .baseline_metrics + .record_poll(Poll::Ready(Some(Ok(out)))); + } + Err(e) => return Poll::Ready(Some(Err(e))), + } + } Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), Poll::Ready(None) => return Poll::Ready(None), Poll::Pending => return Poll::Pending, @@ -376,3 +445,165 @@ impl RecordBatchStream for RankLimitStream { Arc::clone(&self.schema) } } + +#[cfg(test)] +mod tests { + //! Multi-batch state-persistence tests for `RankLimitStream`. The SQL fixtures under + //! `spark/src/test/resources/sql-tests/windows/window_group_limit_*.sql` are single-batch + //! sized under the default `COMET_BATCH_SIZE`, so cross-`poll_next` carry of + //! `prev_partition`, `prev_order`, `rank`, `count`, and `partition_exhausted` is only + //! covered here. + + use super::*; + use arrow::array::{Int32Array, Int64Array}; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::collect; + use datafusion::prelude::SessionContext; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("part", DataType::Int32, false), + Field::new("ord", DataType::Int64, false), + ])) + } + + fn batch(part: Vec, ord: Vec) -> RecordBatch { + RecordBatch::try_new( + schema(), + vec![ + Arc::new(Int32Array::from(part)) as ArrayRef, + Arc::new(Int64Array::from(ord)) as ArrayRef, + ], + ) + .unwrap() + } + + /// `PARTITION BY part ORDER BY ord ASC`. + fn keys() -> (Vec, Vec) { + let partition = vec![PhysicalSortExpr { + expr: Arc::new(Column::new("part", 0)) as Arc, + options: SortOptions::default(), + }]; + let order = vec![PhysicalSortExpr { + expr: Arc::new(Column::new("ord", 1)) as Arc, + options: SortOptions::default(), + }]; + (partition, order) + } + + async fn run(batches: Vec, fetch: usize, kind: WindowFnKind) -> Vec { + let (partition_keys, order_keys) = keys(); + let input = MemorySourceConfig::try_new_exec(&[batches], schema(), None).unwrap(); + let plan = Arc::new( + PartitionedRankLimitExec::try_new(input, partition_keys, order_keys, fetch, kind) + .unwrap(), + ); + collect(plan, SessionContext::new().task_ctx()) + .await + .unwrap() + } + + fn flatten(batches: &[RecordBatch]) -> Vec<(i32, i64)> { + let mut out = vec![]; + for b in batches { + let p = b.column(0).as_any().downcast_ref::().unwrap(); + let o = b.column(1).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + out.push((p.value(i), o.value(i))); + } + } + out + } + + /// Batch 0 ends exactly on the last row of partition 1; batch 1 opens partition 2. Verifies + /// the reset branch fires cleanly at a batch boundary. + #[tokio::test] + async fn partition_boundary_aligned_with_batch_boundary() { + let b0 = batch(vec![1, 1, 1], vec![10, 20, 30]); + let b1 = batch(vec![2, 2, 2], vec![40, 50, 60]); + let out = run(vec![b0, b1], 2, WindowFnKind::RowNumber).await; + assert_eq!(flatten(&out), vec![(1, 10), (1, 20), (2, 40), (2, 50)]); + } + + /// Partition 1 hits the limit inside batch 0 (extra rows must be dropped); batch 1 opens + /// partition 2 whose first row must not inherit partition 1's `rank` / `count`. + #[tokio::test] + async fn limit_hit_mid_batch_new_partition_next_batch() { + let b0 = batch(vec![1, 1, 1, 1], vec![10, 20, 30, 40]); + let b1 = batch(vec![2, 2], vec![50, 60]); + let out = run(vec![b0, b1], 2, WindowFnKind::RowNumber).await; + assert_eq!(flatten(&out), vec![(1, 10), (1, 20), (2, 50), (2, 60)]); + } + + /// An empty batch in the middle of the stream must not advance state and must not surface + /// downstream as a zero-row output batch. + #[tokio::test] + async fn empty_batch_between_real_batches() { + let b0 = batch(vec![1, 1], vec![10, 20]); + let empty = RecordBatch::new_empty(schema()); + let b1 = batch(vec![1, 1], vec![30, 40]); + let out = run(vec![b0, empty, b1], 3, WindowFnKind::RowNumber).await; + assert!(out.iter().all(|b| b.num_rows() > 0)); + assert_eq!(flatten(&out), vec![(1, 10), (1, 20), (1, 30)]); + } + + /// A RANK tie run straddles the batch boundary. `prev_order` must survive across + /// `poll_next` so tied rows keep the same rank as the last row of the prior batch. The + /// second sub-case breaks the tie on the last row of batch 1 with `fetch = 1`: the broken + /// tie ranks 1 (>= fetch) and must be dropped. + #[tokio::test] + async fn rank_tie_run_spans_batch_boundary() { + let out = run( + vec![ + batch(vec![1, 1], vec![10, 10]), + batch(vec![1, 1], vec![10, 10]), + ], + 1, + WindowFnKind::Rank, + ) + .await; + assert_eq!(flatten(&out).len(), 4); + + let out = run( + vec![ + batch(vec![1, 1], vec![10, 10]), + batch(vec![1, 1], vec![10, 20]), + ], + 1, + WindowFnKind::Rank, + ) + .await; + assert_eq!(flatten(&out), vec![(1, 10), (1, 10), (1, 10)]); + } + + /// DENSE_RANK must not skip a rank across a batch-boundary tie run. Two rows tie at rank 1 + /// in batch 0, batch 1 opens a distinct value that must rank 2 (not 3). + #[tokio::test] + async fn dense_rank_no_skip_across_batch_boundary() { + let out = run( + vec![ + batch(vec![1, 1], vec![10, 10]), + batch(vec![1, 1], vec![20, 30]), + ], + 2, + WindowFnKind::DenseRank, + ) + .await; + assert_eq!(flatten(&out), vec![(1, 10), (1, 10), (1, 20)]); + } + + /// A batch that lands entirely inside an already-exhausted partition must not encode + /// order rows. This test only pins the visible behavior (all rows dropped); the cost win + /// is covered by the fact that `order_key.encode` is not called on the fast path. + #[tokio::test] + async fn batch_entirely_inside_exhausted_partition_is_dropped() { + let b0 = batch(vec![1, 1, 1], vec![10, 20, 30]); + // Batch 1 stays inside partition 1 after the limit is hit at rank 1 in batch 0. + let b1 = batch(vec![1, 1, 1], vec![40, 50, 60]); + let out = run(vec![b0, b1], 1, WindowFnKind::RowNumber).await; + assert_eq!(flatten(&out), vec![(1, 10)]); + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 8f9a193ab09..0a99855f256 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2423,46 +2423,34 @@ impl PhysicalPlanner { // Partition keys arrive as bare exprs (no direction). Spark's WGL // requires the child to be sorted by partition-keys ASCENDING then by // order-by keys, so materialize partition keys with SortOptions matching - // Spark's `SortOrder(_, Ascending)` (ascending, nulls_first) and - // concatenate with the ORDER BY exprs into a single LexOrdering - // `[partition_keys, order_keys]`. - let mut sort_exprs: Vec = - Vec::with_capacity(partition_prefix_len + wgl.order_by_list.len()); + // Spark's `SortOrder(_, Ascending)` (ascending, nulls_first). + let mut partition_keys: Vec = + Vec::with_capacity(partition_prefix_len); for expr in &wgl.partition_by_list { let phys = self.create_expr(expr, Arc::clone(&input_schema))?; - sort_exprs.push(PhysicalSortExpr { + partition_keys.push(PhysicalSortExpr { expr: phys, options: SortOptions::default(), }); } + let mut order_keys: Vec = + Vec::with_capacity(wgl.order_by_list.len()); for expr in &wgl.order_by_list { - sort_exprs.push(self.create_sort_expr(expr, Arc::clone(&input_schema))?); + order_keys.push(self.create_sort_expr(expr, Arc::clone(&input_schema))?); } - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - GeneralError("WindowGroupLimit produced empty LexOrdering".to_string()) - })?; - - // `LexOrdering::new` deduplicates by underlying `PhysicalExpr`, so if a - // PARTITION BY column also appears in ORDER BY (e.g. `PARTITION BY t2c - // ORDER BY t2c`, produced by SPARK-46526-shaped scalar-subquery - // rewrites) the ORDER BY suffix collapses away. The streaming operator - // handles that tie-everything degenerate case (every row shares the - // K-th ORDER BY value) the same way it handles any other tie: rank - // stays at 0 for RANK/DENSE_RANK, or increments per row for ROW_NUMBER. - // Always route to the streaming `PartitionedRankLimitExec`. Spark's // `WindowGroupLimitExec.requiredChildOrdering` guarantees the child is // sorted by `[partition_keys..., order_keys...]`, so a single pass - // suffices. `PartitionedTopKExec` (heap-based) is faster asymptotically - // but reorders tied rows, which breaks Spark's `SimpleLimitIterator` / - // `RankLimitIterator` semantics: for `ROW_NUMBER()`, Spark assigns ranks - // in input order to rows tied on the ORDER BY keys, and any deviation - // shows up as a wrong-row diff in SPARK-37099-shaped tests. + // suffices. A heap-based top-K would be faster asymptotically but would + // reorder rows tied on the ORDER BY keys, which breaks Spark's + // `SimpleLimitIterator` / `RankLimitIterator` semantics: for + // `ROW_NUMBER()`, Spark assigns ranks in input order to tied rows, and + // any deviation surfaces as wrong-row diffs in SPARK-37099-shaped tests. Arc::new(PartitionedRankLimitExec::try_new( Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, + partition_keys, + order_keys, fetch, kind, )?) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala index d64d5950b25..0a1a12c40ca 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -46,12 +46,17 @@ import org.apache.comet.shims.ShimCometWindowGroupLimit */ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { - /** Fields extracted from a Spark `WindowGroupLimitExec` (Spark 3.5+). */ + /** + * 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) + limit: Int, + mode: String) override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED) @@ -60,9 +65,15 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { op: SparkPlan, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - // `nativeExecs` only routes here for a real `WindowGroupLimitExec` on Spark 3.5+, so the - // shim always returns Some. - val fields = ShimCometWindowGroupLimit.extract(op).get + // 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. @@ -103,16 +114,21 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { op.output, fields.partitionSpec, fields.orderSpec, + fields.rankLikeFunction, fields.limit, + fields.mode, op.children.head, SerializedPlan(None)) } } /** - * Comet physical plan node for Spark `WindowGroupLimitExec`. The Spark Partial/Final split is - * preserved unchanged in the Spark plan tree (each side is planned as its own native subtree), so - * the case class doesn't carry a mode field. + * Comet physical plan node for Spark `WindowGroupLimitExec`. `rankLikeFunction` and `mode` are + * carried explicitly so `equals` / `hashCode` distinguish Partial vs Final and ROW_NUMBER vs RANK + * vs DENSE_RANK for `ReuseSubquery` / `CacheManager` semantic-hash lookups, and so `stringArgs` / + * explain output renders the two Partial/Final nodes in a Spark plan tree distinguishably (the + * golden files under `tpcds-plan-stability/` for q44 / q67 stack a Final on top of a Partial and + * would otherwise show identical labels for both). */ case class CometWindowGroupLimitExec( override val nativeOp: Operator, @@ -120,7 +136,9 @@ case class CometWindowGroupLimitExec( override val output: Seq[Attribute], partitionSpec: Seq[Expression], orderSpec: Seq[SortOrder], + rankLikeFunction: RankLikeFunction, limit: Int, + mode: String, child: SparkPlan, override val serializedPlanOpt: SerializedPlan) extends CometUnaryExec { @@ -135,19 +153,28 @@ case class CometWindowGroupLimitExec( this.copy(child = newChild) override def stringArgs: Iterator[Any] = - Iterator(output, partitionSpec, orderSpec, limit, child) + Iterator(output, partitionSpec, orderSpec, rankLikeFunction, limit, mode, child) override def equals(obj: Any): Boolean = obj match { case other: CometWindowGroupLimitExec => this.output == other.output && this.partitionSpec == other.partitionSpec && this.orderSpec == other.orderSpec && + this.rankLikeFunction == other.rankLikeFunction && this.limit == other.limit && + this.mode == other.mode && this.child == other.child && this.serializedPlanOpt == other.serializedPlanOpt case _ => false } override def hashCode(): Int = - Objects.hashCode(output, partitionSpec, orderSpec, Integer.valueOf(limit), child) + Objects.hashCode( + output, + partitionSpec, + orderSpec, + rankLikeFunction, + Integer.valueOf(limit), + mode, + child) } diff --git a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala index 4f17f4ec98e..baca02147a7 100644 --- a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -40,11 +40,13 @@ object ShimCometWindowGroupLimit { case _: RowNumber => RankLikeFunction.RowNumber case _: Rank => RankLikeFunction.Rank case _: DenseRank => RankLikeFunction.DenseRank - case other => - throw new IllegalStateException( - s"Unexpected rank-like function in WindowGroupLimitExec: ${other.getClass.getName}") + case _ => + // Future Spark releases could add a fourth rank-like function to + // `InferWindowGroupLimit.support`. Return None so `convert` records a fallback + // reason instead of throwing, keeping a working query working across upgrades. + return None } - Some(Fields(w.partitionSpec, w.orderSpec, fn, w.limit)) + Some(Fields(w.partitionSpec, w.orderSpec, fn, w.limit, w.mode.toString)) case _ => None } diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala index 3175d30f39f..75748945a89 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -41,11 +41,13 @@ object ShimCometWindowGroupLimit { case _: RowNumber => RankLikeFunction.RowNumber case _: Rank => RankLikeFunction.Rank case _: DenseRank => RankLikeFunction.DenseRank - case other => - throw new IllegalStateException( - s"Unexpected rank-like function in WindowGroupLimitExec: ${other.getClass.getName}") + case _ => + // Future Spark releases could add a fourth rank-like function to + // `InferWindowGroupLimit.support`. Return None so `convert` records a fallback + // reason instead of throwing, keeping a working query working across upgrades. + return None } - Some(Fields(w.partitionSpec, w.orderSpec, fn, w.limit)) + Some(Fields(w.partitionSpec, w.orderSpec, fn, w.limit, w.mode.toString)) case _ => None } diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_datatypes.sql similarity index 100% rename from spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql rename to spark/src/test/resources/sql-tests/windows/window_group_limit_datatypes.sql diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_edge.sql similarity index 100% rename from spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql rename to spark/src/test/resources/sql-tests/windows/window_group_limit_edge.sql diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql similarity index 93% rename from spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql rename to spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql index 8530874138b..cf214331dd1 100644 --- a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql @@ -179,10 +179,11 @@ SELECT a, b, score FROM ( -- ================================================================================ -- FP edge cases in the ORDER BY column (NaN / +Inf / -Inf / 0 / NULL). NaN sorts -- greater than any finite value in Spark, and Arrow's row-format total-ordering --- encoding matches -- so NaN gets rank 1 under DESC. Note: -0.0 vs 0.0 is a known --- Spark-vs-Arrow-row divergence (Spark treats them equal in ORDER BY; Arrow's --- bitwise total_cmp treats them distinct), so this test avoids exercising that --- boundary directly by keeping the cutoff (rk <= 3) above the 0-values. +-- encoding matches -- so NaN gets rank 1 under DESC. -0.0 vs 0.0 is a known +-- divergence: Spark's SQLOrderingUtil ties them, Arrow's bitwise total_cmp splits +-- them. The primary test below keeps the cutoff (rk <= 3) above the 0-values so +-- it runs as-is; a second test below with `query ignore(...)` pins the divergent +-- shape directly. See docs/source/user-guide/latest/compatibility/floating-point.md. -- ================================================================================ statement @@ -207,6 +208,26 @@ SELECT part, v FROM ( FROM test_rank_fp ) t WHERE rk <= 3 ORDER BY rk, v DESC NULLS LAST +-- -0.0 vs +0.0 in the ORDER BY column: Spark ties them at rank 1 (both rows +-- survive `rk <= 1`); Comet's row encoder splits them (only -0.0 survives ASC, +-- only +0.0 survives DESC). Kept as `ignore(...)` so the suite passes today and +-- lands green once the divergence is closed. +statement +CREATE TABLE test_rank_fp_zero(part string, v double) USING parquet + +statement +INSERT INTO test_rank_fp_zero VALUES + ('p', 0.0), + ('p', -0.0), + ('p', 1.0) + +query ignore(signed-zero ORDER BY: Spark ties -0.0 with +0.0, Arrow row encoder splits them) +SELECT part, v FROM ( + SELECT part, v, + RANK() OVER (PARTITION BY part ORDER BY v ASC) AS rk + FROM test_rank_fp_zero +) t WHERE rk <= 1 ORDER BY v + -- ================================================================================ -- Decimal / bigint / integer boundary values in ORDER BY column diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_rank_dense_rank.sql similarity index 100% rename from spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql rename to spark/src/test/resources/sql-tests/windows/window_group_limit_rank_dense_rank.sql diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_row_number.sql similarity index 92% rename from spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql rename to spark/src/test/resources/sql-tests/windows/window_group_limit_row_number.sql index bdd5af5a7d2..f96057873fc 100644 --- a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_row_number.sql @@ -21,15 +21,15 @@ -- against an integer literal. Threshold is toggled via ConfigMatrix to exercise both the -- optimized (WGL-inserted) and the naive (Window + Filter) plan shapes. -- --- Comet routes the ROW_NUMBER + non-empty PARTITION BY pushdown case to DataFusion's --- PartitionedTopKExec via `CometWindowGroupLimitExec`. Every pushdown-eligible query below --- runs natively; queries where Spark's optimizer converts to a plain Limit (global top-K) --- or skips WGL pushdown still exercise the ORDER BY + Filter path without a fallback. +-- Comet routes the ROW_NUMBER + non-empty PARTITION BY pushdown case to the streaming +-- `PartitionedRankLimitExec` via `CometWindowGroupLimitExec`. Every pushdown-eligible query +-- below runs natively; queries where Spark's optimizer converts to a plain Limit (global +-- top-K) or skips WGL pushdown still exercise the ORDER BY + Filter path without a fallback. -- -- Tie-breaking note: ROW_NUMBER is non-deterministic when ORDER BY has duplicate values. --- Spark's stream-based SimpleLimitIterator preserves input order; DataFusion's heap-based --- PartitionedTopKExec does not. Every ORDER BY below therefore includes `hire_yr` as a --- secondary key so the tests exercise WGL without hitting tie-breaking non-determinism. +-- Spark's stream-based SimpleLimitIterator preserves input order, which Comet's streaming +-- operator matches. Every ORDER BY below still includes `hire_yr` as a secondary key so the +-- tests are robust against future planning changes. -- MinSparkVersion: 3.5 -- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 diff --git a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_scalar_subquery.sql similarity index 95% rename from spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql rename to spark/src/test/resources/sql-tests/windows/window_group_limit_scalar_subquery.sql index 29e3e2c38db..4482e6baa76 100644 --- a/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_scalar_subquery.sql @@ -20,8 +20,8 @@ -- the ORDER BY column collapses into the correlation (e.g. ORDER BY t2c when -- the correlation predicate is `WHERE t2c = t1c`), the resulting Window may -- have an EMPTY orderSpec but non-empty partitionSpec. Comet's serde must not --- hand DataFusion's `PartitionedTopKExec` an ordering with only partition --- keys -- that panics at execute time. +-- fail on that shape -- every row within a partition ties and every row +-- survives the limit. -- -- Spark 3.5 rejects this correlation form at analysis -- (UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.ACCESSING_OUTER_QUERY_COLUMN_IS_NOT_ALLOWED); From eef7d02058572a52cb24aa2492e9c0396f9a229a Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 21 Aug 2026 15:05:43 -0700 Subject: [PATCH 3/6] feat: support `WindowGroupLimit` --- .../latest/compatibility/operators.md | 11 ++ .../src/execution/operators/rank_limit.rs | 102 +++++++++++------- .../sql/comet/CometWindowGroupLimitExec.scala | 27 +++-- .../windows/window_group_limit_collation.sql | 66 ++++++++++++ 4 files changed, 162 insertions(+), 44 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 2c8eea4a5d3..0894856cc2c 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -74,6 +74,17 @@ incorrect result. When any single window expression in a `WindowExec` falls back `WindowGroupLimitExec` (window-based limit pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`) runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` (default: true). +**Falls back to Spark:** + +- Any `PARTITION BY` or `ORDER BY` key whose type carries a non-default `StringType` collation + (e.g. `UTF8_LCASE`). The native operator compares keys via the Arrow row encoder, which orders + by raw bytes. + +**Known incompatibilities:** + +- Signed-zero ordering (`-0.0` vs `+0.0`) diverges from Spark's `RankLimitIterator`; see + [floating-point ordering](./floating-point.md#ordering-signed-zero-00-vs-00). + ## Round-Robin Partitioning Comet's native shuffle implementation of round-robin partitioning (`df.repartition(n)`) is not compatible with diff --git a/native/core/src/execution/operators/rank_limit.rs b/native/core/src/execution/operators/rank_limit.rs index d6bd0e81e78..8c0dfac3103 100644 --- a/native/core/src/execution/operators/rank_limit.rs +++ b/native/core/src/execution/operators/rank_limit.rs @@ -295,9 +295,14 @@ struct RankLimitStream { // Per-partition streaming state, persisted across batches. prev_partition: Option, prev_order: Option, - /// Rank of the most recently seen row (0-indexed). Only meaningful when `count > 0`. + /// 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, - /// Total rows seen in the current partition (also 0-indexed cursor). + /// 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 @@ -307,10 +312,14 @@ struct RankLimitStream { } impl RankLimitStream { - fn process_batch(&mut self, batch: &RecordBatch) -> Result { + /// 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> { let num_rows = batch.num_rows(); if num_rows == 0 { - return Ok(batch.clone()); + return Ok(None); } let partition_rows = self @@ -325,26 +334,28 @@ impl RankLimitStream { 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 = None; 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 { + // 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; } - 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; } @@ -362,24 +373,22 @@ impl RankLimitStream { (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 this_rank: u64 = match self.kind { + WindowFnKind::RowNumber => self.count, + _ if ties_with_prev => self.rank, + WindowFnKind::Rank => self.count, + WindowFnKind::DenseRank if self.prev_order.is_none() => 0, + WindowFnKind::DenseRank => self.rank + 1, + }; let keep = this_rank < self.limit; mask_builder.append(keep); if keep { kept += 1; } else { + if first_dropped_at.is_none() { + first_dropped_at = Some(i); + } // `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. @@ -399,13 +408,19 @@ impl RankLimitStream { } if kept == num_rows { - return Ok(batch.clone()); + return Ok(Some(batch.clone())); } if kept == 0 { - return Ok(RecordBatch::new_empty(Arc::clone(&self.schema))); + return Ok(None); + } + // Clean prefix: kept rows are `0..kept`, dropped rows are `kept..num_rows`. `slice` is + // O(1) per column (`Arc` reslice), skipping the bitmap `true_count` and the per-column + // filter kernel that `filter_record_batch` would run. + if first_dropped_at == Some(kept) { + return Ok(Some(batch.slice(0, kept))); } let mask = BooleanArray::new(mask_builder.finish(), None); - Ok(filter_record_batch(batch, &mask)?) + Ok(Some(filter_record_batch(batch, &mask)?)) } } @@ -414,17 +429,30 @@ impl Stream for RankLimitStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { + // With no PARTITION BY, `partition_exhausted` never clears -- one virtual + // partition per DF partition. Terminate early instead of pulling the rest + // of the child stream just to drop it (matches how `LocalLimitExec` bails + // once it's satisfied its fetch). + if self.partition_exhausted && self.partition_key.is_none() { + return Poll::Ready(None); + } match self.input.poll_next_unpin(cx) { Poll::Ready(Some(Ok(batch))) => { + // Clone the `Time` metric into a local so the `ScopedTimerGuard` borrows + // the local rather than `self.baseline_metrics`. Otherwise the guard would + // hold an immutable borrow of `self` for the duration of the block and + // `self.process_batch(&batch)` (which needs `&mut self`) would fail with + // E0502. + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let processed = { - let _timer = self.baseline_metrics.elapsed_compute().timer(); + let _timer = elapsed_compute.timer(); self.process_batch(&batch) }; match processed { - // Skip fully-filtered batches so downstream never sees - // spurious empty batches between real ones. - Ok(out) if out.num_rows() == 0 => continue, - Ok(out) => { + // `process_batch` returns `None` when the batch produces no output, + // so downstream never sees a spurious empty batch between real ones. + Ok(None) => continue, + Ok(Some(out)) => { return self .baseline_metrics .record_poll(Poll::Ready(Some(Ok(out)))); diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala index 0a1a12c40ca..3c20f2f5539 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -31,7 +31,7 @@ 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.serde.QueryPlanSerde.{exprToProto, hasNonDefaultStringCollation} import org.apache.comet.shims.ShimCometWindowGroupLimit /** @@ -81,6 +81,22 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { return None } + // The streaming operator compares partition and order keys via the Arrow row encoder, + // which orders by raw bytes. A non-default `StringType` collation (e.g. `UTF8_LCASE`) + // makes Spark's comparison case-insensitive, so byte-ordering would drop rows that + // Spark considers tied. Walk nested types (StructField, ArrayType element, MapType + // key / value) via the shim helper and fall back if any key carries a collation. + val collated = (fields.partitionSpec ++ fields.orderSpec.map(_.child)) + .filter(e => hasNonDefaultStringCollation(e.dataType)) + if (collated.nonEmpty) { + withFallbackReason( + op, + collated + .map(_.sql) + .mkString("WindowGroupLimit: non-default string collation on key(s): ", ", ", "")) + 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)) @@ -103,11 +119,8 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { } override def createExec(nativeOp: Operator, op: SparkPlan): CometNativeExec = { - val fields = ShimCometWindowGroupLimit - .extract(op) - .getOrElse( - throw new IllegalStateException( - "createExec called on a non-WindowGroupLimitExec operator: " + op.nodeName)) + // `convert` above returned `Some`, so `extract` must succeed here -- `.get` is safe. + val fields = ShimCometWindowGroupLimit.extract(op).get CometWindowGroupLimitExec( nativeOp, op, @@ -174,7 +187,7 @@ case class CometWindowGroupLimitExec( partitionSpec, orderSpec, rankLikeFunction, - Integer.valueOf(limit), + limit: java.lang.Integer, mode, child) } diff --git a/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql new file mode 100644 index 00000000000..a0907ca94e4 --- /dev/null +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql @@ -0,0 +1,66 @@ +-- 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. + +-- WindowGroupLimit fallback when a partition or order key carries a non-default string +-- collation. Comet's streaming operator compares row-encoded bytes, which loses the +-- collation semantics that make Spark tie e.g. 'A' with 'a' under UTF8_LCASE. Falling +-- back to Spark keeps peer equality intact. + +-- MinSparkVersion: 4.0 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 + +statement +CREATE TABLE test_wgl_collation(grp int, s string) USING parquet + +statement +INSERT INTO test_wgl_collation VALUES (1, 'A'), (1, 'a'), (1, 'b') + +-- Case-insensitive ORDER BY key: 'A' and 'a' must tie at rank 1 (Spark keeps both). +query expect_fallback(non-default string collation) +SELECT grp, s FROM ( + SELECT grp, s, + RANK() OVER ( + PARTITION BY grp + ORDER BY CAST(s AS STRING COLLATE UTF8_LCASE) + ) AS rk + FROM test_wgl_collation +) t WHERE rk <= 1 ORDER BY grp, s + +-- Same shape with DENSE_RANK to pin both rank functions. +query expect_fallback(non-default string collation) +SELECT grp, s FROM ( + SELECT grp, s, + DENSE_RANK() OVER ( + PARTITION BY grp + ORDER BY CAST(s AS STRING COLLATE UTF8_LCASE) + ) AS rk + FROM test_wgl_collation +) t WHERE rk <= 1 ORDER BY grp, s + +-- Collated PARTITION BY key: 'A' and 'a' belong to the same partition under UTF8_LCASE, +-- so LIMIT 1 in that combined partition survives only one of them for Spark; Comet's +-- byte-order would put them in separate partitions. +query expect_fallback(non-default string collation) +SELECT s, cnt FROM ( + SELECT CAST(s AS STRING COLLATE UTF8_LCASE) AS s, + COUNT(*) OVER (PARTITION BY CAST(s AS STRING COLLATE UTF8_LCASE)) AS cnt, + ROW_NUMBER() OVER ( + PARTITION BY CAST(s AS STRING COLLATE UTF8_LCASE) + ORDER BY s + ) AS rn + FROM test_wgl_collation +) t WHERE rn <= 1 ORDER BY s From faa3564ad7e91134f837e3feb1c123e6703c69b5 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 24 Aug 2026 08:39:05 -0700 Subject: [PATCH 4/6] feat: support `WindowGroupLimit` --- docs/source/contributor-guide/native_shuffle.md | 6 ------ docs/source/user-guide/latest/operators.md | 2 +- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 1e98bc196cb..d9d675e7de3 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -49,7 +49,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition columnar output. Row-based Spark operators require JVM shuffle. 3. **Supported partitioning type**: Native shuffle supports: - - `HashPartitioning` - `RangePartitioning` - `SinglePartition` @@ -129,7 +128,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition 1. **Plan construction**: `CometNativeShuffleWriter` builds a protobuf operator tree with a `ShuffleWriter` operator at the root and `childNativeOp` as its child. `childNativeOp` takes one of two shapes: - - The child plan's `nativeOp` directly, when `CometShuffleExchangeExec`'s child is a `CometNativeExec` subtree. The upstream operators run inside the same `CometExecIterator` as the writer, with no JVM-to-native batch boundary between them. @@ -142,7 +140,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition 2. **Native execution**: A single `CometExecIterator` per partition runs the unified plan. 3. **Partitioning**: `ShuffleWriterExec` receives batches and routes to the appropriate partitioner: - - `MultiPartitionShuffleRepartitioner`: For hash/range/round-robin partitioning - `SinglePartitionShufflePartitioner`: For single partition (simpler path) @@ -150,13 +147,11 @@ Native shuffle (`CometExchange`) is selected when all of the following condition exceeds the threshold, partitions spill to temporary files. 5. **Encoding**: `ShuffleBlockWriter` encodes each partition's data as compressed Arrow IPC: - - Writes compression type header - Writes field count header - Writes compressed IPC stream 6. **Output files**: Two files are produced: - - **Data file**: Concatenated partition data - **Index file**: Array of 8-byte little-endian offsets marking partition boundaries @@ -168,7 +163,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition 1. `CometBlockStoreShuffleReader` fetches shuffle blocks via `ShuffleBlockFetcherIterator`. 2. For each block, `NativeBatchDecoderIterator`: - - Reads the 8-byte compressed length header - Reads the 8-byte field count header - Reads the compressed IPC data diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 20e52a9b1a7..3a18d86606d 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -103,7 +103,7 @@ omitted from the tables below and may be reconsidered based on demand: | Operator | Status | Notes | | ---------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `WindowExec` | ⚠️ | Runs natively and is enabled by default. A broad set of window functions is accelerated; unsupported shapes fall back to Spark. See [window function compatibility](compatibility/operators.md). | -| `WindowGroupLimitExec` | ✅ | Streaming per-partition top-K pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`. | +| `WindowGroupLimitExec` | ✅ | Streaming per-partition top-K pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`. | ## Generators and set operations From 77702dba39c8472094bf55d596f8b0365684c024 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 24 Aug 2026 15:50:11 -0700 Subject: [PATCH 5/6] feat: support `WindowGroupLimit` --- .../latest/compatibility/operators.md | 4 +-- .../sql/comet/CometWindowGroupLimitExec.scala | 12 ++++--- .../windows/window_group_limit_collation.sql | 36 ++++++++++++++----- 3 files changed, 37 insertions(+), 15 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 0894856cc2c..e039df23b76 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -77,8 +77,8 @@ runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` ( **Falls back to Spark:** - Any `PARTITION BY` or `ORDER BY` key whose type carries a non-default `StringType` collation - (e.g. `UTF8_LCASE`). The native operator compares keys via the Arrow row encoder, which orders - by raw bytes. + (e.g. `UTF8_LCASE`). The native operator detects partitions and order-key peer groups by + comparing Arrow row-encoded keys for byte equality, which splits peers that Spark ties. **Known incompatibilities:** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala index 3c20f2f5539..265fe4d8d51 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -81,11 +81,13 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { return None } - // The streaming operator compares partition and order keys via the Arrow row encoder, - // which orders by raw bytes. A non-default `StringType` collation (e.g. `UTF8_LCASE`) - // makes Spark's comparison case-insensitive, so byte-ordering would drop rows that - // Spark considers tied. Walk nested types (StructField, ArrayType element, MapType - // key / value) via the shim helper and fall back if any key carries a collation. + // The streaming operator detects partition boundaries and order-key peer groups by comparing + // Arrow row-encoded keys for byte equality, and relies on the child sort Spark injected to + // have ordered rows the same way. A non-default `StringType` collation (e.g. `UTF8_LCASE`) + // makes Spark's comparison case-insensitive, so byte equality splits a peer group Spark + // considers tied and byte ordering disagrees with the ordering the operator assumes. Walk + // nested types (StructField, ArrayType element, MapType key / value) via the shim helper and + // fall back if any key carries a collation. val collated = (fields.partitionSpec ++ fields.orderSpec.map(_.child)) .filter(e => hasNonDefaultStringCollation(e.dataType)) if (collated.nonEmpty) { diff --git a/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql index a0907ca94e4..cac9e87e501 100644 --- a/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql @@ -16,12 +16,19 @@ -- under the License. -- WindowGroupLimit fallback when a partition or order key carries a non-default string --- collation. Comet's streaming operator compares row-encoded bytes, which loses the +-- collation. Comet's streaming operator compares row-encoded bytes for equality, which loses the -- collation semantics that make Spark tie e.g. 'A' with 'a' under UTF8_LCASE. Falling -- back to Spark keeps peer equality intact. +-- +-- Unlike the other window_group_limit_* fixtures this file pins the threshold instead of running +-- the -1,1000 matrix. `spark.sql.optimizer.windowGroupLimitThreshold=-1` makes +-- `InferWindowGroupLimit` a no-op (see its `apply`), so no `WindowGroupLimit` node is planned and +-- the `expect_fallback` reason below - which only `CometWindowGroupLimitExec.convert` can emit - +-- is unreachable by construction. The plain `WindowExec` path that the -1 arm would exercise has +-- no collation guard of its own and is tracked separately. -- MinSparkVersion: 4.0 --- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 +-- Config: spark.sql.optimizer.windowGroupLimitThreshold=1000 statement CREATE TABLE test_wgl_collation(grp int, s string) USING parquet @@ -29,6 +36,13 @@ CREATE TABLE test_wgl_collation(grp int, s string) USING parquet statement INSERT INTO test_wgl_collation VALUES (1, 'A'), (1, 'a'), (1, 'b') +-- Keep the values byte-order-compatible with UTF8_LCASE order ('A' < 'a' < 'b' holds under both). +-- `QueryPlanSerde.supportedSortType` only rejects collated strings for single-column sorts, so the +-- two-column `Sort [key ASC, s ASC]` that Spark injects below `WindowGroupLimitExec` still runs on +-- Comet and sorts by raw bytes. Values where the two orders disagree (e.g. 'a' before 'B' under +-- UTF8_LCASE but after it by byte value) would feed Spark's fallback operator a wrongly ordered +-- stream and fail for a reason unrelated to what this file pins. + -- Case-insensitive ORDER BY key: 'A' and 'a' must tie at rank 1 (Spark keeps both). query expect_fallback(non-default string collation) SELECT grp, s FROM ( @@ -51,13 +65,19 @@ SELECT grp, s FROM ( FROM test_wgl_collation ) t WHERE rk <= 1 ORDER BY grp, s --- Collated PARTITION BY key: 'A' and 'a' belong to the same partition under UTF8_LCASE, --- so LIMIT 1 in that combined partition survives only one of them for Spark; Comet's --- byte-order would put them in separate partitions. +-- Collated PARTITION BY key: 'A' and 'a' belong to the same partition under UTF8_LCASE, so +-- ROW_NUMBER() <= 1 keeps only one of them for Spark, while Comet's byte-equality partition +-- detection would treat them as two partitions and keep both. +-- +-- This must stay the only window function in the query. A second window over the same collated +-- key (e.g. COUNT(*) OVER (PARTITION BY ...)) lands between the `WindowGroupLimit` and the scan +-- and forces its own Spark shuffle exchange, so no `WindowGroupLimit` node has Comet-native +-- children any more. `CometExecRule.transform` only offers an operator to its serde when every +-- child is a `CometNativeExec`, so the guard below would never run and the reason would never be +-- recorded. query expect_fallback(non-default string collation) -SELECT s, cnt FROM ( - SELECT CAST(s AS STRING COLLATE UTF8_LCASE) AS s, - COUNT(*) OVER (PARTITION BY CAST(s AS STRING COLLATE UTF8_LCASE)) AS cnt, +SELECT s, rn FROM ( + SELECT s, ROW_NUMBER() OVER ( PARTITION BY CAST(s AS STRING COLLATE UTF8_LCASE) ORDER BY s From 8d2ea9830d1b5e0898ac43f349cd835343b7501d Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 25 Aug 2026 13:22:38 -0700 Subject: [PATCH 6/6] feat: support `WindowGroupLimit` --- .../src/execution/operators/rank_limit.rs | 45 ++++++++++++++++--- 1 file changed, 39 insertions(+), 6 deletions(-) diff --git a/native/core/src/execution/operators/rank_limit.rs b/native/core/src/execution/operators/rank_limit.rs index 8c0dfac3103..0f2321c31d3 100644 --- a/native/core/src/execution/operators/rank_limit.rs +++ b/native/core/src/execution/operators/rank_limit.rs @@ -334,9 +334,11 @@ impl RankLimitStream { 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. + // Position of the first dropped row in this batch, recorded by BOTH drop paths (the + // exhausted-partition fast drop inside the loop and the rank-based drop below). 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 = None; for i in 0..num_rows { // Only a PARTITION-BY-shaped stream has partition boundaries. With no @@ -355,6 +357,13 @@ impl RankLimitStream { } if self.partition_exhausted { + // Record the first drop here too. A batch can open inside a partition the + // previous batch already exhausted, dropping these leading rows before the + // rank-based drop below ever runs. The clean-prefix shortcut keys off + // `first_dropped_at`, so leaving it unset here lets a later partition's kept + // row at position `kept` fire `batch.slice(0, kept)` and lose the kept tail + // (see the `exhausted_partition_prefix_then_fresh_partitions` test). + first_dropped_at.get_or_insert(i); mask_builder.append(false); continue; } @@ -386,9 +395,7 @@ impl RankLimitStream { if keep { kept += 1; } else { - if first_dropped_at.is_none() { - first_dropped_at = Some(i); - } + first_dropped_at.get_or_insert(i); // `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. @@ -634,4 +641,30 @@ mod tests { let out = run(vec![b0, b1], 1, WindowFnKind::RowNumber).await; assert_eq!(flatten(&out), vec![(1, 10)]); } + + /// Regression for the clean-prefix shortcut mis-firing after an exhausted-partition + /// prefix. Batch 1 opens with the tail of the partition batch 0 exhausted (dropped), then + /// crosses into two fresh partitions whose top rows must be kept. The mask is + /// `[false, true, false, true]`. If the exhausted-partition branch does not record the + /// leading drops, `first_dropped_at` stays at the rank-based drop (position 2) and equals + /// `kept` (2), so the shortcut returns `batch.slice(0, 2)` -- keeping `(1,50)` and losing + /// partition 3's top row `(3,10)`. Covered for all three window kinds; the ORDER BY values + /// are distinct so RANK, DENSE_RANK, and ROW_NUMBER all yield the same top-1-per-partition. + #[tokio::test] + async fn exhausted_partition_prefix_then_fresh_partitions() { + let b0 = batch(vec![1, 1, 1, 1], vec![10, 20, 30, 40]); + let b1 = batch(vec![1, 2, 2, 3], vec![50, 10, 20, 10]); + for kind in [ + WindowFnKind::RowNumber, + WindowFnKind::Rank, + WindowFnKind::DenseRank, + ] { + let out = run(vec![b0.clone(), b1.clone()], 1, kind).await; + assert_eq!( + flatten(&out), + vec![(1, 10), (2, 10), (3, 10)], + "kind={kind:?}" + ); + } + } }