diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d2c1a5df02..c02482519c 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -268,6 +268,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 25162332fd..bad5781bc4 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -44,6 +44,7 @@ use datafusion::functions_aggregate::min_max::min_udaf; use datafusion::functions_aggregate::percentile_cont::percentile_cont_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::{ @@ -2125,6 +2126,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 2fcfe7f25b..af1a961724 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; } } @@ -491,3 +492,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 cd971d4f92..eb899ee44f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -275,6 +275,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 52f39c59ad..9685a473ad 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 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 08df1c04a0..51279b9c54 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 e527d95f46..447b3cfa6d 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