You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/27 03:56:38 UTC
(datafusion-comet) branch main updated: feat: add opt-in local TopK fusion for native Parquet scans (#6181)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 74fdec5bcd feat: add opt-in local TopK fusion for native Parquet scans (#6181)
74fdec5bcd is described below
commit 74fdec5bcdb95f209104a2b698712f668aea4ae2
Author: Ping Zhang <pi...@gmail.com>
AuthorDate: Sun Sep 27 01:58:46 2026 +0000
feat: add opt-in local TopK fusion for native Parquet scans (#6181)
* feat: add opt-in local TopK fusion for native Parquet scans
* test: cover local TopK fusion boundaries and reversion
* perf: benchmark and document opt-in TopK fusion
* fix: preserve native scan output schema for empty buckets
* fix: retain file paths in fused native execution
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
.../contributor-guide/adding_a_new_operator.md | 15 +
docs/source/user-guide/latest/tuning/operators.md | 18 +
native/core/src/execution/planner.rs | 20 +-
.../execution/planner/tests/empty_native_scan.rs | 173 ++++++++++
.../main/scala/org/apache/comet/CometConf.scala | 11 +
.../org/apache/comet/rules/CometExecRule.scala | 3 +
.../RevertNativeForTransitionHeavyStages.scala | 24 +-
.../apache/spark/sql/comet/CometExecUtils.scala | 31 +-
.../spark/sql/comet/CometLocalTopKExec.scala | 103 ++++++
.../sql/comet/CometTakeOrderedAndProjectExec.scala | 33 +-
.../org/apache/spark/sql/comet/operators.scala | 31 +-
.../org/apache/comet/exec/CometTopKSuite.scala | 374 +++++++++++++++++++++
...RevertNativeForTransitionHeavyStagesSuite.scala | 45 ++-
.../spark/sql/benchmark/CometTopKBenchmark.scala | 206 ++++++++++++
.../shuffle/CometNativeShuffleInputRDDSuite.scala | 8 +-
17 files changed, 1050 insertions(+), 47 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index ea8a23d204..ca2b3e38cb 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -525,6 +525,7 @@ jobs:
org.apache.comet.exec.CometAggregateSuite
org.apache.comet.exec.CometExec3_4PlusSuite
org.apache.comet.exec.CometExecSuite
+ org.apache.comet.exec.CometTopKSuite
org.apache.comet.exec.CometEmptyRelationExecSuite
org.apache.comet.exec.CometInMemoryCacheSuite
org.apache.comet.exec.CometInMemoryCacheKryoSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 3a351d1e6f..74fa0128bf 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -173,6 +173,7 @@ jobs:
org.apache.comet.exec.CometAggregateSuite
org.apache.comet.exec.CometExec3_4PlusSuite
org.apache.comet.exec.CometExecSuite
+ org.apache.comet.exec.CometTopKSuite
org.apache.comet.exec.CometEmptyRelationExecSuite
org.apache.comet.exec.CometInMemoryCacheSuite
org.apache.comet.exec.CometInMemoryCacheKryoSuite
diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md
index 6c3c634e98..c0a20f40aa 100644
--- a/docs/source/contributor-guide/adding_a_new_operator.md
+++ b/docs/source/contributor-guide/adding_a_new_operator.md
@@ -70,6 +70,21 @@ in `CometExecRule` rather than using `CometOperatorSerde`, because they don't ne
Examples: `CometBroadcastExchangeExec`, `CometShuffleExchangeExec`
+#### Local TopK inside a sink
+
+`CometTakeOrderedAndProjectExec` owns local candidate selection, the optional shuffle, and final
+selection with offset and projection. When `spark.comet.exec.topK.fusion.enabled` is enabled for an
+eligible native Parquet scan, its conversion inserts `CometLocalTopKExec` before native blocks are
+serialized. The local node extends `CometUnaryExec` and serializes a bounded sort above the native
+scan, so they execute in one block. It preserves the scan's output and partitioning and advertises
+the local sort order.
+
+The inserted node keeps the original Spark TopK for plan bookkeeping, but only the outer TopK owns
+Spark's offset and projection. Both transition reversion and aggregate-buffer restoration therefore
+remove the inserted local node when restoring Spark execution. Restoring its `originalPlan` would
+apply the global TopK twice. A single input partition uses a final native limit and projection;
+multiple input partitions still require a final TopK after the shuffle.
+
### Choosing the Right Operator Type
When adding a new operator, choose based on these criteria:
diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md
index 843fd3b708..9dffa080ee 100644
--- a/docs/source/user-guide/latest/tuning/operators.md
+++ b/docs/source/user-guide/latest/tuning/operators.md
@@ -99,6 +99,24 @@ SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_thresh
These settings only tune eligible plans. Unsupported accumulators and modes remain
disabled even when configuration overrides are enabled.
+## Local TopK Fusion
+
+Set `spark.comet.exec.topK.fusion.enabled=true` to run an eligible local TopK in the same native
+execution as its Parquet scan. This experimental optimization is disabled by default. It currently
+supports a direct native Parquet scan ordered by one signed integer column (`TINYINT`, `SMALLINT`,
+`INT`, or `BIGINT`). Both sort directions and null orderings are supported. Other inputs use the
+existing TopK execution path.
+
+Each scan partition keeps enough candidates for both `LIMIT` and `OFFSET`. With multiple partitions,
+Comet shuffles those candidates and performs the final TopK. With one partition, it reuses the local
+ordering without building a second heap. The final stage applies the offset and output projection.
+
+Fusion reduces the work of passing scan batches between native execution blocks. It still reads all
+input rows and does not enable TopK reader pruning. It can also reduce overlap between scan decoding
+and TopK processing, so some workloads may run slower. Compare enabled and disabled runs with your
+data layout, payload width, limit, and partition count before enabling it. The
+`CometTopKBenchmark` microbenchmark covers these cases with ascending, descending, and random layouts.
+
## Optimizing Sorting on Floating-Point Values
Comet normalizes NaN payloads and signed zeros in scalar `FLOAT` and `DOUBLE` ordering keys, so `ORDER BY`, window
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index b44420f47f..80f8d890a3 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -94,6 +94,7 @@ use datafusion_comet_spark_expr::{
EvalMode, SparkArraysZipFunc, SparkBloomFilterVersion, SparkListAgg, SparkPercentile,
SumInteger, ToCsv,
};
+use datafusion_datasource::TableSchema;
use iceberg::expr::Bind;
use crate::execution::operators::ExecutionError::GeneralError;
@@ -1635,7 +1636,22 @@ impl PhysicalPlanner {
// Check if this partition has any files (bucketed scan with bucket pruning may have empty partitions)
if partition_files.partitioned_file.is_empty() {
- let empty_exec = Arc::new(EmptyExec::new(required_schema));
+ // Match the nonempty scan's projected table schema, including partition
+ // and constant metadata columns. Fused operators bind against this schema
+ // even when the bucket has no rows.
+ let partition_fields: Vec<FieldRef> = partition_schema
+ .fields()
+ .iter()
+ .map(|f| {
+ Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable()))
+ })
+ .collect();
+ let table_schema = TableSchema::builder(data_schema)
+ .with_table_partition_cols(partition_fields)
+ .build();
+ let output_schema =
+ Arc::new(table_schema.table_schema().project(&projection_vector)?);
+ let empty_exec = Arc::new(EmptyExec::new(output_schema));
return Ok((
vec![],
vec![],
@@ -5101,6 +5117,8 @@ fn needs_fields_coercion(sig: &TypeSignature) -> bool {
#[cfg(test)]
mod tests {
+ mod empty_native_scan;
+
use futures::{poll, StreamExt};
use std::{
sync::{
diff --git a/native/core/src/execution/planner/tests/empty_native_scan.rs b/native/core/src/execution/planner/tests/empty_native_scan.rs
new file mode 100644
index 0000000000..3165b8fc0a
--- /dev/null
+++ b/native/core/src/execution/planner/tests/empty_native_scan.rs
@@ -0,0 +1,173 @@
+// 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.
+
+use std::collections::HashMap;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion_comet_proto::spark_expression::{
+ data_type::DataTypeId, expr::ExprStruct, literal::Value, DataType as SparkDataType, Expr,
+ Literal,
+};
+use datafusion_comet_proto::spark_operator::{
+ operator::OpStruct, NativeScan, NativeScanCommon, Operator, SparkFilePartition,
+ SparkPartitionedFile, SparkStructField,
+};
+
+use crate::execution::planner::PhysicalPlanner;
+
+fn spark_field(
+ name: &str,
+ data_type: DataTypeId,
+ nullable: bool,
+ field_id: &str,
+) -> SparkStructField {
+ SparkStructField {
+ name: name.to_string(),
+ data_type: Some(SparkDataType {
+ type_id: data_type as i32,
+ type_info: None,
+ }),
+ nullable,
+ metadata: HashMap::from([("parquet.field.id".to_string(), field_id.to_string())]),
+ }
+}
+
+fn common(required_data_indices: &[usize], projection: Vec<i64>) -> NativeScanCommon {
+ let data_schema = vec![
+ spark_field("id", DataTypeId::Int32, false, "1"),
+ spark_field("unused", DataTypeId::String, true, "2"),
+ spark_field("payload", DataTypeId::Int64, true, "3"),
+ ];
+ NativeScanCommon {
+ required_schema: required_data_indices
+ .iter()
+ .map(|&index| data_schema[index].clone())
+ .collect(),
+ data_schema,
+ // Constant file metadata is appended alongside partition values by the JVM scan.
+ partition_schema: vec![
+ spark_field("part", DataTypeId::Int32, false, "4"),
+ spark_field("_metadata.file_size", DataTypeId::Int64, false, "5"),
+ ],
+ projection_vector: projection,
+ session_timezone: "UTC".to_string(),
+ ..Default::default()
+ }
+}
+
+fn literal(data_type: DataTypeId, value: Value) -> Expr {
+ Expr {
+ expr_struct: Some(ExprStruct::Literal(Literal {
+ datatype: Some(SparkDataType {
+ type_id: data_type as i32,
+ type_info: None,
+ }),
+ value: Some(value),
+ is_null: false,
+ })),
+ ..Default::default()
+ }
+}
+
+fn plan_schema(common: &NativeScanCommon, empty: bool) -> SchemaRef {
+ let files = if empty {
+ vec![]
+ } else {
+ // Planning does not read the file. Only its presence selects the nonempty path.
+ vec![SparkPartitionedFile {
+ file_path: "file:///tmp/comet-empty-native-scan-schema.parquet".to_string(),
+ start: 0,
+ length: 16,
+ file_size: 16,
+ partition_values: vec![
+ literal(DataTypeId::Int32, Value::IntVal(7)),
+ literal(DataTypeId::Int64, Value::LongVal(16)),
+ ],
+ }]
+ };
+ let operator = Operator {
+ op_struct: Some(OpStruct::NativeScan(NativeScan {
+ common: Some(common.clone()),
+ file_partition: Some(SparkFilePartition {
+ partitioned_file: files,
+ }),
+ ..Default::default()
+ })),
+ ..Default::default()
+ };
+ let (_, _, plan) = PhysicalPlanner::default()
+ .create_plan(&operator, &mut vec![], 1)
+ .unwrap();
+ plan.schema()
+}
+
+fn assert_schema_matches_nonempty(common: NativeScanCommon, expected: Vec<Field>) {
+ let nonempty_schema = plan_schema(&common, false);
+ let empty_schema = plan_schema(&common, true);
+ assert_eq!(nonempty_schema.as_ref(), &Schema::new(expected));
+ assert_eq!(empty_schema, nonempty_schema);
+}
+
+fn id_field() -> Field {
+ Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
+ "parquet.field.id".to_string(),
+ "1".to_string(),
+ )]))
+}
+
+fn payload_field() -> Field {
+ Field::new("payload", DataType::Int64, true).with_metadata(HashMap::from([(
+ "parquet.field.id".to_string(),
+ "3".to_string(),
+ )]))
+}
+
+#[test]
+fn empty_scan_preserves_reordered_and_pruned_data_columns() {
+ assert_schema_matches_nonempty(
+ common(&[0, 2], vec![2, 0]),
+ vec![payload_field(), id_field()],
+ );
+}
+
+#[test]
+fn empty_scan_preserves_partition_and_constant_metadata_projection() {
+ // Match the nonempty path: data-field metadata is preserved, while partition
+ // and constant metadata fields retain their names, types and nullability only.
+ assert_schema_matches_nonempty(
+ common(&[0, 2], vec![3, 2, 4, 0]),
+ vec![
+ Field::new("part", DataType::Int32, false),
+ payload_field(),
+ Field::new("_metadata.file_size", DataType::Int64, false),
+ id_field(),
+ ],
+ );
+}
+
+#[test]
+fn empty_scan_preserves_partition_only_projection() {
+ assert_schema_matches_nonempty(
+ common(&[], vec![3]),
+ vec![Field::new("part", DataType::Int32, false)],
+ );
+}
+
+#[test]
+fn empty_scan_preserves_zero_column_projection() {
+ assert_schema_matches_nonempty(common(&[], vec![]), vec![]);
+}
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 53daeb5710..4193738beb 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -447,6 +447,17 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_EXEC_TOPK_FUSION_ENABLED: ConfigEntry[Boolean] =
+ conf(s"$COMET_EXEC_CONFIG_PREFIX.topK.fusion.enabled")
+ .category(CATEGORY_EXEC)
+ .doc(
+ "Experimental opt-in: run an eligible local TopK in the same native execution as its " +
+ "Parquet scan. Supports one direct signed integer sort key. This changes the local " +
+ "execution pipeline and can reduce scan/TopK overlap, so it may be slower for some " +
+ "workloads.")
+ .booleanConf
+ .createWithDefault(false)
+
val COMET_SCALA_UDF_CODEGEN_ENABLED: ConfigEntry[Boolean] =
conf("spark.comet.exec.scalaUDF.codegen.enabled")
.category(CATEGORY_EXEC)
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 19c88a6a9d..e97ab50dea 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -261,6 +261,9 @@ case class CometExecRule(session: SparkSession)
other
} else {
other match {
+ // Only the outer TopK owns Spark's original offset and projection. If a future
+ // fused input can contain a restored aggregate, remove its inserted local node.
+ case _: CometLocalTopKExec => children.head
// A native ancestor embeds its old child in nativeOp. Replacing only its SparkPlan
// child would leave the incompatible native partial in that serialized plan.
case comet: CometExec =>
diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
index bbdc6db6e5..4e97c0c3c2 100644
--- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
@@ -23,7 +23,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometExec, CometHashAggregateExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometExec, CometHashAggregateExec, CometLocalTopKExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec}
import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.QueryStageExec
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike}
@@ -192,15 +192,19 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan
case sparkToColumnar: CometSparkToColumnarExec => sparkToColumnar.child
case RowToColumnarExec(child) => child
}
- val reverted = transformStageUp(stripped) { case cometExec: CometExec =>
- if (cometExec.originalPlan.children.size == cometExec.children.size) {
- cometExec.originalPlan.withNewChildren(cometExec.children)
- } else {
- logWarning(
- "Comet plan and original have different child count for " +
- s"${cometExec.getClass.getSimpleName}, using originalPlan as-is.")
- cometExec.originalPlan
- }
+ val reverted = transformStageUp(stripped) {
+ // Local candidate selection was inserted by Comet. Only the outer TopK owns
+ // the original Spark operator's offset and projection.
+ case local: CometLocalTopKExec => local.child
+ case cometExec: CometExec =>
+ if (cometExec.originalPlan.children.size == cometExec.children.size) {
+ cometExec.originalPlan.withNewChildren(cometExec.children)
+ } else {
+ logWarning(
+ "Comet plan and original have different child count for " +
+ s"${cometExec.getClass.getSimpleName}, using originalPlan as-is.")
+ cometExec.originalPlan
+ }
}
insertTransitions(reverted)
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala
index 07c44fd195..5628ec039b 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala
@@ -82,29 +82,18 @@ object CometExecUtils {
}
}
- /**
- * Prepare Projection + TopK native plan for CometTakeOrderedAndProjectExec.
- */
+ /** Wrap a native input plan in the requested output projection. */
def getProjectionNativePlan(
projectList: Seq[NamedExpression],
- outputAttributes: Seq[Attribute],
- sortOrder: Seq[SortOrder],
- child: SparkPlan,
- limit: Int,
- offset: Int = 0): Option[Operator] = {
- getTopKNativePlan(outputAttributes, sortOrder, child, limit, offset).flatMap { topK =>
- val exprs = projectList.map(exprToProto(_, child.output))
-
- if (exprs.forall(_.isDefined)) {
- val projectBuilder = OperatorOuterClass.Projection.newBuilder()
- projectBuilder.addAllProjectList(exprs.map(_.get).asJava)
- val opBuilder = OperatorOuterClass.Operator
- .newBuilder()
- .addChildren(topK)
- Some(opBuilder.setProjection(projectBuilder).build())
- } else {
- None
- }
+ inputAttributes: Seq[Attribute],
+ input: Operator): Option[Operator] = {
+ val exprs = projectList.map(exprToProto(_, inputAttributes))
+ if (exprs.forall(_.isDefined)) {
+ val projection = OperatorOuterClass.Projection.newBuilder()
+ projection.addAllProjectList(exprs.map(_.get).asJava)
+ Some(Operator.newBuilder().addChildren(input).setProjection(projection).build())
+ } else {
+ None
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala
new file mode 100644
index 0000000000..b2dcc2599f
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala
@@ -0,0 +1,103 @@
+/*
+ * 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.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.execution.{SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.types.{ByteType, IntegerType, LongType, ShortType}
+
+import com.google.common.base.Objects
+
+import org.apache.comet.CometConf
+import org.apache.comet.serde.OperatorOuterClass.{Operator, Sort}
+import org.apache.comet.serde.QueryPlanSerde.exprToProto
+
+object CometLocalTopKExec {
+
+ /** Insert the local selection before native blocks are serialized, so it owns the scan. */
+ def create(op: TakeOrderedAndProjectExec): Option[CometLocalTopKExec] = {
+ if (!CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.get(op.conf) || op.limit <= 0 ||
+ SortOrder.orderingSatisfies(op.child.outputOrdering, op.sortOrder)) {
+ return None
+ }
+
+ // Initially support one direct signed integer column in a native Parquet scan.
+ // Other plans retain the existing TopKInput execution path.
+ (op.child, op.sortOrder) match {
+ case (scan: CometNativeScanExec, Seq(order))
+ if order.child.isInstanceOf[Attribute] && (order.dataType match {
+ case ByteType | ShortType | IntegerType | LongType => true
+ case _ => false
+ }) =>
+ exprToProto(order, scan.output).map { sortOrder =>
+ // Spark's physical limit already includes the offset. Each partition retains that
+ // many candidates; only the final TopK applies the offset after the shuffle.
+ val sort = Sort.newBuilder().addSortOrders(sortOrder).setFetch(op.limit).setSkip(0)
+ val nativeOp = Operator
+ .newBuilder()
+ .setPlanId(op.id)
+ .setSort(sort)
+ .addChildren(scan.nativeOp)
+ .build()
+ CometLocalTopKExec(
+ nativeOp,
+ op,
+ scan.output,
+ op.limit,
+ op.sortOrder,
+ scan,
+ SerializedPlan(None))
+ }
+ case _ => None
+ }
+ }
+}
+
+/** Selects candidates within one Spark partition in the same native execution as its scan. */
+case class CometLocalTopKExec(
+ override val nativeOp: Operator,
+ override val originalPlan: SparkPlan,
+ override val output: Seq[Attribute],
+ limit: Int,
+ sortOrder: Seq[SortOrder],
+ child: SparkPlan,
+ override val serializedPlanOpt: SerializedPlan)
+ extends CometUnaryExec {
+
+ override def outputPartitioning: Partitioning = child.outputPartitioning
+
+ override def outputOrdering: Seq[SortOrder] = sortOrder
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ copy(child = newChild)
+
+ override def stringArgs: Iterator[Any] = Iterator(limit, sortOrder, output, child)
+
+ override def equals(obj: Any): Boolean = obj match {
+ case other: CometLocalTopKExec =>
+ output == other.output && limit == other.limit && sortOrder == other.sortOrder &&
+ child == other.child && serializedPlanOpt == other.serializedPlanOpt
+ case _ => false
+ }
+
+ override def hashCode(): Int = Objects.hashCode(output, Int.box(limit), sortOrder, child)
+
+}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
index 09dd944d93..6347241c56 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
@@ -79,7 +79,7 @@ object CometTakeOrderedAndProjectExec extends CometSink[TakeOrderedAndProjectExe
op.offset,
op.sortOrder,
op.projectList,
- op.child))
+ CometLocalTopKExec.create(op).getOrElse(op.child)))
}
}
@@ -123,6 +123,23 @@ case class CometTakeOrderedAndProjectExec(
lazy val orderingSatisfies: Boolean =
SortOrder.orderingSatisfies(child.outputOrdering, sortOrder)
+ // Also exposed to tests so they can inspect the plan actually used for final execution.
+ def finalNativePlan(inputPartitions: Int): Option[OperatorOuterClass.Operator] = {
+ val selectedLocally = child match {
+ case local: CometLocalTopKExec =>
+ inputPartitions == 1 && local.limit == limit && local.sortOrder == sortOrder
+ case _ => false
+ }
+ val selection = if (selectedLocally) {
+ // The local TopK has already ordered and bounded the only input partition. Spark's
+ // physical limit includes the offset; the native limit applies that offset once.
+ CometExecUtils.getLimitNativePlan(child.output, limit, offset)
+ } else {
+ CometExecUtils.getTopKNativePlan(child.output, sortOrder, child, limit, offset)
+ }
+ selection.flatMap(CometExecUtils.getProjectionNativePlan(projectList, child.output, _))
+ }
+
protected override def doExecuteColumnar(): RDD[ColumnarBatch] = {
val childRDD = child.executeColumnar()
if (childRDD.getNumPartitions == 0 || limit == 0) {
@@ -131,7 +148,11 @@ case class CometTakeOrderedAndProjectExec(
val singlePartitionRDD = if (childRDD.getNumPartitions == 1) {
childRDD
} else {
- val localTopK = if (orderingSatisfies) {
+ val localTopK = if (child.isInstanceOf[CometLocalTopKExec]) {
+ // A fused child is already ordered and bounded. Check it before orderingSatisfies,
+ // which is also true for that child, to avoid adding a redundant native limit.
+ childRDD
+ } else if (orderingSatisfies) {
CometExecUtils.getNativeLimitRDD(childRDD, child.output, limit)
} else {
val numParts = childRDD.getNumPartitions
@@ -166,10 +187,8 @@ case class CometTakeOrderedAndProjectExec(
}
// Serialize the plan once before mapping to avoid repeated serialization per partition
- val topKAndProjection = CometExecUtils
- .getProjectionNativePlan(projectList, child.output, sortOrder, child, limit, offset)
- .get
- val serializedTopKAndProjection = CometExec.serializeNativePlan(topKAndProjection)
+ val serializedFinalPlan =
+ CometExec.serializeNativePlan(finalNativePlan(childRDD.getNumPartitions).get)
val finalOutputLength = output.length
val finalInputSchema = CometUtils.fromAttributes(child.output)
singlePartitionRDD.mapPartitionsInternal { iter =>
@@ -177,7 +196,7 @@ case class CometTakeOrderedAndProjectExec(
CometArrowStream
.inputObjects(iter, finalInputSchema, "CometTakeOrderedAndProject-final"),
finalOutputLength,
- serializedTopKAndProjection,
+ serializedFinalPlan,
1,
0)
setSubqueries(it.id, this)
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index ed2d62a1a3..2634d89535 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -753,8 +753,8 @@ object CometExec {
* Built once on the driver from the SparkPlan tree, then consumed by either
* [[CometNativeExec.executeColumnarWithContext]] (to build a [[CometExecRDD]]) or the
* native-shuffle path (to drive [[CometNativeShuffleWriter]]). Captures broadcast partition
- * alignment, plan-data, subqueries, and encryption options so each consumer doesn't re-walk the
- * tree.
+ * alignment, plan-data, subqueries, encryption options, and per-partition file paths so each
+ * consumer doesn't re-walk the tree.
*/
private[comet] case class NativeExecContext(
inputs: Seq[RDD[_]],
@@ -770,7 +770,9 @@ private[comet] case class NativeExecContext(
// binary when this context rides on the non-transient CometShuffleDependency.nativeShuffleSpec.
@transient perPartitionByKey: Map[String, Array[Array[Byte]]],
shuffleScanIndices: Set[Int],
- hasScanInput: Boolean) {
+ hasScanInput: Boolean,
+ // Like plan data, file paths are sliced onto each task's Partition on the driver.
+ @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty) {
// Catch shape divergence (e.g. broadcast scans with different partition counts after DPP
// filtering) at construction so consumers don't trip ArrayIndexOutOfBoundsException at
// partition idx access time.
@@ -849,7 +851,8 @@ abstract class CometNativeExec extends CometExec {
ctx.subqueries,
ctx.broadcastedHadoopConfForEncryption,
ctx.encryptedFilePaths,
- ctx.shuffleScanIndices) {
+ ctx.shuffleScanIndices,
+ perPartitionFilePaths = ctx.perPartitionFilePaths) {
override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = {
val res = super.compute(split, context)
if (ctx.hasScanInput) {
@@ -1041,6 +1044,23 @@ abstract class CometNativeExec extends CometExec {
throw new CometRuntimeException(s"No input for CometNativeExec:\n $this")
}
+ // Fused scans do not run CometNativeScanExec.doExecuteColumnar, so carry their file paths
+ // through this execution boundary for read-error diagnostics. Only include scans in this
+ // native block and combine paths at the same partition index, never across partitions.
+ val scanFilePaths = sparkPlans.collect { case scan: CometNativeScanExec =>
+ scan.perPartitionFilePaths
+ }
+ val perPartitionFilePaths = if (scanFilePaths.isEmpty) {
+ Array.empty[Seq[String]]
+ } else {
+ require(
+ scanFilePaths.forall(_.length == firstNonBroadcastPlanNumPartitions),
+ "Native scan file paths must match the execution partition count")
+ Array.tabulate[Seq[String]](firstNonBroadcastPlanNumPartitions) { idx =>
+ scanFilePaths.flatMap(_(idx)).toVector
+ }
+ }
+
NativeExecContext(
inputs = inputs.toSeq,
numPartitions = firstNonBroadcastPlanNumPartitions,
@@ -1057,7 +1077,8 @@ abstract class CometNativeExec extends CometExec {
// reported once the scan is fused into a larger native block, where only the block
// root's `compute` runs. `reportScanInputMetrics` self-filters on the `bytes_scanned`
// metric, so leaves that don't track it are a no-op.
- hasScanInput = sparkPlans.exists(_.isInstanceOf[CometLeafExec]))
+ hasScanInput = sparkPlans.exists(_.isInstanceOf[CometLeafExec]),
+ perPartitionFilePaths = perPartitionFilePaths)
}
/**
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometTopKSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometTopKSuite.scala
new file mode 100644
index 0000000000..1d973ae9bd
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/exec/CometTopKSuite.scala
@@ -0,0 +1,374 @@
+/*
+ * 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.exec
+
+import java.net.URI
+import java.nio.{ByteBuffer, ByteOrder}
+import java.nio.file.Files
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.SparkThrowable
+import org.apache.spark.sql.{CometTestBase, Row}
+import org.apache.spark.sql.comet.{CometLocalTopKExec, CometNativeScanExec, CometTakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
+import org.apache.comet.serde.OperatorOuterClass.Operator
+
+class CometTopKSuite extends CometTestBase {
+
+ private def localTopK(plan: SparkPlan): CometLocalTopKExec = {
+ val nodes = collect(plan) { case topK: CometLocalTopKExec => topK }
+ assert(nodes.size == 1, s"Expected one local TopK:\n$plan")
+ nodes.head
+ }
+
+ private def operators(op: Operator): Seq[Operator] =
+ Seq(op) ++ op.getChildrenList.asScala.toSeq.flatMap(operators)
+
+ for (partitions <- Seq(1, 3); adaptive <- Seq(false, true)) {
+ test(s"local TopK shares the native scan: partitions=$partitions, AQE=$adaptive") {
+ withSQLConf(
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> partitions.toString,
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "4194304") {
+ withTempPath { path =>
+ spark
+ .range(0, 120, 1, partitions)
+ .selectExpr("CAST(119 - id AS INT) AS k", "id AS payload")
+ .write
+ .parquet(path.getCanonicalPath)
+ withParquetTable(path.getCanonicalPath, "topk_input") {
+ val (_, plan) = checkSparkAnswerAndOperator(
+ sql("SELECT payload FROM topk_input ORDER BY k LIMIT 5 OFFSET 7"),
+ Seq(classOf[CometLocalTopKExec], classOf[CometNativeScanExec]))
+ val local = localTopK(plan)
+ assert(local.limit == 12)
+ assert(local.executeColumnar().getNumPartitions == partitions)
+ val block = Operator.parseFrom(local.serializedPlanOpt.plan.get)
+ assert(block.hasSort && block.getSort.getFetch == 12)
+ assert(block.getSort.getSkip == 0)
+ assert(operators(block).count(_.hasNativeScan) == 1)
+ assert(!operators(block).exists(_.hasScan), "Local TopK must not use Arrow input")
+ assert(local.metrics("output_rows").value == 12L * partitions)
+ val scans = collect(plan) { case scan: CometNativeScanExec => scan }
+ assert(scans.head.metrics("output_rows").value == 120L)
+ assert(scans.head.metrics("bytes_scanned").value > 0L)
+ val global = collect(plan) { case topK: CometTakeOrderedAndProjectExec => topK }.head
+ assert(global.limit == 12 && global.offset == 7)
+ val finalPlan = global.finalNativePlan(partitions).get
+ assert(finalPlan.hasProjection)
+ val finalSelection = finalPlan.getChildren(0)
+ if (partitions == 1) {
+ assert(!operators(finalPlan).exists(_.hasSort))
+ assert(finalSelection.hasLimit)
+ assert(finalSelection.getLimit.getLimit == 12)
+ assert(finalSelection.getLimit.getOffset == 7)
+ assert(
+ global
+ .copy(child = local.copy(limit = 13))
+ .finalNativePlan(1)
+ .get
+ .getChildren(0)
+ .hasSort)
+ assert(
+ global
+ .copy(child = local.copy(sortOrder = Seq.empty))
+ .finalNativePlan(1)
+ .get
+ .getChildren(0)
+ .hasSort)
+ } else {
+ assert(finalSelection.hasSort)
+ }
+ assert(local.canonicalized != local.copy(limit = 13).canonicalized)
+ }
+ }
+ }
+ }
+ }
+
+ for (adaptive <- Seq(false, true)) {
+ test(s"disabling fusion retains native TopK: AQE=$adaptive") {
+ assert(!CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.get())
+ withSQLConf(
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) {
+ withParquetTable((0 until 20).map(i => (i, 20 - i)), "topk_input") {
+ val (_, plan) = checkSparkAnswerAndOperator(
+ sql("SELECT _2 FROM topk_input ORDER BY _1 LIMIT 5 OFFSET 2"),
+ Seq(classOf[CometTakeOrderedAndProjectExec], classOf[CometNativeScanExec]))
+ assert(collect(plan) { case local: CometLocalTopKExec => local }.isEmpty)
+ val global = collect(plan) { case topK: CometTakeOrderedAndProjectExec => topK }.head
+ assert(global.finalNativePlan(1).get.getChildren(0).hasSort)
+ }
+ }
+ }
+ }
+
+ for (partitions <- Seq(1, 2)) {
+ test(s"local TopK preserves nulls, ties and small inputs: partitions=$partitions") {
+ withSQLConf(
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> partitions.toString,
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "4194304") {
+ withTempPath { path =>
+ spark
+ .range(0, 24, 1, partitions)
+ .selectExpr("CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id % 5 AS INT) END AS k")
+ .write
+ .parquet(path.getCanonicalPath)
+ withParquetTable(path.getCanonicalPath, "topk_input") {
+ for (direction <- Seq("ASC", "DESC"); nulls <- Seq("FIRST", "LAST")) {
+ for ((limit, offset) <- Seq((3, 0), (3, 4), (50, 0), (5, 30))) {
+ checkSparkAnswerAndOperator(
+ sql(
+ s"SELECT k FROM topk_input ORDER BY k $direction NULLS $nulls " +
+ s"LIMIT $limit OFFSET $offset"),
+ Seq(classOf[CometLocalTopKExec]))
+ }
+ }
+ checkSparkAnswerAndOperator(
+ sql("SELECT k FROM topk_input ORDER BY k LIMIT 0"),
+ classOf[LocalTableScanExec])
+ }
+ }
+ }
+ }
+
+ }
+
+ for (keyType <- Seq("TINYINT", "SMALLINT", "INT", "BIGINT")) {
+ test(s"local TopK supports signed $keyType keys and final projections") {
+ withSQLConf(CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true") {
+ withTempPath { path =>
+ spark
+ .range(-120, 120, 1, 2)
+ .selectExpr(s"CAST(id AS $keyType) AS k", "id * 10 AS payload")
+ .write
+ .parquet(path.getCanonicalPath)
+ withParquetTable(path.getCanonicalPath, "topk_input") {
+ for (direction <- Seq("ASC", "DESC")) {
+ val query = sql(
+ s"SELECT payload FROM topk_input ORDER BY k $direction " +
+ "LIMIT 7 OFFSET 3")
+ checkSparkAnswerAndOperator(query, Seq(classOf[CometLocalTopKExec]))
+ val expected = (if (direction == "ASC") -117L until -110L
+ else
+ 116L to 110L by -1L).map(key => Row(key * 10))
+ assert(query.collect().toSeq == expected)
+ assert(query.collect().toSeq == expected)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("fusion leaves computed projections and unsupported sort keys on the existing path") {
+ withSQLConf(CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true") {
+ withParquetTable((0 until 20).map(i => (i, 20 - i, i.toDouble)), "topk_input") {
+ val queries = Seq("_1 + 1", "_1, _2", "_3")
+ .map(order => s"SELECT * FROM topk_input ORDER BY $order LIMIT 5") :+
+ "SELECT _2 + 1 AS adjusted FROM topk_input ORDER BY _1 LIMIT 5"
+ for (query <- queries) {
+ val (_, plan) =
+ checkSparkAnswerAndOperator(sql(query), Seq(classOf[CometTakeOrderedAndProjectExec]))
+ assert(collect(plan) { case local: CometLocalTopKExec => local }.isEmpty)
+ }
+ }
+ }
+ }
+
+ test("fusion leaves an already ordered input on the existing path") {
+ withSQLConf(
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.catalyst.optimizer.EliminateSorts") {
+ withParquetTable((0 until 20).map(i => (i, 20 - i)), "topk_input") {
+ val query = spark
+ .table("topk_input")
+ .sortWithinPartitions("_1", "_2")
+ .orderBy("_1")
+ .limit(5)
+ val (_, plan) = checkSparkAnswerAndOperator(query)
+ val topK = collect(plan) { case node: CometTakeOrderedAndProjectExec => node }.head
+ assert(topK.orderingSatisfies)
+ assert(collect(plan) { case local: CometLocalTopKExec => local }.isEmpty)
+ }
+ }
+ }
+
+ test("fusion requires a native scan and preserves Spark fallback") {
+ withSQLConf(
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true",
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false") {
+ withParquetTable((0 until 20).map(i => (i, 20 - i)), "topk_input") {
+ val (_, plan) =
+ checkSparkAnswer(sql("SELECT _2 FROM topk_input ORDER BY _1 LIMIT 5 OFFSET 2"))
+ assert(collect(plan) { case topK: TakeOrderedAndProjectExec => topK }.size == 1)
+ assert(collect(plan) { case local: CometLocalTopKExec => local }.isEmpty)
+ }
+ }
+ }
+
+ test("fusion preserves empty Parquet input") {
+ withSQLConf(CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true") {
+ withTempPath { path =>
+ spark.range(0).write.parquet(path.getCanonicalPath)
+ withParquetTable(path.getCanonicalPath, "topk_empty") {
+ val query = sql("SELECT id FROM topk_empty ORDER BY id LIMIT 10 OFFSET 3")
+ checkSparkAnswerAndOperator(query)
+ assert(query.collect().isEmpty)
+ }
+ }
+ }
+ }
+
+ test("local TopK preserves partition columns in empty buckets") {
+ val session = spark
+ import session.implicits._
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.BUCKETING_ENABLED.key -> "true",
+ "spark.sql.sources.bucketing.autoBucketedScan.enabled" -> "false") {
+ withTable("topk_bucketed_partitioned") {
+ // Two rows cannot fill eight buckets, so the native plan must also handle empty scans.
+ Seq((1, 10), (2, 20))
+ .toDF("id", "p")
+ .coalesce(1)
+ .write
+ .format("parquet")
+ .partitionBy("p")
+ .bucketBy(8, "id")
+ .saveAsTable("topk_bucketed_partitioned")
+ for (fusion <- Seq(false, true)) {
+ withSQLConf(CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> fusion.toString) {
+ for ((columns, expected) <- Seq(
+ "id, p" -> Seq(Row(1, 10), Row(2, 20)),
+ "p, id" -> Seq(Row(10, 1), Row(20, 2)))) {
+ val query =
+ sql(s"SELECT $columns FROM topk_bucketed_partitioned ORDER BY p LIMIT 5")
+ val plan = query.queryExecution.executedPlan
+ val scans = collect(plan) { case scan: CometNativeScanExec => scan }
+ assert(scans.size == 1, s"Expected one native scan:\n$plan")
+ assert(scans.head.bucketedScan)
+ assert(scans.head.perPartitionData.length == 8)
+ assert(scans.head.perPartitionFilePaths.count(_.isEmpty) >= 6)
+ val topKs = collect(plan) { case topK: CometTakeOrderedAndProjectExec => topK }
+ assert(topKs.size == 1, s"Expected one native TopK:\n$plan")
+ assert(topKs.head.child.executeColumnar().getNumPartitions == 8)
+ val locals = collect(plan) { case local: CometLocalTopKExec => local }
+ assert(locals.size == (if (fusion) 1 else 0), s"fusion=$fusion:\n$plan")
+ assert(query.collect().toSeq == expected)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ for (corruption <- Seq("footer", "page")) {
+ test(s"local TopK preserves the corrupt Parquet $corruption file in read errors") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "2",
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "4194304",
+ SQLConf.IGNORE_CORRUPT_FILES.key -> "false") {
+ withTempPath { path =>
+ spark
+ .range(0, 1000, 1, 2)
+ .selectExpr("CAST(id AS INT) AS k")
+ .write
+ .option("compression", "uncompressed")
+ .option("parquet.enable.dictionary", "false")
+ .parquet(path.getCanonicalPath)
+ val entries = Files.list(path.toPath)
+ val files =
+ try {
+ entries
+ .iterator()
+ .asScala
+ .filter(_.getFileName.toString.endsWith(".parquet"))
+ .toSeq
+ .sortBy(_.getFileName.toString)
+ } finally entries.close()
+ assert(files.size == 2)
+ val corruptFile = files.last
+ val bytes = Files.readAllBytes(corruptFile)
+ if (corruption == "footer") {
+ // Keep the file length but destroy the footer magic, failing during reader setup.
+ java.util.Arrays.fill(bytes, bytes.length - 4, bytes.length, 0xff.toByte)
+ } else {
+ val footerLength = ByteBuffer
+ .wrap(bytes, bytes.length - 8, 4)
+ .order(ByteOrder.LITTLE_ENDIAN)
+ .getInt()
+ assert(bytes.length - 8 - footerLength > 100)
+ // Keep the complete footer so setup succeeds; fail while decoding the first page.
+ java.util.Arrays.fill(bytes, 4, 100, 0xff.toByte)
+ }
+ Files.write(corruptFile, bytes)
+ withTempView("topk_corrupt_input") {
+ spark.read
+ .schema("k INT")
+ .parquet(path.getCanonicalPath)
+ .createOrReplaceTempView("topk_corrupt_input")
+ for (fusion <- Seq(false, true)) {
+ withSQLConf(CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> fusion.toString) {
+ val query = sql("SELECT k FROM topk_corrupt_input ORDER BY k LIMIT 5")
+ val plan = query.queryExecution.executedPlan
+ val scans = collect(plan) { case scan: CometNativeScanExec => scan }
+ assert(scans.size == 1, s"Expected one native scan:\n$plan")
+ val partitionPaths = scans.head.perPartitionFilePaths
+ assert(partitionPaths.length == 2 && partitionPaths.forall(_.size == 1))
+ val expectedPath = partitionPaths.flatten
+ .find(file => new URI(file) == corruptFile.toUri)
+ .getOrElse(fail(s"Missing corrupt file $corruptFile in scan partitions"))
+ val topKs = collect(plan) { case topK: CometTakeOrderedAndProjectExec => topK }
+ assert(topKs.size == 1, s"Expected one native TopK:\n$plan")
+ assert(topKs.head.child.executeColumnar().getNumPartitions == 2)
+ val locals = collect(plan) { case local: CometLocalTopKExec => local }
+ assert(locals.size == (if (fusion) 1 else 0), s"fusion=$fusion:\n$plan")
+ val failure = intercept[Exception](query.collect())
+ // Spark 3.x uses a legacy error class, but still exposes the structured path.
+ val errorClass =
+ if (isSpark40Plus) "FAILED_READ_FILE.NO_HINT" else "_LEGACY_ERROR_TEMP_2064"
+ val readError = causeChain(failure)
+ .collectFirst {
+ case error: SparkThrowable if error.getErrorClass == errorClass => error
+ }
+ .getOrElse(fail(s"Missing $errorClass in ${causeChain(failure)}"))
+ // A lower-level cause can mention the file even when the structured path is lost.
+ // With one file per partition, the path must name only the corrupt file.
+ assert(readError.getMessageParameters.get("path") == expectedPath)
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala
index 0db2351c72..5c3bfe4cbe 100644
--- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala
@@ -19,7 +19,7 @@
package org.apache.comet.rules
-import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.{CometTestBase, Row}
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
import org.apache.spark.sql.comet._
import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
@@ -141,6 +141,49 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
}
}
+ for (adaptive <- Seq(false, true)) {
+ test(s"transition reversion preserves local TopK: AQE=$adaptive") {
+ withSQLConf(
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1",
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") {
+ withTempPath { path =>
+ spark
+ .range(0, 20, 1, 1)
+ .selectExpr("CAST(id AS INT) AS k", "id * 10 AS payload")
+ .write
+ .parquet(path.getCanonicalPath)
+ withParquetTable(path.getCanonicalPath, "topk_revert") {
+ for (offset <- Seq(2, 7, 0); projection <- Seq("k", "payload")) {
+ val query = s"SELECT $projection FROM topk_revert ORDER BY k LIMIT 5 OFFSET $offset"
+ withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") {
+ val plan = sql(query).queryExecution.executedPlan
+ val local = collect(plan) { case topK: CometLocalTopKExec => topK }
+ assert(local.size == 1, s"Expected a fused TopK before reversion:\n$plan")
+ assert(local.head.child.isInstanceOf[CometNativeScanExec])
+ }
+ val df = sql(query)
+ val expected = (offset until offset + 5).map { key =>
+ if (projection == "k") Row(key) else Row(key * 10L)
+ }
+ withClue(query) {
+ assert(df.collect().toSeq == expected)
+ }
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(countCometExecs(plan) == 0, s"Expected stage reversion:\n$plan")
+ assert(
+ collect(plan) { case topK: TakeOrderedAndProjectExec => topK }.size == 1,
+ s"Reversion must restore offset and projection exactly once:\n$plan")
+ }
+ }
+ }
+ }
+ }
+ }
+
test("revertToSpark removes all Comet operators from a plan with transitions") {
withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometTopKBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometTopKBenchmark.scala
new file mode 100644
index 0000000000..4ceb4679cf
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometTopKBenchmark.scala
@@ -0,0 +1,206 @@
+/*
+ * 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.benchmark
+
+import scala.concurrent.duration._
+
+import org.apache.hadoop.fs.Path
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.comet.{CometLocalTopKExec, CometNativeScanExec, CometTakeOrderedAndProjectExec}
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+import org.apache.comet.serde.OperatorOuterClass.Operator
+
+/**
+ * Compare local TopK fusion enabled and disabled, using a native release build. The make target
+ * builds and packages the release library before running the benchmark:
+ *
+ * {{{
+ * make benchmark-org.apache.spark.sql.benchmark.CometTopKBenchmark -- \
+ * 1048576 1,4 1,16 16,4096 ascending,descending,random 5 off,on
+ * }}}
+ *
+ * Arguments are rows, scan partitions, BIGINT payload columns, K values, layouts, minimum
+ * measured iterations, and mode order. Defaults match the example. Use on,off for a second run
+ * with the opposite case order. Selectors between rows and iterations accept comma-separated
+ * lists. For a quick smoke run, use 16384 rows and three iterations. Each case warms up for at
+ * least 500 ms and measures for at least 500 ms. Results include query planning and collection;
+ * fixtures, Spark result checks, and native plan checks run outside the timings. The wide payload
+ * and descending layout expose the possible cost of losing scan/TopK overlap.
+ */
+object CometTopKBenchmark extends CometBenchmarkBase {
+
+ private val nativeConfigs = Seq(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true",
+ CometConf.COMET_BATCH_SIZE.key -> "4096",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "4194304")
+
+ private def keyExpression(layout: String, rows: Int): String = {
+ val unsigned = layout match {
+ case "ascending" => "id"
+ case "descending" => s"$rows - id - 1"
+ // An odd multiplier permutes all values modulo a power of two, without duplicate keys.
+ case "random" => s"pmod(id * 104729, $rows)"
+ }
+ s"CAST(($unsigned) - ${rows / 2} AS INT)"
+ }
+
+ private def writeInput(
+ directory: String,
+ rows: Int,
+ partitions: Int,
+ payloadColumns: Int,
+ layout: String): Long = {
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val payload = (0 until payloadColumns).map(i => s"id * ${17 + i * 2} AS payload_$i")
+ spark
+ .range(0, rows, 1, partitions)
+ .selectExpr((Seq(s"${keyExpression(layout, rows)} AS k") ++ payload): _*)
+ .write
+ .option("compression", "snappy")
+ .option("parquet.enable.dictionary", "false")
+ .parquet(directory)
+ }
+ val path = new Path(directory)
+ val files = path
+ .getFileSystem(spark.sessionState.newHadoopConf())
+ .listStatus(path)
+ .filter(_.getPath.getName.endsWith(".parquet"))
+ require(files.length == partitions, s"Expected $partitions files, found ${files.length}")
+ files.map(_.getLen).max
+ }
+
+ private def verifyPlanAndResult(
+ benchmark: Benchmark,
+ query: String,
+ expected: Seq[Row],
+ fusion: Boolean,
+ k: Int,
+ partitions: Int,
+ expectedInputRows: Int): Unit = {
+ val df = spark.sql(query)
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ val local = collect(plan) { case node: CometLocalTopKExec => node }
+ val scans = collect(plan) { case node: CometNativeScanExec => node }
+ val outer = collect(plan) { case node: CometTakeOrderedAndProjectExec => node }
+ require(scans.size == 1 && outer.size == 1, s"Expected native TopK and scan:\n$plan")
+ require(local.size == (if (fusion) 1 else 0), s"Unexpected fusion=$fusion:\n$plan")
+ require(outer.head.child.executeColumnar().getNumPartitions == partitions)
+ if (fusion) {
+ require(local.head.limit == k)
+ val block = Operator.parseFrom(local.head.serializedPlanOpt.plan.get)
+ require(block.hasSort && block.getSort.getFetch == k && block.getSort.getSkip == 0)
+ require(block.getChildrenCount == 1 && block.getChildren(0).hasNativeScan)
+ }
+ val finalPlan = outer.head.finalNativePlan(partitions).get
+ require(finalPlan.hasProjection && finalPlan.getChildrenCount == 1)
+ val selection = finalPlan.getChildren(0)
+ if (fusion && partitions == 1) {
+ require(selection.hasLimit, "Single-partition fusion must avoid a second TopK heap")
+ require(selection.getLimit.getLimit == k && selection.getLimit.getOffset == 0)
+ } else {
+ require(selection.hasSort && selection.getSort.getFetch == k)
+ }
+ require(df.collect().toSeq == expected, s"Incorrect results with fusion=$fusion")
+ val scanMetrics = scans.head.metrics
+ require(
+ scanMetrics.get("output_rows").exists(_.value == expectedInputRows.toLong),
+ "TopK fusion must still decode all input rows")
+ val readerWork = Seq("output_rows", "bytes_scanned", "row_groups_pruned_dynamic_filter")
+ .flatMap(name => scanMetrics.get(name).map(metric => s"$name=${metric.value}"))
+ .mkString(", ")
+ benchmark.out.println(s"fusion=$fusion, $readerWork")
+ }
+
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ require(
+ mainArgs.length <= 7,
+ "Usage: CometTopKBenchmark [rows] [partitions] [payload-columns] [ks] [layouts] " +
+ "[iterations] [off,on|on,off]")
+ val rows = mainArgs.headOption.map(_.toInt).getOrElse(1 << 20)
+ val partitions = mainArgs.lift(1).getOrElse("1,4").split(",").map(_.toInt).toSeq
+ val payloadColumns = mainArgs.lift(2).getOrElse("1,16").split(",").map(_.toInt).toSeq
+ val ks = mainArgs.lift(3).getOrElse("16,4096").split(",").map(_.toInt).toSeq
+ val layouts = mainArgs.lift(4).getOrElse("ascending,descending,random").split(",").toSeq
+ val iterations = mainArgs.lift(5).map(_.toInt).getOrElse(5)
+ val modeOrder = mainArgs.lift(6).getOrElse("off,on")
+ require(Set("off,on", "on,off").contains(modeOrder), "mode order must be off,on or on,off")
+ val fusionModes = modeOrder.split(",").map(_ == "on").toSeq
+ require(rows >= 16384 && (rows & (rows - 1)) == 0, "rows must be a power of two >= 16384")
+ require(partitions.forall(p => p > 0 && rows % p == 0), "partitions must divide rows")
+ require(payloadColumns.forall(_ > 0), "payload-columns must be positive")
+ require(ks.forall(k => k > 0 && k <= rows), "ks must be positive and at most rows")
+ require(layouts.forall(Set("ascending", "descending", "random")), "Unknown layout")
+ require(iterations > 0, "iterations must be positive")
+
+ withTempPath { root =>
+ withTempTable("topk_input") {
+ for (p <- partitions; columns <- payloadColumns; layout <- layouts) {
+ val directory = s"${root.getCanonicalPath}/$layout-p$p-c$columns"
+ val maxFileBytes = writeInput(directory, rows, p, columns, layout)
+ // Keep each file in its own scan partition without inserting a repartition operator,
+ // which would interrupt fusion. Verify the actual partition count for every case.
+ val scanConfigs = nativeConfigs :+
+ (SQLConf.FILES_MAX_PARTITION_BYTES.key -> (maxFileBytes + 4194304L).toString)
+ spark.read.parquet(directory).createOrReplaceTempView("topk_input")
+ for (k <- ks) {
+ val query = s"SELECT * FROM topk_input ORDER BY k LIMIT $k"
+ var expected = Seq.empty[Row]
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ expected = spark.sql(query).collect().toSeq
+ }
+ require(expected.map(_.getInt(0)) == (-rows / 2 until -rows / 2 + k))
+ val name = s"TopK $layout, partitions=$p, payloadColumns=$columns, K=$k"
+ runBenchmark(name) {
+ val benchmark = new Benchmark(
+ name,
+ rows,
+ minNumIters = iterations,
+ warmupTime = 500.millis,
+ minTime = 500.millis,
+ output = output)
+ withSQLConf(scanConfigs: _*) {
+ fusionModes.foreach { fusion =>
+ val fusionConfig =
+ CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> fusion.toString
+ withSQLConf(fusionConfig) {
+ verifyPlanAndResult(benchmark, query, expected, fusion, k, p, rows)
+ }
+ benchmark.addCase(if (fusion) "fusion enabled" else "fusion disabled") { _ =>
+ withSQLConf(fusionConfig) {
+ spark.sql(query).collect()
+ }
+ }
+ }
+ benchmark.run()
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
index e2894d1891..871255b1ec 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
@@ -228,7 +228,10 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase {
commonByKey = Map.empty,
perPartitionByKey = perPartitionByKey,
shuffleScanIndices = Set.empty,
- hasScanInput = false)
+ hasScanInput = false,
+ perPartitionFilePaths = Array.tabulate(numPartitions) { idx =>
+ Seq(s"file:/tmp/part-$idx.parquet")
+ })
val spec = NativeShuffleSpec(Operator.getDefaultInstance, childMetricNode, execContext)
val dep = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch](
_rdd = rdd,
@@ -241,6 +244,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase {
// Pre-fix this pair grew from ~13KB to ~10MB between 10 and 10000 partitions. With the map held
// @transient on both the RDD and the NativeExecContext, the pair stays roughly constant.
+ // The diagnostic file paths must also stay on the driver, rather than travel with every task.
val (_, smallDep) = build(10)
val (largeRdd, largeDep) = build(10000)
val smallPair = ser.serialize((smallDep.rdd, smallDep)).limit()
@@ -248,7 +252,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase {
assert(
math.abs(largePair - smallPair) < 100 * 1024,
s"serialized (rdd, dep) grew with partition count (small=$smallPair, large=$largePair); " +
- "the per-partition plan-data map is leaking into the broadcast task binary")
+ "per-partition plan data or file paths are leaking into the broadcast task binary")
// Each task's Partition object still carries its own slice so the writer can inject plan data.
val part = largeRdd.partitions(7).asInstanceOf[CometNativeShuffleInputPartition]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org