From 8de60acdaece042e4ac42dce4d5e70f321a9c8ba Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 8 Jul 2026 13:26:38 -0700 Subject: [PATCH 01/13] feat: support `WindowGroupLimit` --- native/core/src/execution/jni_api.rs | 1 + native/core/src/execution/planner.rs | 68 +++++++ .../execution/planner/operator_registry.rs | 1 + native/proto/src/proto/operator.proto | 15 ++ .../scala/org/apache/comet/CometConf.scala | 2 + .../apache/comet/rules/CometExecRule.scala | 12 +- .../sql/comet/CometWindowGroupLimit.scala | 43 +++++ .../sql/comet/CometWindowGroupLimitExec.scala | 167 ++++++++++++++++++ .../shims/ShimCometWindowGroupLimit.scala | 40 +++++ .../shims/ShimCometWindowGroupLimit.scala | 49 +++++ .../shims/ShimCometWindowGroupLimit.scala | 50 ++++++ .../window/window_group_limit_datatypes.sql | 138 +++++++++++++++ .../window/window_group_limit_edge.sql | 126 +++++++++++++ .../window_group_limit_rank_dense_rank.sql | 105 +++++++++++ .../window/window_group_limit_row_number.sql | 164 +++++++++++++++++ 15 files changed, 978 insertions(+), 3 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimit.scala 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_dense_rank.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 9a82eb6746..cdf15f1796 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -270,6 +270,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::CsvScan(_) => "CsvScan", OpStruct::ShuffleScan(_) => "ShuffleScan", OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin", + OpStruct::WindowGroupLimit(_) => "WindowGroupLimit", } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 077df787c8..61af63977d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -46,6 +46,7 @@ use datafusion::functions_aggregate::min_max::max_udaf; use datafusion::functions_aggregate::min_max::min_udaf; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; +use datafusion::physical_plan::sorts::partitioned_topk::PartitionedTopKExec; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion::physical_plan::InputOrderMode; use datafusion::{ @@ -2186,6 +2187,73 @@ impl PhysicalPlanner { Arc::new(SparkPlan::new(spark_plan.plan_id, final_plan, vec![child])), )) } + OpStruct::WindowGroupLimit(wgl) => { + // Comet only serializes the ROW_NUMBER pushdown case with a non-empty + // PARTITION BY (see CometWindowGroupLimitExec.convert). RANK / DENSE_RANK and + // empty PARTITION BY fall back to Spark, so the proto here is guaranteed to + // have a non-empty partition_by_list. Route to DataFusion's + // `PartitionedTopKExec`, which maintains one bounded heap per distinct + // partition key and emits the top `limit` rows per partition — exactly the + // semantics of Spark's SimpleLimitIterator over a partition-sorted stream. + assert_eq!(children.len(), 1); + let (scans, shuffle_scans, child) = + self.create_plan(&children[0], inputs, partition_count)?; + let input_schema = child.schema(); + + if wgl.partition_by_list.is_empty() { + return Err(GeneralError( + "WindowGroupLimit with empty partition_by_list should have fallen back \ + to Spark; native path only supports the partitioned ROW_NUMBER top-K" + .to_string(), + )); + } + + // 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, and + // PartitionedTopKExec expects `[partition_keys, order_keys]` as a single + // LexOrdering. Materialize partition keys with SortOptions matching Spark's + // `SortOrder(_, Ascending)` (ascending, nulls_first) so the row-encoding used + // for per-partition equality inside PartitionedTopKExec matches. + let partition_prefix_len = wgl.partition_by_list.len(); + let partition_sort_options = SortOptions { + descending: false, + nulls_first: true, + }; + 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: partition_sort_options, + }); + } + 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()) + })?; + + let fetch = wgl.limit as usize; + let topk: Arc = Arc::new(PartitionedTopKExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?); + + 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 dc98a32fb0..6f6cd9a4a2 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -152,5 +152,6 @@ fn get_operator_type(spark_operator: &Operator) -> Option { OpStruct::CsvScan(_) => Some(OperatorType::CsvScan), OpStruct::ShuffleScan(_) => None, // Not yet in OperatorType enum OpStruct::BroadcastNestedLoopJoin(_) => 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 5493d921b8..7664cf51c8 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -55,6 +55,7 @@ message Operator { CsvScan csv_scan = 115; ShuffleScan shuffle_scan = 116; BroadcastNestedLoopJoin broadcast_nested_loop_join = 117; + WindowGroupLimit window_group_limit = 118; } } @@ -577,3 +578,17 @@ 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+). Comet currently supports the ROW_NUMBER pushdown case only, +// which routes to DataFusion's PartitionedTopKExec. Fields: +// partition_by_list: partition keys (must be non-empty for the ROW_NUMBER +// pushdown path — global top-K should be lowered to a +// SortExec with fetch instead). +// order_by_list: ORDER BY sort expressions. +// limit: the K in "top-K per partition". +message WindowGroupLimit { + repeated spark.spark_expression.Expr partition_by_list = 1; + repeated spark.spark_expression.Expr order_by_list = 2; + int32 limit = 3; +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 958a22cc2b..e07a1feccf 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -229,6 +229,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 bd3b92f2cc..4b035f640f 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 { @@ -70,8 +70,8 @@ object CometExecRule { /** * Fully native operators. */ - val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = - Map( + val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = { + val base: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = Map( classOf[ProjectExec] -> CometProjectExec, classOf[FilterExec] -> CometFilterExec, classOf[LocalLimitExec] -> CometLocalLimitExec, @@ -87,6 +87,12 @@ object CometExecRule { classOf[SortExec] -> CometSortExec, classOf[LocalTableScanExec] -> CometLocalTableScanExec, classOf[WindowExec] -> CometWindowExec) + // WindowGroupLimitExec exists only on Spark 3.5+; the shim returns None on 3.4. + ShimCometWindowGroupLimit.windowGroupLimitClass match { + case Some(cls) => base + (cls -> CometWindowGroupLimitExec) + case None => base + } + } /** * Sinks that have a native plan of ScanExec. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimit.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimit.scala new file mode 100644 index 0000000000..1a12a2532a --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimit.scala @@ -0,0 +1,43 @@ +/* + * 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 org.apache.spark.sql.catalyst.expressions.{Expression, SortOrder} + +/** + * Common types used by [[CometWindowGroupLimitExec]] and the per-Spark-version + * [[org.apache.comet.shims.ShimCometWindowGroupLimit]] shims. Kept out of the shim to avoid a + * sealed-trait-per-shim class-identity mismatch. + */ +object CometWindowGroupLimit { + + /** Which rank-like function drives the top-K semantics. Mirrors Spark's supported set. */ + sealed trait RankLikeKind + case object RowNumberKind extends RankLikeKind + case object RankKind extends RankLikeKind + case object DenseRankKind extends RankLikeKind + + /** Fields extracted from a Spark `WindowGroupLimitExec` (Spark 3.5+). */ + case class Fields( + partitionSpec: Seq[Expression], + orderSpec: Seq[SortOrder], + rankLikeKind: RankLikeKind, + limit: Int) +} 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 0000000000..9447785675 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -0,0 +1,167 @@ +/* + * 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 +import org.apache.comet.serde.QueryPlanSerde.exprToProto +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** + * Serde for Spark's `WindowGroupLimitExec` (Spark 3.5+, SPARK-37099). Handles the ROW_NUMBER + * pushdown case with a non-empty PARTITION BY only -- that maps directly onto DataFusion's + * `PartitionedTopKExec`. RANK / DENSE_RANK and empty PARTITION BY are rejected here and fall back + * to Spark; add a native implementation before enabling them. + * + * 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] { + import CometWindowGroupLimit._ + + 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] = { + val fields = ShimCometWindowGroupLimit.extract(op).getOrElse { + // Defensive: nativeExecs only routes here for a real WindowGroupLimitExec on Spark 3.5+. + withFallbackReason(op, "WindowGroupLimit: unexpected operator shape") + return None + } + + // The DataFusion `PartitionedTopKExec` requires a non-empty partition prefix and only + // handles ROW_NUMBER. RANK / DENSE_RANK ties and global (no-partition) top-K need either a + // custom Comet operator or extended DataFusion support and are out of scope here. + fields.rankLikeKind match { + case RowNumberKind => // supported + case RankKind => + withFallbackReason(op, "WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)") + return None + case DenseRankKind => + withFallbackReason(op, "WindowGroupLimit: DENSE_RANK not yet supported (only ROW_NUMBER)") + return None + } + if (fields.partitionSpec.isEmpty) { + withFallbackReason( + op, + "WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)") + return None + } + if (fields.limit <= 0) { + // Spark's optimizer collapses limit <= 0 to an empty LocalRelation, but guard anyway. + withFallbackReason(op, s"WindowGroupLimit: non-positive limit ${fields.limit}") + return None + } + + val childOutput = op.children.head.output + val partitionExprs = fields.partitionSpec.map(exprToProto(_, childOutput)) + val orderExprs = fields.orderSpec.map(exprToProto(_, childOutput)) + + if (partitionExprs.forall(_.isDefined) && orderExprs.forall(_.isDefined)) { + val wglBuilder = OperatorOuterClass.WindowGroupLimit + .newBuilder() + .setLimit(fields.limit) + wglBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava) + wglBuilder.addAllOrderByList(orderExprs.map(_.get).asJava) + Some(builder.setWindowGroupLimit(wglBuilder).build()) + } else { + val failing = + fields.partitionSpec.zip(partitionExprs).collect { case (e, None) => e } ++ + fields.orderSpec.zip(orderExprs).collect { case (e, None) => e } + withFallbackReason(op, failing: _*) + None + } + } + + 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`. Executes as DataFusion's + * `PartitionedTopKExec`; 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 0000000000..57d138e3b4 --- /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.CometWindowGroupLimit.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 0000000000..54ae656b13 --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -0,0 +1,49 @@ +/* + * 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.CometWindowGroupLimit.{DenseRankKind, Fields, RankKind, RowNumberKind} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.window.WindowGroupLimitExec + +/** + * 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 kind = w.rankLikeFunction match { + case _: RowNumber => RowNumberKind + case _: Rank => RankKind + case _: DenseRank => DenseRankKind + case other => + throw new IllegalStateException( + s"Unexpected rank-like function in WindowGroupLimitExec: ${other.getClass.getName}") + } + Some(Fields(w.partitionSpec, w.orderSpec, kind, 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 0000000000..55bb71c1fe --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWindowGroupLimit.scala @@ -0,0 +1,50 @@ +/* + * 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.CometWindowGroupLimit.{DenseRankKind, Fields, RankKind, RowNumberKind} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.window.WindowGroupLimitExec + +/** + * 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 kind = w.rankLikeFunction match { + case _: RowNumber => RowNumberKind + case _: Rank => RankKind + case _: DenseRank => DenseRankKind + case other => + throw new IllegalStateException( + s"Unexpected rank-like function in WindowGroupLimitExec: ${other.getClass.getName}") + } + Some(Fields(w.partitionSpec, w.orderSpec, kind, 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 0000000000..c4bdd56cca --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_datatypes.sql @@ -0,0 +1,138 @@ +-- 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. +-- ROW_NUMBER cases execute natively via `CometWindowGroupLimitExec`; the RANK and DENSE_RANK +-- cases are exercised here for coverage and remain fall-back until we grow a rank/dense-rank +-- aware native operator. +-- +-- 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 Spark's stream-based +-- SimpleLimitIterator and DataFusion's heap-based PartitionedTopKExec (see comment in +-- window_group_limit_row_number.sql). + +-- 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. RANK is not yet supported natively; falls back to Spark under +-- threshold=1000, runs via CometWindow+Filter under threshold=-1. Parity-only assertion. +query spark_answer_only +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. Parity-only (see comment above). +query spark_answer_only +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 0000000000..6122117e4e --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_edge.sql @@ -0,0 +1,126 @@ +-- 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 ROW_NUMBER. (Query kept malformed by the maintainer; skip.) +query ignore(https://github.com/apache/datafusion-comet/issues/wgl-broken-sql-placeholder) +SELECT a, b, rn FROM ( + SELECT a, b, + FROM test_wgl_all_null +) t WHERE rn <= 2 ORDER BY rn + +-- All-null order column with RANK: every row shares rank 1. Under threshold=1000 Comet +-- rejects RANK in the WGL serde and falls back; under threshold=-1 it runs natively. +query spark_answer_only +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. Parity-only. +query spark_answer_only +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. Parity-only (RANK fallback under 1000). +query spark_answer_only +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. Parity-only. +query spark_answer_only +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 + +-- SizeBasedWindowFunction sibling (percent_rank). SPARK-46941 added a guard in later Spark +-- versions to block WGL pushdown here, but Spark 3.5 still inserts WGL for the RANK arm; +-- Comet's serde rejects RANK and falls back. Assert result parity only. +query spark_answer_only +SELECT a, b FROM ( + SELECT a, b, + RANK() OVER (PARTITION BY a ORDER BY b DESC) AS rk, + PERCENT_RANK() OVER (PARTITION BY a ORDER BY b DESC) AS pr + FROM test_wgl_one +) t WHERE rk <= 2 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 0000000000..bd671c676f --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank_dense_rank.sql @@ -0,0 +1,105 @@ +-- 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. These are the two +-- rank-like functions that DataFusion's PartitionedTopKExec explicitly rejects, so Comet +-- support for these cases requires a custom native operator (or continued fallback). +-- All queries here use `spark_answer_only` because the plan varies with the threshold: at +-- threshold=-1 no WGL fires and Comet runs the whole thing natively; at threshold=1000 the +-- rule inserts a WGL that the Comet serde rejects, forcing a fallback to Spark. Both plans +-- must return the same rows, which is what we assert. + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.optimizer.windowGroupLimitThreshold=-1,1000 + +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 spark_answer_only +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 spark_answer_only +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 spark_answer_only +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 spark_answer_only +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 spark_answer_only +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 spark_answer_only +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 spark_answer_only +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 spark_answer_only +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 0000000000..58426cc060 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_row_number.sql @@ -0,0 +1,164 @@ +-- 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 + +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 via Spark's InferWindowGroupLimit rule under +-- threshold=1000, and Spark rewrites that to a LocalTableScan (which Comet does not +-- natively support). Under threshold=-1 the query runs via Window+Filter and returns +-- empty. spark_answer_only tolerates both plans. +query spark_answer_only +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 From 5e17824fb6032e05148714c39eb851154eb9066a Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 9 Jul 2026 14:34:01 -0700 Subject: [PATCH 02/13] feat: support `WindowGroupLimit` --- .../approved-plans-v1_4-spark3_5/q44/extended.txt | 8 ++++---- .../approved-plans-v1_4-spark3_5/q67/extended.txt | 4 ++-- .../approved-plans-v1_4-spark3_5/q70/extended.txt | 2 +- .../approved-plans-v1_4-spark4_0/q44/extended.txt | 8 ++++---- .../approved-plans-v2_7-spark3_5/q67a/extended.txt | 4 ++-- .../approved-plans-v2_7-spark3_5/q70a/extended.txt | 6 +++--- 6 files changed, 16 insertions(+), 16 deletions(-) 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 3ee645114d..d6a61aace8 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 @@ -9,11 +9,11 @@ TakeOrderedAndProject : : : +- Project : : : +- Filter : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometFilter @@ -35,11 +35,11 @@ TakeOrderedAndProject : : +- Project : : +- Filter : : +- Window - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometFilter 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 3c1bc08e97..5042e13876 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,11 +1,11 @@ TakeOrderedAndProject +- Filter +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] +- CometNativeColumnarToRow +- CometSort +- CometColumnarExchange - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] +- CometNativeColumnarToRow +- CometSort +- CometHashAggregate 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 f9fb64d9e4..f040933046 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 @@ -34,7 +34,7 @@ CometNativeColumnarToRow +- Project +- Filter +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] +- CometNativeColumnarToRow +- CometSort +- CometHashAggregate 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 c1b85fc86e..d375af1323 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 @@ -11,11 +11,11 @@ CometNativeColumnarToRow : : : +- Project : : : +- Filter : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometFilter @@ -38,11 +38,11 @@ CometNativeColumnarToRow : : +- Project : : +- Filter : : +- Window - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometFilter 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 8abd1fc15f..fc53f99aed 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,11 +1,11 @@ TakeOrderedAndProject +- Filter +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] +- CometNativeColumnarToRow +- CometSort +- CometColumnarExchange - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] +- CometNativeColumnarToRow +- CometSort +- CometUnion 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 2aac2972c3..dd683a9b49 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 @@ -37,7 +37,7 @@ CometNativeColumnarToRow : +- Project : +- Filter : +- Window - : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : +- CometNativeColumnarToRow : +- CometSort : +- CometHashAggregate @@ -90,7 +90,7 @@ CometNativeColumnarToRow : +- Project : +- Filter : +- Window - : +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] : +- CometNativeColumnarToRow : +- CometSort : +- CometHashAggregate @@ -143,7 +143,7 @@ CometNativeColumnarToRow +- Project +- Filter +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit is not supported] + +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] +- CometNativeColumnarToRow +- CometSort +- CometHashAggregate From c844121c76c0172b3410a40e1b84100ed59f4705 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 12:56:16 -0700 Subject: [PATCH 03/13] feat: support `WindowGroupLimit` --- native/core/src/execution/operators/mod.rs | 2 + .../src/execution/operators/rank_limit.rs | 351 +++++++++++ native/core/src/execution/planner.rs | 68 ++- native/proto/src/proto/operator.proto | 26 +- .../sql/comet/CometWindowGroupLimitExec.scala | 25 +- .../window/window_group_limit_rank.sql | 334 ++++++++++ .../q44/extended.txt | 8 +- .../q67/extended.txt | 77 ++- .../q70/extended.txt | 99 ++- .../q44/extended.txt | 8 +- .../q67a/extended.txt | 571 +++++++++--------- .../q70a/extended.txt | 287 +++++---- 12 files changed, 1286 insertions(+), 570 deletions(-) create mode 100644 native/core/src/execution/operators/rank_limit.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 6d27a17042..7adbeebbef 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -33,6 +33,8 @@ mod parquet_writer; pub use parquet_writer::{ParquetCompression, ParquetWriterExec}; mod csv_scan; pub mod projection; +mod rank_limit; +pub use rank_limit::PartitionedRankLimitExec; 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 0000000000..ff4a8cd29c --- /dev/null +++ b/native/core/src/execution/operators/rank_limit.rs @@ -0,0 +1,351 @@ +// 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 with `RANK()` semantics. +//! +//! Comet routes Spark's `WindowGroupLimitExec` with `RANK()` to this operator. +//! Unlike DataFusion's heap-based `PartitionedTopKExec` (which handles +//! `ROW_NUMBER` only in the pinned DF 54.1.0 release), this operator relies on +//! the fact that Spark's `WindowGroupLimitExec.requiredChildOrdering` guarantees +//! the child is already sorted by `[partition_keys..., order_keys...]` — so a +//! single streaming pass decides, row-by-row, whether to emit or drop, matching +//! the exact tie behavior of Spark's `RankLimitIterator`. + +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::{ArrayRef, RecordBatch, UInt32Array}; +use arrow::compute::take_record_batch; +use arrow::datatypes::SchemaRef; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion::common::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{LexOrdering, LexRequirement, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, +}; +use futures::{Stream, StreamExt}; + +/// Streaming top-K per partition operator matching Spark's +/// `WindowGroupLimitExec` with `RANK()`. +/// +/// # Preconditions +/// The child stream MUST be sorted by `[partition_keys..., order_keys...]`. +/// Spark's `WindowGroupLimitExec.requiredChildOrdering` guarantees this via +/// `EnsureRequirements`, so Comet's serde relies on the injected sort rather +/// than declaring a hard ordering requirement here that would trigger a second +/// DataFusion-side sort. +/// +/// # Semantics +/// Row-by-row, we maintain per-partition state: `rank` (0-indexed rank of the +/// last emitted row) and `count` (total rows seen in the partition). On each +/// input row: +/// * Reset state on a partition-key change. +/// * `this_rank = 0` for the first row of a partition; else `this_rank = +/// rank` if the ORDER BY key matches the previous row, else `count`. +/// * Emit iff `this_rank < limit`. +/// +/// Consequence: all rows tied at the K-th ORDER BY value ARE emitted — RANK can +/// return more than K rows when ties straddle the boundary, exactly like +/// Spark's `RankLimitIterator`. +#[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. + partition_prefix_len: usize, + /// `K` in "top-K per partition". + fetch: usize, + cache: Arc, +} + +impl PartitionedRankLimitExec { + pub fn try_new( + input: Arc, + expr: LexOrdering, + partition_prefix_len: usize, + fetch: usize, + ) -> Result { + assert!( + partition_prefix_len > 0, + "PartitionedRankLimitExec requires a non-empty partition prefix" + ); + assert!( + partition_prefix_len < expr.len(), + "PartitionedRankLimitExec requires at least one ORDER BY expression after the partition prefix" + ); + assert!(fetch > 0, "PartitionedRankLimitExec requires fetch > 0"); + let cache = Arc::new(Self::compute_properties(&input, expr.clone())?); + Ok(Self { + input, + expr, + partition_prefix_len, + fetch, + cache, + }) + } + + fn compute_properties( + input: &Arc, + sort_exprs: LexOrdering, + ) -> Result { + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.reorder(sort_exprs)?; + Ok(PlanProperties::new( + eq_properties, + input.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + + pub fn input(&self) -> &Arc { + &self.input + } + pub fn expr(&self) -> &LexOrdering { + &self.expr + } + pub fn partition_prefix_len(&self) -> usize { + self.partition_prefix_len + } + pub fn fetch(&self) -> usize { + self.fetch + } +} + +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: fetch={}, partition_prefix_len={}, expr=[{}]", + 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, + )?)) + } + + // 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> { + let req = LexRequirement::from(self.expr.clone()); + vec![Some(OrderingRequirements::new(req))] + } + + 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(); + + // Two independent RowConverters: one for the partition prefix (used + // only for equality against the previous row) and one for the ORDER BY + // suffix (used to detect a rank-changing tie boundary). + let partition_sort_fields = build_sort_fields( + &self.expr[..self.partition_prefix_len], + &schema, + )?; + let order_sort_fields = build_sort_fields( + &self.expr[self.partition_prefix_len..], + &schema, + )?; + + let partition_converter = RowConverter::new(partition_sort_fields)?; + let order_converter = RowConverter::new(order_sort_fields)?; + + let partition_exprs: Vec> = self.expr + [..self.partition_prefix_len] + .iter() + .map(|e| Arc::clone(&e.expr)) + .collect(); + let order_exprs: Vec> = self.expr + [self.partition_prefix_len..] + .iter() + .map(|e| Arc::clone(&e.expr)) + .collect(); + + Ok(Box::pin(RankLimitStream { + input, + schema, + partition_converter, + order_converter, + partition_exprs, + order_exprs, + limit: self.fetch, + prev_partition: None, + prev_order: None, + rank: 0, + count: 0, + })) + } +} + +fn build_sort_fields( + ordering: &[datafusion::physical_expr::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, + partition_converter: RowConverter, + order_converter: RowConverter, + partition_exprs: Vec>, + order_exprs: Vec>, + limit: usize, + + // 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_cols: Vec = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + let order_cols: Vec = self + .order_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + + let partition_rows = self.partition_converter.convert_columns(&partition_cols)?; + let order_rows = self.order_converter.convert_columns(&order_cols)?; + + let mut keep: Vec = Vec::with_capacity(num_rows); + for i in 0..num_rows { + let p_row = partition_rows.row(i); + let o_row = order_rows.row(i); + + let same_partition = matches!(&self.prev_partition, Some(prev) if prev.row() == p_row); + if !same_partition { + self.prev_partition = Some(p_row.owned()); + self.prev_order = None; + self.rank = 0; + self.count = 0; + } + + let this_rank: u64 = match &self.prev_order { + None => 0, + Some(prev_o) if prev_o.row() == o_row => self.rank, + Some(_) => self.count, + }; + + if (this_rank as usize) < self.limit { + keep.push(i as u32); + } + + self.rank = this_rank; + self.prev_order = Some(o_row.owned()); + self.count += 1; + } + + let indices = UInt32Array::from(keep); + Ok(take_record_batch(batch, &indices)?) + } +} + +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 61af63977d..16c43671e5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -24,6 +24,7 @@ pub mod operator_registry; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::AlignedArrowStreamReader; use crate::execution::operators::IcebergScanExec; +use crate::execution::operators::PartitionedRankLimitExec; use crate::execution::{ expressions::list_positions::ListPositionsExpr, expressions::subquery::Subquery, @@ -125,7 +126,8 @@ use datafusion_comet_proto::{ }, spark_operator::{ self, lower_window_frame_bound::LowerFrameBoundStruct, operator::OpStruct, - upper_window_frame_bound::UpperFrameBoundStruct, AggregateMode as ProtoAggregateMode, + upper_window_frame_bound::UpperFrameBoundStruct, + window_group_limit::RankType as ProtoRankType, AggregateMode as ProtoAggregateMode, BuildSide, CompressionCodec as SparkCompressionCodec, JoinType, Operator, WindowFrameType, }, spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, @@ -2188,13 +2190,17 @@ impl PhysicalPlanner { )) } OpStruct::WindowGroupLimit(wgl) => { - // Comet only serializes the ROW_NUMBER pushdown case with a non-empty - // PARTITION BY (see CometWindowGroupLimitExec.convert). RANK / DENSE_RANK and - // empty PARTITION BY fall back to Spark, so the proto here is guaranteed to - // have a non-empty partition_by_list. Route to DataFusion's - // `PartitionedTopKExec`, which maintains one bounded heap per distinct - // partition key and emits the top `limit` rows per partition — exactly the - // semantics of Spark's SimpleLimitIterator over a partition-sorted stream. + // Comet serializes the ROW_NUMBER and RANK pushdown cases with a non-empty + // PARTITION BY (see CometWindowGroupLimitExec.convert). DENSE_RANK and empty + // PARTITION BY still fall back to Spark, so the proto here is guaranteed to + // have a non-empty partition_by_list. Route: + // * ROW_NUMBER -> DataFusion's `PartitionedTopKExec`, which maintains one + // bounded heap per distinct partition key and emits the top `limit` rows + // per partition — matches Spark's SimpleLimitIterator. + // * RANK -> Comet's streaming `PartitionedRankLimitExec`, which relies on + // the Spark-injected sort of `[partition_keys, order_keys]` and retains + // all rows tied at the K-th ORDER BY value — matches Spark's + // RankLimitIterator (`hasNext = rank < limit && input.hasNext`). assert_eq!(children.len(), 1); let (scans, shuffle_scans, child) = self.create_plan(&children[0], inputs, partition_count)?; @@ -2203,17 +2209,30 @@ impl PhysicalPlanner { if wgl.partition_by_list.is_empty() { return Err(GeneralError( "WindowGroupLimit with empty partition_by_list should have fallen back \ - to Spark; native path only supports the partitioned ROW_NUMBER top-K" + to Spark; native path only supports the partitioned top-K case" + .to_string(), + )); + } + + let rank_type = ProtoRankType::try_from(wgl.rank_type).map_err(|_| { + GeneralError(format!( + "Unsupported WindowGroupLimit rank_type: {}", + wgl.rank_type + )) + })?; + if rank_type == ProtoRankType::DenseRank { + return Err(GeneralError( + "WindowGroupLimit with DENSE_RANK should have fallen back to Spark; \ + native path does not implement DENSE_RANK yet" .to_string(), )); } // 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, and - // PartitionedTopKExec expects `[partition_keys, order_keys]` as a single - // LexOrdering. Materialize partition keys with SortOptions matching Spark's - // `SortOrder(_, Ascending)` (ascending, nulls_first) so the row-encoding used - // for per-partition equality inside PartitionedTopKExec matches. + // 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 partition_prefix_len = wgl.partition_by_list.len(); let partition_sort_options = SortOptions { descending: false, @@ -2237,12 +2256,21 @@ impl PhysicalPlanner { })?; let fetch = wgl.limit as usize; - let topk: Arc = Arc::new(PartitionedTopKExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - )?); + let topk: Arc = match rank_type { + ProtoRankType::RowNumber => Arc::new(PartitionedTopKExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?), + ProtoRankType::Rank => Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?), + ProtoRankType::DenseRank => unreachable!("guarded above"), + }; Ok(( scans, diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 7664cf51c8..4f99d0948d 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -580,15 +580,31 @@ message Window { } // Top-K rows per partition group. Corresponds to Spark's WindowGroupLimitExec -// (Spark 3.5+). Comet currently supports the ROW_NUMBER pushdown case only, -// which routes to DataFusion's PartitionedTopKExec. Fields: -// partition_by_list: partition keys (must be non-empty for the ROW_NUMBER -// pushdown path — global top-K should be lowered to a -// SortExec with fetch instead). +// (Spark 3.5+). Fields: +// partition_by_list: partition keys (must be non-empty; empty PARTITION BY / +// global top-K falls back to Spark). // order_by_list: ORDER BY sort expressions. // limit: the K in "top-K per partition". +// rank_type: which rank-like function drives semantics. ROW_NUMBER +// routes to DataFusion's PartitionedTopKExec (heap-based, +// exactly K per partition). RANK routes to Comet's +// streaming PartitionedRankLimitExec, which relies on the +// Spark-injected sort of `[partition_by_list, order_by_list]` +// and retains all rows tied at the K-th ORDER BY value +// (matching Spark's RankLimitIterator). DENSE_RANK still +// falls back on the Scala side today. message WindowGroupLimit { repeated spark.spark_expression.Expr partition_by_list = 1; repeated spark.spark_expression.Expr order_by_list = 2; int32 limit = 3; + RankType rank_type = 4; + + enum RankType { + // Default zero value keeps the wire format backward compatible: older + // Scala serde that omits the field decodes as ROW_NUMBER, matching the + // pre-existing pushdown behavior. + ROW_NUMBER = 0; + RANK = 1; + DENSE_RANK = 2; + } } 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 9447785675..9d354f98f8 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 @@ -35,10 +35,11 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProto import org.apache.comet.shims.ShimCometWindowGroupLimit /** - * Serde for Spark's `WindowGroupLimitExec` (Spark 3.5+, SPARK-37099). Handles the ROW_NUMBER - * pushdown case with a non-empty PARTITION BY only -- that maps directly onto DataFusion's - * `PartitionedTopKExec`. RANK / DENSE_RANK and empty PARTITION BY are rejected here and fall back - * to Spark; add a native implementation before enabling them. + * Serde for Spark's `WindowGroupLimitExec` (Spark 3.5+, SPARK-37099). Handles the ROW_NUMBER and + * RANK pushdown cases with a non-empty PARTITION BY -- ROW_NUMBER maps onto DataFusion's + * `PartitionedTopKExec`, RANK maps onto Comet's streaming `PartitionedRankLimitExec`. DENSE_RANK + * and empty PARTITION BY are rejected here and fall back to Spark; add a native implementation + * before enabling them. * * 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 @@ -60,16 +61,13 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { return None } - // The DataFusion `PartitionedTopKExec` requires a non-empty partition prefix and only - // handles ROW_NUMBER. RANK / DENSE_RANK ties and global (no-partition) top-K need either a - // custom Comet operator or extended DataFusion support and are out of scope here. - fields.rankLikeKind match { - case RowNumberKind => // supported - case RankKind => - withFallbackReason(op, "WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)") - return None + // DENSE_RANK and empty (global) PARTITION BY are still Spark-only. RANK is native via the + // streaming operator, which relies on the Spark-injected `[partition, order]` sort. + val rankType = fields.rankLikeKind match { + case RowNumberKind => OperatorOuterClass.WindowGroupLimit.RankType.ROW_NUMBER + case RankKind => OperatorOuterClass.WindowGroupLimit.RankType.RANK case DenseRankKind => - withFallbackReason(op, "WindowGroupLimit: DENSE_RANK not yet supported (only ROW_NUMBER)") + withFallbackReason(op, "WindowGroupLimit: DENSE_RANK not yet supported") return None } if (fields.partitionSpec.isEmpty) { @@ -92,6 +90,7 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { val wglBuilder = OperatorOuterClass.WindowGroupLimit .newBuilder() .setLimit(fields.limit) + .setRankType(rankType) wglBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava) wglBuilder.addAllOrderByList(orderExprs.map(_.get).asJava) Some(builder.setWindowGroupLimit(wglBuilder).build()) 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 0000000000..9a549b2df1 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_rank.sql @@ -0,0 +1,334 @@ +-- 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 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 d6a61aace8..1f7f0e50e3 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 @@ -9,11 +9,11 @@ TakeOrderedAndProject : : : +- Project : : : +- Filter : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometFilter @@ -35,11 +35,11 @@ TakeOrderedAndProject : : +- Project : : +- Filter : : +- Window - : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometFilter 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 5042e13876..5d01271217 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: RANK not yet supported (only ROW_NUMBER)] - +- CometNativeColumnarToRow +CometNativeColumnarToRow ++- CometTakeOrderedAndProject + +- CometFilter + +- CometWindowExec + +- CometWindowGroupLimitExec +- CometSort - +- CometColumnarExchange - +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] - +- CometNativeColumnarToRow - +- 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 31 out of 37 eligible operators (83%). Final plan contains 2 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 36 out of 37 eligible operators (97%). Final plan contains 1 transitions between Spark and Comet. \ 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 f040933046..cc1d439368 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 @@ CometNativeColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange - +- HashAggregate - +- Expand - +- Project - +- BroadcastHashJoin - :- CometNativeColumnarToRow - : +- 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 - :- CometNativeColumnarToRow - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- BroadcastExchange - +- Project - +- Filter - +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] - +- CometNativeColumnarToRow - +- 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 39 out of 53 eligible operators (73%). Final plan contains 4 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 51 out of 53 eligible operators (96%). Final plan contains 1 transitions between Spark and Comet. \ 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 d375af1323..a952838468 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 @@ -11,11 +11,11 @@ CometNativeColumnarToRow : : : +- Project : : : +- Filter : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : : +- CometNativeColumnarToRow : : : +- CometSort : : : +- CometFilter @@ -38,11 +38,11 @@ CometNativeColumnarToRow : : +- Project : : +- Filter : : +- Window - : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] + : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] : : +- CometNativeColumnarToRow : : +- CometSort : : +- CometFilter 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 fc53f99aed..37a61e983d 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: RANK not yet supported (only ROW_NUMBER)] - +- CometNativeColumnarToRow +CometNativeColumnarToRow ++- CometTakeOrderedAndProject + +- CometFilter + +- CometWindowExec + +- CometWindowGroupLimitExec +- CometSort - +- CometColumnarExchange - +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] - +- CometNativeColumnarToRow - +- 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 271 out of 285 eligible operators (95%). Final plan contains 2 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 276 out of 285 eligible operators (96%). Final plan contains 1 transitions between Spark and Comet. \ 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 dd683a9b49..e5ccab39b3 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 @@ CometNativeColumnarToRow +- CometHashAggregate +- CometUnion :- CometHashAggregate - : +- CometColumnarExchange - : +- HashAggregate - : +- Project - : +- BroadcastHashJoin - : :- CometNativeColumnarToRow - : : +- 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 - : :- CometNativeColumnarToRow - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- BroadcastExchange - : +- Project - : +- Filter - : +- Window - : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] - : +- CometNativeColumnarToRow - : +- 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 - : :- CometNativeColumnarToRow - : : +- 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 - : :- CometNativeColumnarToRow - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store - : +- BroadcastExchange - : +- Project - : +- Filter - : +- Window - : +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] - : +- CometNativeColumnarToRow - : +- 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 - :- CometNativeColumnarToRow - : +- 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 - :- CometNativeColumnarToRow - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.store - +- BroadcastExchange - +- Project - +- Filter - +- Window - +- WindowGroupLimit [COMET: WindowGroupLimit: RANK not yet supported (only ROW_NUMBER)] - +- CometNativeColumnarToRow - +- 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 117 out of 156 eligible operators (75%). Final plan contains 10 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 150 out of 156 eligible operators (96%). Final plan contains 1 transitions between Spark and Comet. \ No newline at end of file From f314a1a7b55e216fadaa82b8dc74b5917ca3528e Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 13:08:07 -0700 Subject: [PATCH 04/13] feat: support `WindowGroupLimit` --- .../core/src/execution/operators/rank_limit.rs | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/native/core/src/execution/operators/rank_limit.rs b/native/core/src/execution/operators/rank_limit.rs index ff4a8cd29c..f96621f492 100644 --- a/native/core/src/execution/operators/rank_limit.rs +++ b/native/core/src/execution/operators/rank_limit.rs @@ -196,25 +196,19 @@ impl ExecutionPlan for PartitionedRankLimitExec { // Two independent RowConverters: one for the partition prefix (used // only for equality against the previous row) and one for the ORDER BY // suffix (used to detect a rank-changing tie boundary). - let partition_sort_fields = build_sort_fields( - &self.expr[..self.partition_prefix_len], - &schema, - )?; - let order_sort_fields = build_sort_fields( - &self.expr[self.partition_prefix_len..], - &schema, - )?; + let partition_sort_fields = + build_sort_fields(&self.expr[..self.partition_prefix_len], &schema)?; + let order_sort_fields = + build_sort_fields(&self.expr[self.partition_prefix_len..], &schema)?; let partition_converter = RowConverter::new(partition_sort_fields)?; let order_converter = RowConverter::new(order_sort_fields)?; - let partition_exprs: Vec> = self.expr - [..self.partition_prefix_len] + let partition_exprs: Vec> = self.expr[..self.partition_prefix_len] .iter() .map(|e| Arc::clone(&e.expr)) .collect(); - let order_exprs: Vec> = self.expr - [self.partition_prefix_len..] + let order_exprs: Vec> = self.expr[self.partition_prefix_len..] .iter() .map(|e| Arc::clone(&e.expr)) .collect(); From 612b2e7316a1c7c327009c9240c34b3f2e8b1889 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 14:34:58 -0700 Subject: [PATCH 05/13] feat: support `WindowGroupLimit` --- .../src/execution/operators/rank_limit.rs | 82 +++++++++---- native/core/src/execution/planner.rs | 70 +++++++---- .../sql/comet/CometWindowGroupLimitExec.scala | 20 ++- .../window/window_group_limit_rank.sql | 60 +++++++++ .../q44/extended.txt | 115 +++++++++--------- .../q44/extended.txt | 84 ++++++------- 6 files changed, 260 insertions(+), 171 deletions(-) diff --git a/native/core/src/execution/operators/rank_limit.rs b/native/core/src/execution/operators/rank_limit.rs index f96621f492..fba8d2ee1f 100644 --- a/native/core/src/execution/operators/rank_limit.rs +++ b/native/core/src/execution/operators/rank_limit.rs @@ -54,6 +54,12 @@ use futures::{Stream, StreamExt}; /// than declaring a hard ordering requirement here that would trigger a second /// DataFusion-side sort. /// +/// When `partition_prefix_len == 0` there is no PARTITION BY (global top-K): +/// state accumulates across all rows in the DF partition and never resets, +/// so each DF input partition produces its own local RANK top-K (matching +/// Spark's Partial WGL semantics; the Final WGL then runs on the merged +/// single-partition stream downstream). +/// /// # Semantics /// Row-by-row, we maintain per-partition state: `rank` (0-indexed rank of the /// last emitted row) and `count` (total rows seen in the partition). On each @@ -72,6 +78,7 @@ pub struct PartitionedRankLimitExec { /// 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). partition_prefix_len: usize, /// `K` in "top-K per partition". fetch: usize, @@ -85,10 +92,6 @@ impl PartitionedRankLimitExec { partition_prefix_len: usize, fetch: usize, ) -> Result { - assert!( - partition_prefix_len > 0, - "PartitionedRankLimitExec requires a non-empty partition prefix" - ); assert!( partition_prefix_len < expr.len(), "PartitionedRankLimitExec requires at least one ORDER BY expression after the partition prefix" @@ -193,26 +196,30 @@ impl ExecutionPlan for PartitionedRankLimitExec { let input = self.input.execute(partition, context)?; let schema = input.schema(); - // Two independent RowConverters: one for the partition prefix (used - // only for equality against the previous row) and one for the ORDER BY - // suffix (used to detect a rank-changing tie boundary). - let partition_sort_fields = - build_sort_fields(&self.expr[..self.partition_prefix_len], &schema)?; + // Order-by suffix drives tie detection; always required. let order_sort_fields = build_sort_fields(&self.expr[self.partition_prefix_len..], &schema)?; - - let partition_converter = RowConverter::new(partition_sort_fields)?; let order_converter = RowConverter::new(order_sort_fields)?; - - let partition_exprs: Vec> = self.expr[..self.partition_prefix_len] - .iter() - .map(|e| Arc::clone(&e.expr)) - .collect(); let order_exprs: Vec> = self.expr[self.partition_prefix_len..] .iter() .map(|e| Arc::clone(&e.expr)) .collect(); + // Partition prefix is optional: `PARTITION BY` may be empty, in which + // case state accumulates across every row in the DF partition. + let (partition_converter, partition_exprs) = if self.partition_prefix_len == 0 { + (None, Vec::new()) + } else { + let partition_sort_fields = + build_sort_fields(&self.expr[..self.partition_prefix_len], &schema)?; + let converter = RowConverter::new(partition_sort_fields)?; + let exprs: Vec> = self.expr[..self.partition_prefix_len] + .iter() + .map(|e| Arc::clone(&e.expr)) + .collect(); + (Some(converter), exprs) + }; + Ok(Box::pin(RankLimitStream { input, schema, @@ -225,6 +232,7 @@ impl ExecutionPlan for PartitionedRankLimitExec { prev_order: None, rank: 0, count: 0, + started: false, })) } } @@ -247,8 +255,10 @@ fn build_sort_fields( struct RankLimitStream { input: SendableRecordBatchStream, schema: SchemaRef, - partition_converter: RowConverter, + /// `None` when there is no PARTITION BY (global top-K per DF input partition). + partition_converter: Option, order_converter: RowConverter, + /// Empty when `partition_converter` is `None`. partition_exprs: Vec>, order_exprs: Vec>, limit: usize, @@ -261,6 +271,9 @@ struct RankLimitStream { rank: u64, /// Total rows seen in the current partition (also 0-indexed cursor). count: u64, + /// Whether any row has been observed yet. Only used when there is no + /// PARTITION BY: on the first row we still need to initialize state. + started: bool, } impl RankLimitStream { @@ -270,32 +283,47 @@ impl RankLimitStream { return Ok(batch.clone()); } - let partition_cols: Vec = self - .partition_exprs - .iter() - .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) - .collect::>()?; + let partition_rows = match &self.partition_converter { + Some(converter) => { + let partition_cols: Vec = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + Some(converter.convert_columns(&partition_cols)?) + } + None => None, + }; + let order_cols: Vec = self .order_exprs .iter() .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) .collect::>()?; - - let partition_rows = self.partition_converter.convert_columns(&partition_cols)?; let order_rows = self.order_converter.convert_columns(&order_cols)?; let mut keep: Vec = Vec::with_capacity(num_rows); for i in 0..num_rows { - let p_row = partition_rows.row(i); let o_row = order_rows.row(i); - let same_partition = matches!(&self.prev_partition, Some(prev) if prev.row() == p_row); + let same_partition = match &partition_rows { + Some(pr) => { + let p_row = pr.row(i); + matches!(&self.prev_partition, Some(prev) if prev.row() == p_row) + } + // No PARTITION BY: state accumulates across every row, resetting + // only on the very first row of the stream. + None => self.started, + }; if !same_partition { - self.prev_partition = Some(p_row.owned()); + if let Some(pr) = &partition_rows { + self.prev_partition = Some(pr.row(i).owned()); + } self.prev_order = None; self.rank = 0; self.count = 0; } + self.started = true; let this_rank: u64 = match &self.prev_order { None => 0, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 16c43671e5..1000d7b57c 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2190,30 +2190,25 @@ impl PhysicalPlanner { )) } OpStruct::WindowGroupLimit(wgl) => { - // Comet serializes the ROW_NUMBER and RANK pushdown cases with a non-empty - // PARTITION BY (see CometWindowGroupLimitExec.convert). DENSE_RANK and empty - // PARTITION BY still fall back to Spark, so the proto here is guaranteed to - // have a non-empty partition_by_list. Route: - // * ROW_NUMBER -> DataFusion's `PartitionedTopKExec`, which maintains one - // bounded heap per distinct partition key and emits the top `limit` rows - // per partition — matches Spark's SimpleLimitIterator. - // * RANK -> Comet's streaming `PartitionedRankLimitExec`, which relies on - // the Spark-injected sort of `[partition_keys, order_keys]` and retains - // all rows tied at the K-th ORDER BY value — matches Spark's - // RankLimitIterator (`hasNext = rank < limit && input.hasNext`). + // Comet serializes the ROW_NUMBER and RANK cases here (DENSE_RANK still + // falls back on the Scala side). Route: + // * ROW_NUMBER + non-empty PARTITION BY -> DataFusion's `PartitionedTopKExec` + // (heap-based per-partition top-K; matches Spark's SimpleLimitIterator). + // * 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. + // * RANK (empty or non-empty PARTITION BY) -> Comet's streaming + // `PartitionedRankLimitExec`. Relies on the Spark-injected sort of + // `[partition_keys, order_keys]` and retains all rows tied at the K-th + // ORDER BY value — matches Spark's RankLimitIterator (`hasNext = rank < + // limit && input.hasNext`). `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(); - if wgl.partition_by_list.is_empty() { - return Err(GeneralError( - "WindowGroupLimit with empty partition_by_list should have fallen back \ - to Spark; native path only supports the partitioned top-K case" - .to_string(), - )); - } - let rank_type = ProtoRankType::try_from(wgl.rank_type).map_err(|_| { GeneralError(format!( "Unsupported WindowGroupLimit rank_type: {}", @@ -2228,12 +2223,31 @@ impl PhysicalPlanner { )); } + 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. + if partition_prefix_len == 0 && rank_type == ProtoRankType::RowNumber { + let topk: Arc = + Arc::new(LocalLimitExec::new(Arc::clone(&child.native_plan), fetch)); + return Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new( + spark_plan.plan_id, + topk, + vec![Arc::clone(&child)], + )), + )); + } + // 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 partition_prefix_len = wgl.partition_by_list.len(); let partition_sort_options = SortOptions { descending: false, nulls_first: true, @@ -2255,14 +2269,16 @@ impl PhysicalPlanner { GeneralError("WindowGroupLimit produced empty LexOrdering".to_string()) })?; - let fetch = wgl.limit as usize; let topk: Arc = match rank_type { - ProtoRankType::RowNumber => Arc::new(PartitionedTopKExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - )?), + ProtoRankType::RowNumber => { + // Non-empty partition prefix (the empty case is handled above). + Arc::new(PartitionedTopKExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?) + } ProtoRankType::Rank => Arc::new(PartitionedRankLimitExec::try_new( Arc::clone(&child.native_plan), ordering, 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 9d354f98f8..f417c2a20a 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 @@ -35,11 +35,12 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProto import org.apache.comet.shims.ShimCometWindowGroupLimit /** - * Serde for Spark's `WindowGroupLimitExec` (Spark 3.5+, SPARK-37099). Handles the ROW_NUMBER and - * RANK pushdown cases with a non-empty PARTITION BY -- ROW_NUMBER maps onto DataFusion's - * `PartitionedTopKExec`, RANK maps onto Comet's streaming `PartitionedRankLimitExec`. DENSE_RANK - * and empty PARTITION BY are rejected here and fall back to Spark; add a native implementation - * before enabling them. + * Serde for Spark's `WindowGroupLimitExec` (Spark 3.5+, SPARK-37099). Handles ROW_NUMBER and RANK + * natively -- ROW_NUMBER with a non-empty PARTITION BY maps onto DataFusion's + * `PartitionedTopKExec`, ROW_NUMBER without PARTITION BY collapses to a `LocalLimitExec` over the + * Spark-sorted child, and RANK (partitioned or global) maps onto Comet's streaming + * `PartitionedRankLimitExec`. DENSE_RANK still falls back to Spark until we implement it + * natively. * * 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 @@ -61,8 +62,7 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { return None } - // DENSE_RANK and empty (global) PARTITION BY are still Spark-only. RANK is native via the - // streaming operator, which relies on the Spark-injected `[partition, order]` sort. + // DENSE_RANK is still Spark-only. ROW_NUMBER and RANK cover partitioned and global cases. val rankType = fields.rankLikeKind match { case RowNumberKind => OperatorOuterClass.WindowGroupLimit.RankType.ROW_NUMBER case RankKind => OperatorOuterClass.WindowGroupLimit.RankType.RANK @@ -70,12 +70,6 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { withFallbackReason(op, "WindowGroupLimit: DENSE_RANK not yet supported") return None } - if (fields.partitionSpec.isEmpty) { - withFallbackReason( - op, - "WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)") - return None - } if (fields.limit <= 0) { // Spark's optimizer collapses limit <= 0 to an empty LocalRelation, but guard anyway. withFallbackReason(op, s"WindowGroupLimit: non-positive limit ${fields.limit}") 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 index 9a549b2df1..8530874138 100644 --- 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 @@ -332,3 +332,63 @@ SELECT category, item, total FROM ( 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/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 1f7f0e50e3..68f5b9424f 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: empty PARTITION BY not yet supported (global top-K)] - : : : +- CometNativeColumnarToRow - : : : +- CometSort - : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] - : : : +- CometNativeColumnarToRow - : : : +- CometSort - : : : +- CometFilter - : : : : +- Subquery - : : : : +- CometNativeColumnarToRow - : : : : +- 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: empty PARTITION BY not yet supported (global top-K)] - : : +- CometNativeColumnarToRow - : : +- CometSort - : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] - : : +- CometNativeColumnarToRow - : : +- CometSort - : : +- CometFilter - : : : +- ReusedSubquery - : : +- CometHashAggregate - : : +- CometExchange - : : +- CometHashAggregate - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.store_sales - : +- BroadcastExchange - : +- CometNativeColumnarToRow - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - +- BroadcastExchange - +- CometNativeColumnarToRow +CometNativeColumnarToRow ++- CometTakeOrderedAndProject + +- CometProject + +- CometBroadcastHashJoin + :- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometSortMergeJoin + : : :- CometSort + : : : +- CometProject + : : : +- CometFilter + : : : +- CometWindowExec + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometExchange + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometFilter + : : : : +- Subquery + : : : : +- CometNativeColumnarToRow + : : : : +- 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 55 eligible operators (58%). Final plan contains 7 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 53 out of 55 eligible operators (96%). Final plan contains 2 transitions between Spark and Comet. \ 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 a952838468..04a6e7e45c 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 @@ CometNativeColumnarToRow : :- CometProject : : +- CometSortMergeJoin : : :- CometSort - : : : +- CometColumnarExchange - : : : +- Project - : : : +- Filter - : : : +- Window - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] - : : : +- CometNativeColumnarToRow - : : : +- CometSort - : : : +- CometColumnarExchange - : : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] - : : : +- CometNativeColumnarToRow - : : : +- CometSort - : : : +- CometFilter - : : : : +- Subquery - : : : : +- CometNativeColumnarToRow - : : : : +- CometHashAggregate - : : : : +- CometExchange - : : : : +- CometHashAggregate - : : : : +- CometProject - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.store_sales + : : : +- CometExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometWindowExec + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometExchange + : : : +- CometWindowGroupLimitExec + : : : +- CometSort + : : : +- CometFilter + : : : : +- Subquery + : : : : +- CometNativeColumnarToRow + : : : : +- 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: empty PARTITION BY not yet supported (global top-K)] - : : +- CometNativeColumnarToRow - : : +- CometSort - : : +- CometColumnarExchange - : : +- WindowGroupLimit [COMET: WindowGroupLimit: empty PARTITION BY not yet supported (global top-K)] - : : +- CometNativeColumnarToRow - : : +- 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 @@ CometNativeColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.item -Comet accelerated 45 out of 57 eligible operators (78%). Final plan contains 6 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 55 out of 57 eligible operators (96%). Final plan contains 2 transitions between Spark and Comet. \ No newline at end of file From 5eb5b41755be07cf6a883cb8d3ebb86f57875d69 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 16:45:51 -0700 Subject: [PATCH 06/13] feat: support `WindowGroupLimit` --- native/core/src/execution/planner.rs | 77 ++++++++++++------- native/proto/src/proto/operator.proto | 35 ++++----- .../sql/comet/CometWindowGroupLimitExec.scala | 17 ++-- 3 files changed, 74 insertions(+), 55 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 1000d7b57c..aa387ce800 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -127,8 +127,8 @@ use datafusion_comet_proto::{ spark_operator::{ self, lower_window_frame_bound::LowerFrameBoundStruct, operator::OpStruct, upper_window_frame_bound::UpperFrameBoundStruct, - window_group_limit::RankType as ProtoRankType, AggregateMode as ProtoAggregateMode, - BuildSide, CompressionCodec as SparkCompressionCodec, JoinType, Operator, WindowFrameType, + AggregateMode as ProtoAggregateMode, BuildSide, + CompressionCodec as SparkCompressionCodec, JoinType, Operator, WindowFrameType, }, spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, }; @@ -2191,7 +2191,10 @@ impl PhysicalPlanner { } OpStruct::WindowGroupLimit(wgl) => { // Comet serializes the ROW_NUMBER and RANK cases here (DENSE_RANK still - // falls back on the Scala side). Route: + // falls back on the Scala side). The `built_in_window_function` field is a + // `ScalarFunc` expression whose `func` name selects the rank semantics, + // matching how `WindowExec` encodes its own `built_in_window_function`. + // Route: // * ROW_NUMBER + non-empty PARTITION BY -> DataFusion's `PartitionedTopKExec` // (heap-based per-partition top-K; matches Spark's SimpleLimitIterator). // * ROW_NUMBER + empty PARTITION BY -> `LocalLimitExec`. Spark's WGL requires @@ -2209,19 +2212,37 @@ impl PhysicalPlanner { self.create_plan(&children[0], inputs, partition_count)?; let input_schema = child.schema(); - let rank_type = ProtoRankType::try_from(wgl.rank_type).map_err(|_| { - GeneralError(format!( - "Unsupported WindowGroupLimit rank_type: {}", - wgl.rank_type - )) - })?; - if rank_type == ProtoRankType::DenseRank { - return Err(GeneralError( - "WindowGroupLimit with DENSE_RANK should have fallen back to Spark; \ - native path does not implement DENSE_RANK yet" - .to_string(), - )); + let rank_func_name = wgl + .built_in_window_function + .as_ref() + .and_then(|e| e.expr_struct.as_ref()) + .and_then(|s| match s { + ExprStruct::ScalarFunc(f) => Some(f.func.as_str()), + _ => None, + }) + .ok_or_else(|| { + GeneralError( + "WindowGroupLimit is missing built_in_window_function (expected a \ + ScalarFunc named row_number/rank/dense_rank)" + .to_string(), + ) + })?; + match rank_func_name { + "row_number" | "rank" => {} + "dense_rank" => { + return Err(GeneralError( + "WindowGroupLimit with dense_rank should have fallen back to Spark; \ + native path does not implement DENSE_RANK yet" + .to_string(), + )); + } + other => { + return Err(GeneralError(format!( + "Unsupported WindowGroupLimit rank function: {other}" + ))); + } } + let is_row_number = rank_func_name == "row_number"; let partition_prefix_len = wgl.partition_by_list.len(); let fetch = wgl.limit as usize; @@ -2229,7 +2250,7 @@ impl PhysicalPlanner { // 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. - if partition_prefix_len == 0 && rank_type == ProtoRankType::RowNumber { + if partition_prefix_len == 0 && is_row_number { let topk: Arc = Arc::new(LocalLimitExec::new(Arc::clone(&child.native_plan), fetch)); return Ok(( @@ -2269,23 +2290,21 @@ impl PhysicalPlanner { GeneralError("WindowGroupLimit produced empty LexOrdering".to_string()) })?; - let topk: Arc = match rank_type { - ProtoRankType::RowNumber => { - // Non-empty partition prefix (the empty case is handled above). - Arc::new(PartitionedTopKExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - )?) - } - ProtoRankType::Rank => Arc::new(PartitionedRankLimitExec::try_new( + let topk: Arc = if is_row_number { + // Non-empty partition prefix (the empty case is handled above). + Arc::new(PartitionedTopKExec::try_new( Arc::clone(&child.native_plan), ordering, partition_prefix_len, fetch, - )?), - ProtoRankType::DenseRank => unreachable!("guarded above"), + )?) + } else { + Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?) }; Ok(( diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 4f99d0948d..92dc0329aa 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -581,30 +581,25 @@ message Window { // Top-K rows per partition group. Corresponds to Spark's WindowGroupLimitExec // (Spark 3.5+). Fields: -// partition_by_list: partition keys (must be non-empty; empty PARTITION BY / -// global top-K falls back to Spark). +// partition_by_list: partition keys (empty for global top-K). // order_by_list: ORDER BY sort expressions. // limit: the K in "top-K per partition". -// rank_type: which rank-like function drives semantics. ROW_NUMBER -// routes to DataFusion's PartitionedTopKExec (heap-based, -// exactly K per partition). RANK routes to Comet's -// streaming PartitionedRankLimitExec, which relies on the -// Spark-injected sort of `[partition_by_list, order_by_list]` -// and retains all rows tied at the K-th ORDER BY value -// (matching Spark's RankLimitIterator). DENSE_RANK still -// falls back on the Scala side today. +// built_in_window_function: the rank-like window function driving semantics, +// encoded the same way as `WindowExpr.built_in_window_function` +// -- a ScalarFunc expr named "row_number", "rank", or +// "dense_rank". Reusing the WindowExec encoding keeps the +// proto surface consistent with the Window operator and +// lets both sides reuse the existing scalar-function +// serde helpers. +// DENSE_RANK is still rejected on the Scala side today, so +// the native path only sees "row_number" or "rank"; empty +// PARTITION BY routes to LocalLimitExec (row_number) or +// PartitionedRankLimitExec (rank), and non-empty routes to +// DataFusion's PartitionedTopKExec (row_number) or the +// streaming operator (rank). message WindowGroupLimit { repeated spark.spark_expression.Expr partition_by_list = 1; repeated spark.spark_expression.Expr order_by_list = 2; int32 limit = 3; - RankType rank_type = 4; - - enum RankType { - // Default zero value keeps the wire format backward compatible: older - // Scala serde that omits the field decodes as ROW_NUMBER, matching the - // pre-existing pushdown behavior. - ROW_NUMBER = 0; - RANK = 1; - DENSE_RANK = 2; - } + spark.spark_expression.Expr built_in_window_function = 4; } 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 f417c2a20a..eb5ce99102 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 -import org.apache.comet.serde.QueryPlanSerde.exprToProto +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, scalarFunctionExprToProto} import org.apache.comet.shims.ShimCometWindowGroupLimit /** @@ -63,9 +63,12 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { } // DENSE_RANK is still Spark-only. ROW_NUMBER and RANK cover partitioned and global cases. - val rankType = fields.rankLikeKind match { - case RowNumberKind => OperatorOuterClass.WindowGroupLimit.RankType.ROW_NUMBER - case RankKind => OperatorOuterClass.WindowGroupLimit.RankType.RANK + // Encode the rank-like function the same way CometWindowExec does (`ScalarFunc` with a + // canonical function name) so both operators share `WindowExpr.built_in_window_function` + // semantics on the wire and the same `scalarFunctionExprToProto` helper. + val rankFuncName = fields.rankLikeKind match { + case RowNumberKind => "row_number" + case RankKind => "rank" case DenseRankKind => withFallbackReason(op, "WindowGroupLimit: DENSE_RANK not yet supported") return None @@ -79,12 +82,14 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { val childOutput = op.children.head.output val partitionExprs = fields.partitionSpec.map(exprToProto(_, childOutput)) val orderExprs = fields.orderSpec.map(exprToProto(_, childOutput)) + val rankFuncExpr = scalarFunctionExprToProto(rankFuncName) - if (partitionExprs.forall(_.isDefined) && orderExprs.forall(_.isDefined)) { + if (partitionExprs.forall(_.isDefined) && orderExprs.forall(_.isDefined) && + rankFuncExpr.isDefined) { val wglBuilder = OperatorOuterClass.WindowGroupLimit .newBuilder() .setLimit(fields.limit) - .setRankType(rankType) + .setBuiltInWindowFunction(rankFuncExpr.get) wglBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava) wglBuilder.addAllOrderByList(orderExprs.map(_.get).asJava) Some(builder.setWindowGroupLimit(wglBuilder).build()) From 3d56501fcc1a552e006d502dd183f8c93a5bfbcb Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 17:35:36 -0700 Subject: [PATCH 07/13] feat: support `WindowGroupLimit` --- native/core/src/execution/operators/mod.rs | 2 +- .../src/execution/operators/rank_limit.rs | 157 +++++++++++++----- native/core/src/execution/planner.rs | 48 ++++-- 3 files changed, 148 insertions(+), 59 deletions(-) diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 7adbeebbef..883598b860 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -34,7 +34,7 @@ pub use parquet_writer::{ParquetCompression, ParquetWriterExec}; mod csv_scan; pub mod projection; mod rank_limit; -pub use rank_limit::PartitionedRankLimitExec; +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 index fba8d2ee1f..bae08a8116 100644 --- a/native/core/src/execution/operators/rank_limit.rs +++ b/native/core/src/execution/operators/rank_limit.rs @@ -15,15 +15,19 @@ // specific language governing permissions and limitations // under the License. -//! Streaming top-K per partition operator with `RANK()` semantics. +//! Streaming top-K per partition operator for Spark's `WindowGroupLimitExec`. //! -//! Comet routes Spark's `WindowGroupLimitExec` with `RANK()` to this operator. -//! Unlike DataFusion's heap-based `PartitionedTopKExec` (which handles -//! `ROW_NUMBER` only in the pinned DF 54.1.0 release), this operator relies on -//! the fact that Spark's `WindowGroupLimitExec.requiredChildOrdering` guarantees -//! the child is already sorted by `[partition_keys..., order_keys...]` — so a -//! single streaming pass decides, row-by-row, whether to emit or drop, matching -//! the exact tie behavior of Spark's `RankLimitIterator`. +//! Comet routes both `RANK()` and the degenerate `ROW_NUMBER()` cases (empty +//! effective ORDER BY suffix) to this operator. It relies on the fact that +//! Spark's `WindowGroupLimitExec.requiredChildOrdering` guarantees the child is +//! already sorted by `[partition_keys..., order_keys...]`, so a single +//! streaming pass decides, row-by-row, whether to emit or drop. Tie behavior +//! matches Spark's `RankLimitIterator` / `SimpleLimitIterator` exactly. +//! +//! For the common non-degenerate `ROW_NUMBER()` case (non-empty PARTITION BY +//! AND non-empty ORDER BY that isn't fully covered by PARTITION BY), the +//! planner still routes to DataFusion's heap-based `PartitionedTopKExec` which +//! is O(N log K) and doesn't require sorted input. use std::fmt::Formatter; use std::pin::Pin; @@ -44,8 +48,22 @@ use datafusion::physical_plan::{ }; use futures::{Stream, StreamExt}; +/// Which rank-like window function drives the top-K semantics. +/// +/// * [`RowNumber`](Self::RowNumber): each row within a partition gets a unique +/// monotonic 1-indexed rank (0-indexed internally). No tie handling; first K +/// rows are kept. +/// * [`Rank`](Self::Rank): rows tied on ORDER BY share the same rank, then rank +/// jumps to the running count. All rows with rank <= K are kept, which can +/// emit more than K rows when ties straddle the K-th boundary. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WindowFnKind { + RowNumber, + Rank, +} + /// Streaming top-K per partition operator matching Spark's -/// `WindowGroupLimitExec` with `RANK()`. +/// `WindowGroupLimitExec`. /// /// # Preconditions /// The child stream MUST be sorted by `[partition_keys..., order_keys...]`. @@ -56,22 +74,33 @@ use futures::{Stream, StreamExt}; /// /// When `partition_prefix_len == 0` there is no PARTITION BY (global top-K): /// state accumulates across all rows in the DF partition and never resets, -/// so each DF input partition produces its own local RANK top-K (matching +/// so each DF input partition produces its own local top-K (matching /// Spark's Partial WGL semantics; the Final WGL then runs on the merged /// single-partition stream downstream). /// +/// When `partition_prefix_len == expr.len()` the ORDER BY suffix is empty +/// (either the query has no ORDER BY, or every ORDER BY column was +/// deduplicated with a PARTITION BY column by `LexOrdering::new`). In that +/// case: +/// * `Rank`: every row within a partition ties at rank 0. All rows are +/// emitted when `limit >= 1`. +/// * `RowNumber`: rank still increments per row, so the first `limit` rows +/// per partition are emitted (matching Spark's `SimpleLimitIterator`). +/// /// # Semantics /// Row-by-row, we maintain per-partition state: `rank` (0-indexed rank of the /// last emitted row) and `count` (total rows seen in the partition). On each /// input row: /// * Reset state on a partition-key change. -/// * `this_rank = 0` for the first row of a partition; else `this_rank = -/// rank` if the ORDER BY key matches the previous row, else `count`. +/// * Compute `this_rank` according to [`WindowFnKind`]: +/// * `RowNumber`: `this_rank = count` (each row's own position). +/// * `Rank`: `0` for the first row of a partition; `rank` if the ORDER BY +/// key matches the previous row (or there is no ORDER BY suffix), else +/// `count`. /// * Emit iff `this_rank < limit`. /// -/// Consequence: all rows tied at the K-th ORDER BY value ARE emitted — RANK can -/// return more than K rows when ties straddle the boundary, exactly like -/// Spark's `RankLimitIterator`. +/// Consequence for `Rank`: all rows tied at the K-th ORDER BY value ARE emitted +/// — matches Spark's `RankLimitIterator` exactly. #[derive(Debug)] pub struct PartitionedRankLimitExec { input: Arc, @@ -79,9 +108,12 @@ pub struct PartitionedRankLimitExec { 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 the ORDER BY suffix collapsed away. partition_prefix_len: usize, /// `K` in "top-K per partition". fetch: usize, + /// Which rank-like function drives the emit/skip decision. + kind: WindowFnKind, cache: Arc, } @@ -91,10 +123,11 @@ impl PartitionedRankLimitExec { expr: LexOrdering, partition_prefix_len: usize, fetch: usize, + kind: WindowFnKind, ) -> Result { assert!( - partition_prefix_len < expr.len(), - "PartitionedRankLimitExec requires at least one ORDER BY expression after the partition prefix" + partition_prefix_len <= expr.len(), + "PartitionedRankLimitExec: partition prefix must not exceed the ordering length" ); assert!(fetch > 0, "PartitionedRankLimitExec requires fetch > 0"); let cache = Arc::new(Self::compute_properties(&input, expr.clone())?); @@ -103,6 +136,7 @@ impl PartitionedRankLimitExec { expr, partition_prefix_len, fetch, + kind, cache, }) } @@ -133,6 +167,9 @@ impl PartitionedRankLimitExec { pub fn fetch(&self) -> usize { self.fetch } + pub fn kind(&self) -> WindowFnKind { + self.kind + } } impl DisplayAs for PartitionedRankLimitExec { @@ -140,8 +177,9 @@ impl DisplayAs for PartitionedRankLimitExec { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => write!( f, - "CometPartitionedRankLimitExec: fetch={}, partition_prefix_len={}, expr=[{}]", - self.fetch, self.partition_prefix_len, self.expr + "CometPartitionedRankLimitExec: kind={:?}, fetch={}, partition_prefix_len={}, \ + expr=[{}]", + self.kind, self.fetch, self.partition_prefix_len, self.expr ), DisplayFormatType::TreeRender => unimplemented!(), } @@ -171,6 +209,7 @@ impl ExecutionPlan for PartitionedRankLimitExec { self.expr.clone(), self.partition_prefix_len, self.fetch, + self.kind, )?)) } @@ -180,8 +219,12 @@ impl ExecutionPlan for PartitionedRankLimitExec { // during plan construction, though in practice Spark's Catalyst already // injects the sort above `WindowGroupLimitExec`. fn required_input_ordering(&self) -> Vec> { - let req = LexRequirement::from(self.expr.clone()); - vec![Some(OrderingRequirements::new(req))] + if self.expr.is_empty() { + vec![None] + } else { + let req = LexRequirement::from(self.expr.clone()); + vec![Some(OrderingRequirements::new(req))] + } } fn maintains_input_order(&self) -> Vec { @@ -196,14 +239,23 @@ impl ExecutionPlan for PartitionedRankLimitExec { let input = self.input.execute(partition, context)?; let schema = input.schema(); - // Order-by suffix drives tie detection; always required. - let order_sort_fields = - build_sort_fields(&self.expr[self.partition_prefix_len..], &schema)?; - let order_converter = RowConverter::new(order_sort_fields)?; - let order_exprs: Vec> = self.expr[self.partition_prefix_len..] - .iter() - .map(|e| Arc::clone(&e.expr)) - .collect(); + // Order-by suffix drives tie detection for `Rank`. When the suffix is + // empty (either the query has no ORDER BY, or every ORDER BY column + // was deduplicated with a PARTITION BY column by `LexOrdering::new`), + // we skip the order converter and treat every row within a partition + // as a tie. + let (order_converter, order_exprs) = if self.partition_prefix_len < self.expr.len() { + let order_sort_fields = + build_sort_fields(&self.expr[self.partition_prefix_len..], &schema)?; + let converter = RowConverter::new(order_sort_fields)?; + let exprs: Vec> = self.expr[self.partition_prefix_len..] + .iter() + .map(|e| Arc::clone(&e.expr)) + .collect(); + (Some(converter), exprs) + } else { + (None, Vec::new()) + }; // Partition prefix is optional: `PARTITION BY` may be empty, in which // case state accumulates across every row in the DF partition. @@ -228,6 +280,7 @@ impl ExecutionPlan for PartitionedRankLimitExec { partition_exprs, order_exprs, limit: self.fetch, + kind: self.kind, prev_partition: None, prev_order: None, rank: 0, @@ -257,11 +310,16 @@ struct RankLimitStream { schema: SchemaRef, /// `None` when there is no PARTITION BY (global top-K per DF input partition). partition_converter: Option, - order_converter: RowConverter, + /// `None` when the ORDER BY suffix is empty (fully covered by PARTITION BY + /// or absent entirely). In that case every row within a partition ties + /// under `Rank`, or ranks monotonically by position under `RowNumber`. + order_converter: Option, /// Empty when `partition_converter` is `None`. partition_exprs: Vec>, + /// Empty when `order_converter` is `None`. order_exprs: Vec>, limit: usize, + kind: WindowFnKind, // Per-partition streaming state, persisted across batches. prev_partition: Option, @@ -295,17 +353,20 @@ impl RankLimitStream { None => None, }; - let order_cols: Vec = self - .order_exprs - .iter() - .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) - .collect::>()?; - let order_rows = self.order_converter.convert_columns(&order_cols)?; + let order_rows = match &self.order_converter { + Some(converter) => { + let order_cols: Vec = self + .order_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + Some(converter.convert_columns(&order_cols)?) + } + None => None, + }; let mut keep: Vec = Vec::with_capacity(num_rows); for i in 0..num_rows { - let o_row = order_rows.row(i); - let same_partition = match &partition_rows { Some(pr) => { let p_row = pr.row(i); @@ -325,10 +386,18 @@ impl RankLimitStream { } self.started = true; - let this_rank: u64 = match &self.prev_order { - None => 0, - Some(prev_o) if prev_o.row() == o_row => self.rank, - Some(_) => self.count, + let this_rank: u64 = match self.kind { + // ROW_NUMBER: 0-indexed position within the partition. + WindowFnKind::RowNumber => self.count, + // RANK: 0 on first row of partition; else same as prior rank + // when ORDER BY key ties (or when there is no ORDER BY), + // otherwise jump to the running count. + WindowFnKind::Rank => match (&self.prev_order, &order_rows) { + (None, _) => 0, + (Some(_), None) => self.rank, + (Some(prev_o), Some(rows)) if prev_o.row() == rows.row(i) => self.rank, + (Some(_), Some(_)) => self.count, + }, }; if (this_rank as usize) < self.limit { @@ -336,7 +405,9 @@ impl RankLimitStream { } self.rank = this_rank; - self.prev_order = Some(o_row.owned()); + if let Some(rows) = &order_rows { + self.prev_order = Some(rows.row(i).owned()); + } self.count += 1; } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index aa387ce800..85686635d0 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -24,7 +24,7 @@ pub mod operator_registry; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::AlignedArrowStreamReader; use crate::execution::operators::IcebergScanExec; -use crate::execution::operators::PartitionedRankLimitExec; +use crate::execution::operators::{PartitionedRankLimitExec, WindowFnKind}; use crate::execution::{ expressions::list_positions::ListPositionsExpr, expressions::subquery::Subquery, @@ -2290,23 +2290,41 @@ impl PhysicalPlanner { GeneralError("WindowGroupLimit produced empty LexOrdering".to_string()) })?; - let topk: Arc = if is_row_number { - // Non-empty partition prefix (the empty case is handled above). - Arc::new(PartitionedTopKExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - )?) + // `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. `PartitionedTopKExec::execute` in + // DF 54.1.0 panics on an empty ORDER BY suffix, so any WGL whose effective + // ORDER BY is fully covered by PARTITION BY routes to Comet's streaming + // operator instead, which handles the tie-everything degenerate case. + let effective_order_by_len = ordering.len().saturating_sub(partition_prefix_len); + let kind = if is_row_number { + WindowFnKind::RowNumber } else { - Arc::new(PartitionedRankLimitExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - )?) + WindowFnKind::Rank }; + let topk: Arc = + if is_row_number && effective_order_by_len > 0 { + // Common fast path: heap-based per-partition top-K. + Arc::new(PartitionedTopKExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?) + } else { + // Streaming operator handles: any RANK, plus ROW_NUMBER with an + // empty effective ORDER BY suffix. + Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + kind, + )?) + }; + Ok(( scans, shuffle_scans, From 7ad467f79befa69e6c18c7a2cf69ed2a70474492 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 18:13:54 -0700 Subject: [PATCH 08/13] feat: support `WindowGroupLimit` --- .../window_group_limit_scalar_subquery.sql | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.sql 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 0000000000..7ba5c7cf85 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/window/window_group_limit_scalar_subquery.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. + +-- 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. + +-- MinSparkVersion: 3.5 + +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 From acf57e6d2941d014d844e98cc9ea46bba4898ee6 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 23 Jul 2026 18:19:47 -0700 Subject: [PATCH 09/13] feat: support `WindowGroupLimit` --- native/core/src/execution/planner.rs | 44 +++++++++++++--------------- 1 file changed, 21 insertions(+), 23 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 85686635d0..69a8f70992 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -126,9 +126,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, + upper_window_frame_bound::UpperFrameBoundStruct, AggregateMode as ProtoAggregateMode, + BuildSide, CompressionCodec as SparkCompressionCodec, JoinType, Operator, WindowFrameType, }, spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, }; @@ -2304,26 +2303,25 @@ impl PhysicalPlanner { WindowFnKind::Rank }; - let topk: Arc = - if is_row_number && effective_order_by_len > 0 { - // Common fast path: heap-based per-partition top-K. - Arc::new(PartitionedTopKExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - )?) - } else { - // Streaming operator handles: any RANK, plus ROW_NUMBER with an - // empty effective ORDER BY suffix. - Arc::new(PartitionedRankLimitExec::try_new( - Arc::clone(&child.native_plan), - ordering, - partition_prefix_len, - fetch, - kind, - )?) - }; + let topk: Arc = if is_row_number && effective_order_by_len > 0 { + // Common fast path: heap-based per-partition top-K. + Arc::new(PartitionedTopKExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + )?) + } else { + // Streaming operator handles: any RANK, plus ROW_NUMBER with an + // empty effective ORDER BY suffix. + Arc::new(PartitionedRankLimitExec::try_new( + Arc::clone(&child.native_plan), + ordering, + partition_prefix_len, + fetch, + kind, + )?) + }; Ok(( scans, From 95dd96e887d86315734af12e90f1849e760e3146 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 24 Jul 2026 07:55:41 -0700 Subject: [PATCH 10/13] feat: support `WindowGroupLimit` --- .../window/window_group_limit_scalar_subquery.sql | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 index 7ba5c7cf85..29e3e2c38d 100644 --- 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 @@ -22,8 +22,12 @@ -- 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: 3.5 +-- MinSparkVersion: 4.0 statement CREATE TABLE t1(t1a string, t1b smallint, t1c int, t1d int) USING parquet From 035a53cec06566a4f120bd5968f7149e7ae2c0fe Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 24 Jul 2026 10:23:36 -0700 Subject: [PATCH 11/13] feat: support `WindowGroupLimit` --- dev/diffs/3.5.8.diff | 75 ++++++++++++++++---------------------------- dev/diffs/4.0.2.diff | 75 ++++++++++++++++---------------------------- dev/diffs/4.1.2.diff | 75 ++++++++++++++++---------------------------- 3 files changed, 81 insertions(+), 144 deletions(-) diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 428e964bad..4dceccf604 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -644,54 +644,6 @@ index 93275487f29..a5208b8d54b 100644 }.flatten assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L)))) } -diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala -new file mode 100644 -index 00000000000..a42533c7c84 ---- /dev/null -+++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala -@@ -0,0 +1,42 @@ -+/* -+ * 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 -+ -+import org.scalactic.source.Position -+import org.scalatest.Tag -+ -+import org.apache.spark.sql.test.SQLTestUtils -+ -+/** -+ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`). -+ */ -+case class IgnoreComet(reason: String) extends Tag("DisableComet") -+ -+/** -+ * Helper trait that disables Comet for all tests regardless of default config values. -+ */ -+trait IgnoreCometSuite extends SQLTestUtils { -+ override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit -+ pos: Position): Unit = { -+ if (isCometEnabled) { -+ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun) -+ } else { -+ super.test(testName, testTags: _*)(testFun) -+ } -+ } -+} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala index 7af826583bd..3c3def1eb67 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala @@ -1426,6 +1378,33 @@ 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..242ed80d4ee 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,13 @@ 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) ++ // Comet may replace Spark's `WindowGroupLimitExec` with `CometWindowGroupLimitExec`; ++ // count both so the assertion measures the redundancy-elimination rule's effect rather ++ // than accelerator coverage. ++ 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.2.diff b/dev/diffs/4.0.2.diff index 8ee55fea1d..1b9c9acb57 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -791,54 +791,6 @@ index 9c529d14221..ab2850b5d68 100644 }.flatten assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L)))) } -diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala -new file mode 100644 -index 00000000000..4b31bea33de ---- /dev/null -+++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala -@@ -0,0 +1,42 @@ -+/* -+ * 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 -+ -+import org.scalactic.source.Position -+import org.scalatest.Tag -+ -+import org.apache.spark.sql.test.SQLTestUtils -+ -+/** -+ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`). -+ */ -+case class IgnoreComet(reason: String) extends Tag("DisableComet") -+ -+/** -+ * Helper trait that disables Comet for all tests regardless of default config values. -+ */ -+trait IgnoreCometSuite extends SQLTestUtils { -+ override protected def test(testName: String, testTags: Tag*)(testFun: => Any) -+ (implicit pos: Position): Unit = { -+ if (isCometEnabled) { -+ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun) -+ } else { -+ super.test(testName, testTags: _*)(testFun) -+ } -+ } -+} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala index 53e47f428c3..a55d8f0c161 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala @@ -1828,6 +1780,33 @@ 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..242ed80d4ee 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,13 @@ 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) ++ // Comet may replace Spark's `WindowGroupLimitExec` with `CometWindowGroupLimitExec`; ++ // count both so the assertion measures the redundancy-elimination rule's effect rather ++ // than accelerator coverage. ++ 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.2.diff b/dev/diffs/4.1.2.diff index 599e2cd077..21ba207d34 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -862,54 +862,6 @@ index 95e86fe4311..fb2b6363af6 100644 }.flatten assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L)))) } -diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala -new file mode 100644 -index 00000000000..4b31bea33de ---- /dev/null -+++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala -@@ -0,0 +1,42 @@ -+/* -+ * 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 -+ -+import org.scalactic.source.Position -+import org.scalatest.Tag -+ -+import org.apache.spark.sql.test.SQLTestUtils -+ -+/** -+ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`). -+ */ -+case class IgnoreComet(reason: String) extends Tag("DisableComet") -+ -+/** -+ * Helper trait that disables Comet for all tests regardless of default config values. -+ */ -+trait IgnoreCometSuite extends SQLTestUtils { -+ override protected def test(testName: String, testTags: Tag*)(testFun: => Any) -+ (implicit pos: Position): Unit = { -+ if (isCometEnabled) { -+ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun) -+ } else { -+ super.test(testName, testTags: _*)(testFun) -+ } -+ } -+} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala index 53e47f428c3..a55d8f0c161 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala @@ -1951,6 +1903,33 @@ 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..242ed80d4ee 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,13 @@ 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) ++ // Comet may replace Spark's `WindowGroupLimitExec` with `CometWindowGroupLimitExec`; ++ // count both so the assertion measures the redundancy-elimination rule's effect rather ++ // than accelerator coverage. ++ 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 From f3c2878ef1985f4193345d65792a78eb397b0cef Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 25 Jul 2026 11:12:03 -0700 Subject: [PATCH 12/13] feat: support `WindowGroupLimit` --- dev/diffs/3.5.8.diff | 75 ++++++++++++++++++++++++++++---------------- dev/diffs/4.0.2.diff | 75 ++++++++++++++++++++++++++++---------------- dev/diffs/4.1.2.diff | 75 ++++++++++++++++++++++++++++---------------- 3 files changed, 144 insertions(+), 81 deletions(-) diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 4dceccf604..428e964bad 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -644,6 +644,54 @@ index 93275487f29..a5208b8d54b 100644 }.flatten assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L)))) } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala +new file mode 100644 +index 00000000000..a42533c7c84 +--- /dev/null ++++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala +@@ -0,0 +1,42 @@ ++/* ++ * 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 ++ ++import org.scalactic.source.Position ++import org.scalatest.Tag ++ ++import org.apache.spark.sql.test.SQLTestUtils ++ ++/** ++ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`). ++ */ ++case class IgnoreComet(reason: String) extends Tag("DisableComet") ++ ++/** ++ * Helper trait that disables Comet for all tests regardless of default config values. ++ */ ++trait IgnoreCometSuite extends SQLTestUtils { ++ override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit ++ pos: Position): Unit = { ++ if (isCometEnabled) { ++ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun) ++ } else { ++ super.test(testName, testTags: _*)(testFun) ++ } ++ } ++} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala index 7af826583bd..3c3def1eb67 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala @@ -1378,33 +1426,6 @@ 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..242ed80d4ee 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,13 @@ 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) -+ // Comet may replace Spark's `WindowGroupLimitExec` with `CometWindowGroupLimitExec`; -+ // count both so the assertion measures the redundancy-elimination rule's effect rather -+ // than accelerator coverage. -+ 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.2.diff b/dev/diffs/4.0.2.diff index 1b9c9acb57..8ee55fea1d 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -791,6 +791,54 @@ index 9c529d14221..ab2850b5d68 100644 }.flatten assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L)))) } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala +new file mode 100644 +index 00000000000..4b31bea33de +--- /dev/null ++++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala +@@ -0,0 +1,42 @@ ++/* ++ * 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 ++ ++import org.scalactic.source.Position ++import org.scalatest.Tag ++ ++import org.apache.spark.sql.test.SQLTestUtils ++ ++/** ++ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`). ++ */ ++case class IgnoreComet(reason: String) extends Tag("DisableComet") ++ ++/** ++ * Helper trait that disables Comet for all tests regardless of default config values. ++ */ ++trait IgnoreCometSuite extends SQLTestUtils { ++ override protected def test(testName: String, testTags: Tag*)(testFun: => Any) ++ (implicit pos: Position): Unit = { ++ if (isCometEnabled) { ++ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun) ++ } else { ++ super.test(testName, testTags: _*)(testFun) ++ } ++ } ++} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala index 53e47f428c3..a55d8f0c161 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala @@ -1780,33 +1828,6 @@ 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..242ed80d4ee 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,13 @@ 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) -+ // Comet may replace Spark's `WindowGroupLimitExec` with `CometWindowGroupLimitExec`; -+ // count both so the assertion measures the redundancy-elimination rule's effect rather -+ // than accelerator coverage. -+ 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.2.diff b/dev/diffs/4.1.2.diff index 21ba207d34..599e2cd077 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -862,6 +862,54 @@ index 95e86fe4311..fb2b6363af6 100644 }.flatten assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L)))) } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala +new file mode 100644 +index 00000000000..4b31bea33de +--- /dev/null ++++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala +@@ -0,0 +1,42 @@ ++/* ++ * 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 ++ ++import org.scalactic.source.Position ++import org.scalatest.Tag ++ ++import org.apache.spark.sql.test.SQLTestUtils ++ ++/** ++ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`). ++ */ ++case class IgnoreComet(reason: String) extends Tag("DisableComet") ++ ++/** ++ * Helper trait that disables Comet for all tests regardless of default config values. ++ */ ++trait IgnoreCometSuite extends SQLTestUtils { ++ override protected def test(testName: String, testTags: Tag*)(testFun: => Any) ++ (implicit pos: Position): Unit = { ++ if (isCometEnabled) { ++ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun) ++ } else { ++ super.test(testName, testTags: _*)(testFun) ++ } ++ } ++} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala index 53e47f428c3..a55d8f0c161 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala @@ -1903,33 +1951,6 @@ 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..242ed80d4ee 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,13 @@ 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) -+ // Comet may replace Spark's `WindowGroupLimitExec` with `CometWindowGroupLimitExec`; -+ // count both so the assertion measures the redundancy-elimination rule's effect rather -+ // than accelerator coverage. -+ 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 From d781362f4ce284d5e87997be385fd68d31860b30 Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 25 Jul 2026 11:37:16 -0700 Subject: [PATCH 13/13] feat: support `WindowGroupLimit` --- dev/diffs/3.5.8.diff | 24 ++++++++++++++++++++++++ dev/diffs/4.0.2.diff | 24 ++++++++++++++++++++++++ dev/diffs/4.1.2.diff | 24 ++++++++++++++++++++++++ 3 files changed, 72 insertions(+) diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 428e964bad..235028481d 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -1426,6 +1426,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.2.diff b/dev/diffs/4.0.2.diff index 8ee55fea1d..661845c278 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -1828,6 +1828,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.2.diff b/dev/diffs/4.1.2.diff index 599e2cd077..8be5737cdf 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -1951,6 +1951,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