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/30 02:12:08 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6297-5d6902508f76dea6da1ac054ccf10bb3556e1964 created (now 12e85575df)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6297-5d6902508f76dea6da1ac054ccf10bb3556e1964
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 12e85575df fix: key each scan's planning data by its plan node in a native block (#6297)
This branch includes the following new commits:
new 12e85575df fix: key each scan's planning data by its plan node in a native block (#6297)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: key each scan's planning data by its plan node in a native block (#6297)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6297-5d6902508f76dea6da1ac054ccf10bb3556e1964
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 12e85575df13c161527866612b2d42aedbf7a668
Author: dustin <dw...@users.noreply.github.com>
AuthorDate: Wed Sep 30 00:50:31 2026 +0000
fix: key each scan's planning data by its plan node in a native block (#6297)
Scans in one native block have their planning data collected under a key
and merged with toMap, so two scans with the same key both read the last
one's files. The Iceberg key (metadata location plus the SparkScan hash)
and the Parquet key (source plus a hash of schema, data filters and
projection) leave out what can tell two scans of one table apart: storage
partitioned join grouping, partition filters and DPP. A storage partitioned
self-join with partially clustered distribution returned duplicate rows,
and a bucketed Parquet self-join without pushed filters read the wrong
partitions.
Both keys now include the scan operator's plan_id, which is the id of the
Spark plan node the scan was converted from. The merge in findAllPlanData
keeps one copy of identical data under a key and fails on different data
instead of dropping one side.
Closes #6278.
---
native/proto/src/proto/operator.proto | 11 +--
.../sql/comet/CometIcebergNativeScanExec.scala | 6 +-
.../spark/sql/comet/CometNativeScanExec.scala | 9 ++-
.../org/apache/spark/sql/comet/operators.scala | 90 ++++++++++++++++++----
.../org/apache/comet/CometIcebergNativeSuite.scala | 64 ++++++++++++++-
.../org/apache/comet/exec/CometJoinSuite.scala | 75 +++++++++++++++++-
.../sql/comet/CometScanWithPlanDataSuite.scala | 42 +++++++++-
.../spark/sql/comet/PlanDataInjectorSuite.scala | 68 +++++++++++++---
.../shuffle/CometCelebornShuffleManagerSuite.scala | 2 +-
9 files changed, 325 insertions(+), 42 deletions(-)
diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto
index 7b3a7aec0d..af5a4d263b 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -191,9 +191,10 @@ message NativeScan {
// Single partition's file list (injected at execution time)
SparkFilePartition file_partition = 2;
- // Hash half of the key under which this scan's planning data is stored and looked up at
- // execution time; the other half is common.source. Derived once on the driver
- // (CometNativeScanExec) so executors rebuild the key instead of hashing the common per task.
+ // Hash part of the key under which this scan's planning data is stored and looked up at
+ // execution time; the other parts are common.source and the operator's plan_id. Derived once
+ // on the driver (CometNativeScanExec) so executors rebuild the key instead of hashing the
+ // common per task.
// JVM-consumed only; the native side ignores it.
optional int32 source_key_hash = 3;
}
@@ -352,8 +353,8 @@ message IcebergScanCommon {
// Iceberg's SparkScan.hashCode() (folds in pushed filters, snapshot, branch, and read schema).
// A self-join/self-merge can read the same table (same metadata_location) twice with different
- // projections or filters, so PlanDataInjector keys planning data by
- // (metadata_location, scan_hash_code) rather than metadata_location alone.
+ // projections or filters, so PlanDataInjector keys planning data by metadata_location,
+ // scan_hash_code and the operator's plan_id rather than metadata_location alone.
int32 scan_hash_code = 15;
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala
index fbc90be15b..c944d3af8c 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala
@@ -271,9 +271,9 @@ case class CometIcebergNativeScanExec(
sendDriverMetrics()
val nativeMetrics = CometMetricNode.fromCometPlan(this)
val serializedPlan = CometExec.serializeNativePlan(nativeOp)
- // Key by the same (metadata_location, scan_hash_code) pair PlanDataInjector.injectPlanData
- // looks up via IcebergPlanDataInjector.getKey (see the scan_hash_code field comment in
- // operator.proto for why metadata_location alone cannot identify a scan).
+ // Key by metadata_location, scan_hash_code and the op's plan_id, the same key
+ // PlanDataInjector.injectPlanData looks up via IcebergPlanDataInjector.getKey (see
+ // PlanDataInjector.withPlanId for why the content parts alone cannot identify a scan).
val injectorKey = IcebergPlanDataInjector.getKey(nativeOp).getOrElse(metadataLocation)
new CometExecRDD(
sparkContext,
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
index 74cdbdc694..8ea48c1002 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
@@ -364,12 +364,17 @@ object CometNativeScanExec {
session: SparkSession,
scan: CometScanExec): CometNativeScanExec = {
// Generate unique key for this scan so PlanDataInjector can match common+partition data.
- // Multiple scans of same table with different projections/filters get different keys.
+ // Multiple scans of same table with different projections/filters get different keys, and
+ // the op's plan_id separates scans that match on those but read different files (see
+ // PlanDataInjector.withPlanId). It must be the same key NativeScanPlanDataInjector.getKey
+ // rebuilds from nativeOp.
// The hash is computed once here and embedded in the NativeScan proto, so executors
// (including the native shuffle writer) rebuild the key instead of hashing per task.
val common = nativeOp.getNativeScan.getCommon
val sourceKeyHash = NativeScanPlanDataInjector.sourceKeyHash(common)
- val sourceKey = NativeScanPlanDataInjector.sourceKey(common.getSource, sourceKeyHash)
+ val sourceKey = PlanDataInjector.withPlanId(
+ NativeScanPlanDataInjector.sourceKey(common.getSource, sourceKeyHash),
+ nativeOp.getPlanId)
val opWithKey = nativeOp.toBuilder
.setNativeScan(nativeOp.getNativeScan.toBuilder.setSourceKeyHash(sourceKeyHash))
.build()
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 ffb393a294..b6d16bfb22 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
@@ -81,7 +81,12 @@ private[comet] trait PlanDataInjector {
/** Check if this injector can handle the given operator. */
def canInject(op: Operator): Boolean
- /** Extract the key used to look up planning data for this operator. */
+ /**
+ * Extract the key used to look up planning data for this operator. The key must be unique per
+ * scan node within a native plan, so include the scan op's plan_id via
+ * `PlanDataInjector.withPlanId(..., op.getPlanId)`. findAllPlanData throws when two scans share
+ * a key with different data.
+ */
def getKey(op: Operator): Option[String]
/**
@@ -325,6 +330,21 @@ private[comet] object PlanDataInjector extends Logging {
private[comet] def preparedKey(injector: PlanDataInjector, key: String): String =
s"${injector.getClass.getName}:$key"
+ /**
+ * A scan's planning data key: its content key plus the scan op's `plan_id`. Two scans of one
+ * table can match on the content key yet read different files, such as the two sides of a
+ * bucketed or storage-partitioned self-join.
+ *
+ * The key is stable: it comes from the scan's own nativeOp `plan_id`, which parent blocks embed
+ * unchanged. The driver and executors therefore compute the same key, and so do copies of the
+ * exec, since a copy keeps its nativeOp.
+ *
+ * The key is unique: every scan in a native block was converted from its own CometScanExec or
+ * BatchScanExec. Spark assigns `SparkPlan.id` per plan-node instance from a JVM-wide counter,
+ * so two scans in one block never share it.
+ */
+ private[comet] def withPlanId(key: String, planId: Int): String = s"$key#$planId"
+
private def injectPlanData(
op: Operator,
commonByKey: Map[String, Array[Byte]],
@@ -418,9 +438,10 @@ private[comet] object PlanDataInjector extends Logging {
* by a unique identifier: one for common data (shared across partitions) and one for
* per-partition data.
*
- * Recognises Iceberg scans (keyed by metadata_location) plus any leaf scan that surfaces its
- * data via the [[CometScanWithPlanData]] trait (`CometNativeScanExec` and out-of-tree contrib
- * scans such as the Delta contrib's `CometDeltaNativeScanExec`).
+ * Recognises Iceberg scans (keyed by metadata_location, scan_hash_code and the op's plan_id,
+ * see IcebergPlanDataInjector.getKey) plus any leaf scan that surfaces its data via the
+ * [[CometScanWithPlanData]] trait (`CometNativeScanExec` and out-of-tree contrib scans such as
+ * the Delta contrib's `CometDeltaNativeScanExec`).
*
* Stops at stage boundaries (shuffle exchanges, etc.) because partition indices are only valid
* within the same stage.
@@ -446,10 +467,10 @@ private[comet] object PlanDataInjector extends Logging {
iceberg.sendDriverMetrics()
if (iceberg.commonData.nonEmpty && iceberg.perPartitionData.nonEmpty) {
// A self-join/self-merge can put two scans of the same table (same metadata_location)
- // in one native plan. Computing the key via IcebergPlanDataInjector.getKey, the same
- // function injectPlanData uses to look it up, keeps the two sides from drifting apart
- // (see the scan_hash_code field comment in operator.proto for why metadata_location
- // alone cannot distinguish them).
+ // in one native plan, and they can match on scan_hash_code too while reading
+ // different files, so the key also carries the op's plan_id (see withPlanId).
+ // Computing it via IcebergPlanDataInjector.getKey, the same function injectPlanData
+ // uses to look it up, keeps the two sides from drifting apart.
IcebergPlanDataInjector.getKey(iceberg.nativeOp) match {
case Some(key) =>
(Map(key -> iceberg.commonData), Map(key -> iceberg.perPartitionData))
@@ -491,9 +512,33 @@ private[comet] object PlanDataInjector extends Logging {
// Continue searching through other operators, combining results from all children
case _ =>
val results = plan.children.map(findAllPlanData)
- (results.flatMap(_._1).toMap, results.flatMap(_._2).toMap)
+ (
+ mergeByKey(results.map(_._1))(java.util.Arrays.equals(_, _)),
+ mergeByKey(results.map(_._2))(samePartitionData))
}
}
+
+ private def samePartitionData(a: Array[Array[Byte]], b: Array[Array[Byte]]): Boolean =
+ a.length == b.length && a.indices.forall(i => java.util.Arrays.equals(a(i), b(i)))
+
+ /**
+ * Merges the planning data found under sibling subtrees. The injector looks data up by key, so
+ * two scans sharing a key would both read whichever entry survived. The same data reached twice
+ * is kept once; different data under one key is a planning bug and fails the query rather than
+ * returning another scan's rows.
+ */
+ private def mergeByKey[V](maps: Seq[Map[String, V]])(same: (V, V) => Boolean): Map[String, V] =
+ maps.foldLeft(Map.empty[String, V]) { (merged, next) =>
+ next.foldLeft(merged) { case (acc, (key, value)) =>
+ acc.get(key) match {
+ case None => acc + (key -> value)
+ case Some(existing) if same(existing, value) => acc
+ case Some(_) =>
+ throw new CometRuntimeException(
+ s"Two scans in one native plan have different planning data under key $key")
+ }
+ }
+ }
}
/**
@@ -512,7 +557,9 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector {
override def getKey(op: Operator): Option[String] = {
val common = op.getIcebergScan.getCommon
- Some(s"${common.getMetadataLocation}_${common.getScanHashCode}")
+ Some(
+ PlanDataInjector
+ .withPlanId(s"${common.getMetadataLocation}_${common.getScanHashCode}", op.getPlanId))
}
// Parsed once per stage or shuffle by injectPlanData's memo, never per partition.
@@ -550,17 +597,20 @@ private[comet] object NativeScanPlanDataInjector extends PlanDataInjector {
val scan = op.getNativeScan
// The driver hashes the common once and ships the hash inside the plan
// (CometNativeScanExec.apply), so no per-task hashing happens here; deriving from the
- // common is only a fallback for plans built without one.
- Some(
+ // common is only a fallback for plans built without one. The plan_id matches the one
+ // CometNativeScanExec.apply folds into sourceKey.
+ val key =
if (scan.hasSourceKeyHash) sourceKey(scan.getCommon.getSource, scan.getSourceKeyHash)
- else sourceKey(scan.getCommon))
+ else sourceKey(scan.getCommon)
+ Some(PlanDataInjector.withPlanId(key, op.getPlanId))
}
/**
- * The key under which a native scan's planning data is stored and looked up: the scan's source
- * plus [[sourceKeyHash]]. `CometNativeScanExec.apply` computes the hash once on the driver and
- * embeds it in the NativeScan proto, so [[getKey]] rebuilds the identical string from the
- * source already carried by the common.
+ * A native scan's content key: the scan's source plus [[sourceKeyHash]]. The planning data is
+ * stored and looked up under this key plus the op's plan_id (PlanDataInjector.withPlanId).
+ * `CometNativeScanExec.apply` computes the hash once on the driver and embeds it in the
+ * NativeScan proto, so [[getKey]] rebuilds the identical string from the source already carried
+ * by the common.
*/
private[comet] def sourceKey(common: OperatorOuterClass.NativeScanCommon): String =
sourceKey(common.getSource, sourceKeyHash(common))
@@ -1263,6 +1313,12 @@ abstract class CometLeafExec extends CometNativeExec with LeafExecNode {
* runtime "not a leaf" fallback.
*/
trait CometScanWithPlanData { self: CometLeafExec =>
+
+ /**
+ * The key for this scan's planning data. It must be unique per scan node within a native plan,
+ * so include the scan op's plan_id via `PlanDataInjector.withPlanId(..., op.getPlanId)`.
+ * findAllPlanData throws when two scans share a key with different data.
+ */
def sourceKey: String
def commonData: Array[Byte]
def perPartitionData: Array[Array[Byte]]
diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala
index 2882e778bc..531bb77894 100644
--- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.comet._
import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.execution.{InSubqueryExec, ReusedSubqueryExec, SparkPlan, SubqueryExec}
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, BroadcastQueryStageExec}
-import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, StringType, StructType, TimestampType}
@@ -4403,6 +4403,68 @@ class CometIcebergNativeSuite
}
}
+ test("storage-partitioned self-join with partially clustered distribution (#6278)") {
+ assume(icebergAvailable, "Iceberg not available")
+
+ withTempIcebergDir { warehouseDir =>
+ withSQLConf(
+ "spark.sql.catalog.spj_cat" -> "org.apache.iceberg.spark.SparkCatalog",
+ "spark.sql.catalog.spj_cat.type" -> "hadoop",
+ "spark.sql.catalog.spj_cat.warehouse" -> warehouseDir.getAbsolutePath,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") {
+
+ // The tiny split size makes every data file its own task, so each bucket spans several
+ // input partitions. Partial clustering then splits one side of the join by file and
+ // replicates the other, giving the two scans different files per partition.
+ spark.sql("""
+ CREATE TABLE spj_cat.db.spj_self (id INT, v STRING) USING iceberg
+ PARTITIONED BY (bucket(4, id))
+ TBLPROPERTIES ('read.split.target-size'='1', 'read.split.open-file-cost'='1',
+ 'format-version'='2')
+ """)
+ spark.sql(
+ "INSERT INTO spj_cat.db.spj_self SELECT CAST(id AS INT), concat('a', id) " +
+ "FROM range(0, 40)")
+ spark.sql(
+ "INSERT INTO spj_cat.db.spj_self SELECT CAST(id AS INT), concat('b', id) " +
+ "FROM range(0, 40, 3)")
+ spark.sql(
+ "INSERT INTO spj_cat.db.spj_self SELECT CAST(id AS INT), concat('c', id) " +
+ "FROM range(0, 12)")
+
+ val query =
+ "SELECT a.id, a.v, b.v FROM spj_cat.db.spj_self a JOIN spj_cat.db.spj_self b " +
+ "ON a.id = b.id"
+
+ for (partiallyClustered <- Seq(false, true)) {
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
+ partiallyClustered.toString,
+ "spark.sql.requireAllClusterKeysForCoPartition" -> "false",
+ "spark.sql.iceberg.planning.preserve-data-grouping" -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(query)
+
+ // Both scans run in one native plan with no exchange between them, which is where
+ // each scan must be handed its own files rather than the other side's.
+ assert(
+ collectIcebergNativeScans(cometPlan).length == 2,
+ s"Expected 2 CometIcebergNativeScanExec. Plan:\n$cometPlan")
+ assert(
+ collect(cometPlan) { case e: ShuffleExchangeLike => e }.isEmpty,
+ s"Expected a storage-partitioned join without a shuffle. Plan:\n$cometPlan")
+ }
+ }
+
+ spark.sql("DROP TABLE spj_cat.db.spj_self")
+ }
+ }
+ }
+
// ---- AQE DPP broadcast reuse tests ----
private def collectIcebergDPPSubqueries(plan: SparkPlan): Seq[SparkPlan] = {
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
index 0083ab9222..06da27766d 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
@@ -29,17 +29,18 @@ import org.apache.spark.SparkException
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
-import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, IsNotNull}
+import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec}
import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec}
import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException
-import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType}
import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus
class CometJoinSuite extends CometTestBase {
@@ -1572,4 +1573,74 @@ class CometJoinSuite extends CometTestBase {
}
}
}
+
+ test("scans of one bucketed table in a single native plan each read their own files") {
+ withTable("bucketed_self", "bucketed_dim") {
+ val rows = for {
+ p <- 1 to 3
+ i <- 0 until 20
+ } yield (i % 10, s"p${p}_$i", p)
+ rows
+ .toDF("k", "v", "p")
+ .write
+ .format("parquet")
+ .partitionBy("p")
+ .bucketBy(4, "k")
+ .saveAsTable("bucketed_self")
+ Seq((1, "one"), (2, "two"), (3, "three"))
+ .toDF("p", "x")
+ .write
+ .format("parquet")
+ .saveAsTable("bucketed_dim")
+
+ // A bucketed sort-merge self-join needs no exchange, so both scans run in one native plan
+ // and must each get their own files, though they read the same table and columns.
+ def assertScansShareOneNativePlan(plan: SparkPlan): Seq[CometNativeScanExec] = {
+ val join = collectFirst(plan) { case j: CometSortMergeJoinExec => j }
+ .getOrElse(fail(s"Expected a native sort-merge join:\n$plan"))
+ val scans = collect(join) { case s: CometNativeScanExec => s }
+ assert(scans.size == 2, s"Expected both scans under the join:\n$plan")
+ assert(
+ collect(join) { case e: ShuffleExchangeLike => e }.isEmpty,
+ s"Expected no exchange between the join and its scans:\n$plan")
+ scans
+ }
+
+ for (adaptive <- Seq(false, true)) {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.BUCKETING_ENABLED.key -> "true",
+ SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ CometConf.COMET_EXEC_SORT_MERGE_JOIN_ENABLED.key -> "true") {
+ // In a full outer join neither side gets an inferred isnotnull data filter, so only
+ // the partition filters tell the two scans apart.
+ val (_, fullOuter) = checkSparkAnswerAndOperator(
+ "SELECT a.k, a.v, b.v FROM (SELECT * FROM bucketed_self WHERE p = 1) a " +
+ "FULL OUTER JOIN (SELECT * FROM bucketed_self WHERE p = 2) b ON a.k = b.k")
+ assertScansShareOneNativePlan(fullOuter)
+
+ // Without filter pushdown the scans carry no data filters, and dynamic partition
+ // pruning picks a different partition for each side. Spark 3.4 only runs this with
+ // AQE off: CometScanRule.transformV1Scan leaves a V1 scan with AQE dynamic pruning in
+ // Spark there (case 2 in CometSpark34AqeDppFallbackRule).
+ if (!adaptive || isSpark35Plus) {
+ withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false") {
+ val (_, pruned) = checkSparkAnswer(
+ "SELECT /*+ BROADCAST(d1), BROADCAST(d2) */ a.k, a.v, b.v " +
+ "FROM bucketed_self a JOIN bucketed_self b ON a.k = b.k " +
+ "JOIN bucketed_dim d1 ON a.p = d1.p JOIN bucketed_dim d2 ON b.p = d2.p " +
+ "WHERE d1.x = 'one' AND d2.x = 'two'")
+ assertScansShareOneNativePlan(pruned).foreach { scan =>
+ assert(
+ scan.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]),
+ s"Expected each scan to be pruned dynamically:\n$pruned")
+ }
+ }
+ }
+ }
+ }
+ }
+ }
}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala
index 0023940334..9bf49125d7 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometScanWithPlanDataSuite.scala
@@ -29,9 +29,12 @@ import scala.jdk.CollectionConverters._
import org.scalatest.funsuite.AnyFunSuite
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan}
+import org.apache.comet.CometRuntimeException
import org.apache.comet.serde.OperatorOuterClass.Operator
/**
@@ -80,6 +83,16 @@ class CometScanWithPlanDataSuite extends AnyFunSuite {
override def perPartitionData: Array[Array[Byte]] = perPart
}
+ /** A non-boundary parent, like a join, that puts two scans in one native plan. */
+ private case class StubParent(left: SparkPlan, right: SparkPlan) extends BinaryExecNode {
+ override def output: Seq[Attribute] = Seq.empty
+ override protected def doExecute(): RDD[InternalRow] =
+ throw new UnsupportedOperationException("not executed")
+ override protected def withNewChildrenInternal(
+ newLeft: SparkPlan,
+ newRight: SparkPlan): StubParent = copy(left = newLeft, right = newRight)
+ }
+
test(
"CometScanWithPlanData defaults: no DPP filters and withDynamicPruningFilters is unsupported") {
val scan = StubScan()
@@ -122,6 +135,33 @@ class CometScanWithPlanDataSuite extends AnyFunSuite {
assert(perPartitionByKey.isEmpty)
}
+ test("findAllPlanData keeps one copy when two scans carry the same data under one key") {
+ val left = StubScan("shared-key", Array[Byte](1, 2), Array(Array[Byte](10), Array[Byte](20)))
+ val right = StubScan("shared-key", Array[Byte](1, 2), Array(Array[Byte](10), Array[Byte](20)))
+
+ val (commonByKey, perPartitionByKey) =
+ PlanDataInjector.findAllPlanData(StubParent(left, right))
+
+ assert(commonByKey.keySet == Set("shared-key"))
+ assert(commonByKey("shared-key") sameElements Array[Byte](1, 2))
+ assert(
+ perPartitionByKey("shared-key").map(_.toSeq).toSeq == Seq(Seq[Byte](10), Seq[Byte](20)))
+ }
+
+ test("findAllPlanData rejects two scans with different data under one key") {
+ // Keeping either side would hand both scans the same files and return wrong rows.
+ val base = StubScan("clash-key", Array[Byte](1), Array(Array[Byte](10), Array[Byte](20)))
+ val otherFiles = base.copy(perPart = Array(Array[Byte](10), Array[Byte](30)))
+ val otherCommon = base.copy(common = Array[Byte](2))
+
+ Seq(otherFiles, otherCommon).foreach { other =>
+ val e = intercept[CometRuntimeException] {
+ PlanDataInjector.findAllPlanData(StubParent(base, other))
+ }
+ assert(e.getMessage.contains("clash-key"), s"message should name the key: ${e.getMessage}")
+ }
+ }
+
test("PlanDataInjector registry contains only built-in injectors on a default build") {
// No contrib ships a META-INF/services/...PlanDataInjector on the default test classpath, so
// ServiceLoader discovery adds nothing and the registry is exactly the two built-ins.
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
index 9b93421d7d..bb165f5834 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/PlanDataInjectorSuite.scala
@@ -30,14 +30,17 @@ import org.apache.comet.serde.OperatorOuterClass.Operator
class PlanDataInjectorSuite extends AnyFunSuite {
/** Builds an un-injected IcebergScan operator: hasCommon, zero file_scan_tasks. */
- private def icebergScanOp(metadataLocation: String, scanHashCode: Int): Operator = {
+ private def icebergScanOp(
+ metadataLocation: String,
+ scanHashCode: Int,
+ planId: Int = 0): Operator = {
val common = OperatorOuterClass.IcebergScanCommon
.newBuilder()
.setMetadataLocation(metadataLocation)
.setScanHashCode(scanHashCode)
.build()
val icebergScan = OperatorOuterClass.IcebergScan.newBuilder().setCommon(common).build()
- Operator.newBuilder().setIcebergScan(icebergScan).build()
+ Operator.newBuilder().setPlanId(planId).setIcebergScan(icebergScan).build()
}
/** Serialized (commonBytes, partitionBytes) a real CometIcebergNativeScanExec would produce. */
@@ -144,21 +147,36 @@ class PlanDataInjectorSuite extends AnyFunSuite {
assert(IcebergPlanDataInjector.getKey(targetOp) != IcebergPlanDataInjector.getKey(sourceOp))
}
- test("two Iceberg scans of the same table with equal scan_hash_code get the same key") {
- val opA = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111)
- val opB = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111)
+ test("an Iceberg scan with equal scan_hash_code and plan_id gets the same key") {
+ // The same scan node's op, for example on the driver and again in a task.
+ val opA = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111, planId = 7)
+ val opB = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111, planId = 7)
assert(IcebergPlanDataInjector.getKey(opA) == IcebergPlanDataInjector.getKey(opB))
}
+ test("two Iceberg scans with equal scan_hash_code but different plan_id get distinct keys") {
+ // The two sides of a storage-partitioned self-join share table, snapshot and read schema,
+ // yet read different files per partition.
+ val location = "s3://table/metadata/v1.json"
+ val opA = icebergScanOp(location, scanHashCode = 111, planId = 1)
+ val opB = icebergScanOp(location, scanHashCode = 111, planId = 2)
+
+ assert(IcebergPlanDataInjector.getKey(opA) != IcebergPlanDataInjector.getKey(opB))
+ }
+
/**
* Builds an un-injected NativeScan operator the way the driver ships it: hasCommon, no
* file_partition, source_key_hash embedded (see CometNativeScanExec.apply).
*/
- private def nativeScanOp(source: String, columnNames: Seq[String]): Operator = {
+ private def nativeScanOp(
+ source: String,
+ columnNames: Seq[String],
+ planId: Int = 0): Operator = {
val common = nativeScanCommon(source, columnNames)
Operator
.newBuilder()
+ .setPlanId(planId)
.setNativeScan(
OperatorOuterClass.NativeScan
.newBuilder()
@@ -622,11 +640,38 @@ class PlanDataInjectorSuite extends AnyFunSuite {
assert(injectedB.getNativeScan.getCommon == commonB)
}
+ test("NativeScan scans with the same content but different plan_id read their own files") {
+ // Two scans of one table and schema, such as different partitions of a bucketed self-join,
+ // sit under one parent. Each must be injected with its own file list.
+ val scanA = nativeScanOp("file:///self-join-tbl", Seq("k", "v"), planId = 1)
+ val scanB = nativeScanOp("file:///self-join-tbl", Seq("k", "v"), planId = 2)
+ val keyA = NativeScanPlanDataInjector.getKey(scanA).get
+ val keyB = NativeScanPlanDataInjector.getKey(scanB).get
+ assert(keyA != keyB)
+ val sameNode = nativeScanOp("file:///self-join-tbl", Seq("k", "v"), planId = 1)
+ assert(NativeScanPlanDataInjector.getKey(sameNode).contains(keyA))
+
+ val common = scanA.getNativeScan.getCommon.toByteArray
+ val root = Operator.newBuilder().setPlanId(1).addChildren(scanA).addChildren(scanB).build()
+ val injected = PlanDataInjector.injectPlanData(
+ parseBasePlan(root.toByteArray),
+ Map(keyA -> common, keyB -> common),
+ Map(
+ keyA -> nativeScanPartitionBytes("p=1.parquet"),
+ keyB -> nativeScanPartitionBytes("p=2.parquet")))
+
+ def filePath(scan: Operator): String =
+ scan.getNativeScan.getFilePartition.getPartitionedFile(0).getFilePath
+ assert(filePath(injected.getChildren(0)) == "p=1.parquet")
+ assert(filePath(injected.getChildren(1)) == "p=2.parquet")
+ }
+
test("NativeScan getKey rebuilds the key from the source and the transported hash") {
// Only the hash travels in the plan; the source is already in the common next to it.
val common = nativeScanCommon("file:///transported-tbl", Seq("id", "v"))
val op = Operator
.newBuilder()
+ .setPlanId(5)
.setNativeScan(
OperatorOuterClass.NativeScan
.newBuilder()
@@ -634,7 +679,7 @@ class PlanDataInjectorSuite extends AnyFunSuite {
.setSourceKeyHash(1234))
.build()
- assert(NativeScanPlanDataInjector.getKey(op).contains("file:///transported-tbl_1234"))
+ assert(NativeScanPlanDataInjector.getKey(op).contains("file:///transported-tbl_1234#5"))
assert(
NativeScanPlanDataInjector.sourceKey(common) ==
s"file:///transported-tbl_${NativeScanPlanDataInjector.sourceKeyHash(common)}",
@@ -645,21 +690,23 @@ class PlanDataInjectorSuite extends AnyFunSuite {
val common = nativeScanCommon("file:///zero-hash-tbl", Seq("id"))
val op = Operator
.newBuilder()
+ .setPlanId(5)
.setNativeScan(
OperatorOuterClass.NativeScan.newBuilder().setCommon(common).setSourceKeyHash(0))
.build()
- assert(NativeScanPlanDataInjector.getKey(op).contains("file:///zero-hash-tbl_0"))
+ assert(NativeScanPlanDataInjector.getKey(op).contains("file:///zero-hash-tbl_0#5"))
}
test("NativeScan getKey derives the key only when the plan carries none") {
val common = nativeScanCommon("file:///fallback-tbl", Seq("id", "v", "w"))
val op = Operator
.newBuilder()
+ .setPlanId(5)
.setNativeScan(OperatorOuterClass.NativeScan.newBuilder().setCommon(common))
.build()
- val derived = NativeScanPlanDataInjector.sourceKey(common)
+ val derived = PlanDataInjector.withPlanId(NativeScanPlanDataInjector.sourceKey(common), 5)
assert(NativeScanPlanDataInjector.getKey(op).contains(derived))
}
@@ -669,10 +716,11 @@ class PlanDataInjectorSuite extends AnyFunSuite {
val common = nativeScanCommon("file:///shuffle-tbl", Seq("id"))
val op = Operator
.newBuilder()
+ .setPlanId(9)
.setNativeScan(
OperatorOuterClass.NativeScan.newBuilder().setCommon(common).setSourceKeyHash(77))
.build()
- val key = "file:///shuffle-tbl_77"
+ val key = "file:///shuffle-tbl_77#9"
val injected = PlanDataInjector.injectPlanDataForShuffle(
424242,
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
index 7c35b974ee..de3cc5de31 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
@@ -871,7 +871,6 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
/** Caches one prepared scan under `shuffleId` as a native map task would. */
private def prepareShuffleScan(shuffleId: Int, source: String): String = {
val common = OperatorOuterClass.NativeScanCommon.newBuilder().setSource(source).build()
- val key = NativeScanPlanDataInjector.sourceKey(common)
val scanOp = OperatorOuterClass.Operator
.newBuilder()
.setNativeScan(
@@ -880,6 +879,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
.setCommon(common)
.setSourceKeyHash(NativeScanPlanDataInjector.sourceKeyHash(common)))
.build()
+ val key = NativeScanPlanDataInjector.getKey(scanOp).get
val partition = OperatorOuterClass.NativeScan
.newBuilder()
.setFilePartition(OperatorOuterClass.SparkFilePartition.newBuilder())
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org