From 867a767e4ff00370ad959d1fa358b3e637dd6518 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Fri, 31 Jul 2026 16:59:55 -0400 Subject: [PATCH] fix disambiguate scans with same metadata location --- native/proto/src/proto/operator.proto | 6 ++ .../operator/CometIcebergNativeScan.scala | 8 +- .../comet/CometIcebergNativeScanExec.scala | 8 +- .../apache/spark/sql/comet/operators.scala | 20 +++-- .../sql/comet/PlanDataInjectorSuite.scala | 90 +++++++++++++++++++ 5 files changed, 122 insertions(+), 10 deletions(-) diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 1633429986..dae7eed9fb 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -283,6 +283,12 @@ message IcebergScanCommon { // delete file here (rather than embedding copies per delete_files_pool set) keeps the broadcast // common message from growing with the number of references. repeated IcebergDeleteFile delete_file_pool = 14; + + // 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. + int32 scan_hash_code = 15; } message IcebergScan { diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index c2b792c36c..f61f1368b2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -799,10 +799,12 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val icebergScanBuilder = OperatorOuterClass.IcebergScan.newBuilder() val commonBuilder = OperatorOuterClass.IcebergScanCommon.newBuilder() - // Only set metadata_location - used for matching in PlanDataInjector. - // All other fields (catalog_properties, required_schema, pools) are set by - // serializePartitions() at execution time, so setting them here would be wasted work. + // metadata_location and scan_hash_code are set here for PlanDataInjector's key (see the + // scan_hash_code field comment in operator.proto). Everything else (catalog_properties, + // required_schema, pools) is set by serializePartitions() at execution time, so setting it + // here would be wasted work. commonBuilder.setMetadataLocation(metadata.metadataLocation) + commonBuilder.setScanHashCode(scan.scan.hashCode()) icebergScanBuilder.setCommon(commonBuilder.build()) // partition field intentionally empty - will be populated at execution time 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 95cf2fc9fd..3182a9daeb 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 @@ -220,11 +220,15 @@ case class CometIcebergNativeScanExec( override def doExecuteColumnar(): RDD[ColumnarBatch] = { 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). + val injectorKey = IcebergPlanDataInjector.getKey(nativeOp).getOrElse(metadataLocation) new CometExecRDD( sparkContext, inputRDDs = Seq.empty, - commonByKey = Map(metadataLocation -> commonData), - perPartitionByKey = Map(metadataLocation -> perPartitionData), + commonByKey = Map(injectorKey -> commonData), + perPartitionByKey = Map(injectorKey -> perPartitionData), serializedPlan = serializedPlan, defaultNumPartitions = perPartitionData.length, numOutputCols = output.length, 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 fcd48d69b1..e6ab81e024 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 @@ -203,9 +203,17 @@ private[comet] object PlanDataInjector extends Logging { // the serializedPartitionData lazy val initializer (a known deadlock surface). iceberg.ensureSubqueriesResolved() if (iceberg.commonData.nonEmpty && iceberg.perPartitionData.nonEmpty) { - ( - Map(iceberg.metadataLocation -> iceberg.commonData), - Map(iceberg.metadataLocation -> iceberg.perPartitionData)) + // 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). + IcebergPlanDataInjector.getKey(iceberg.nativeOp) match { + case Some(key) => + (Map(key -> iceberg.commonData), Map(key -> iceberg.perPartitionData)) + case None => + (Map.empty, Map.empty) + } } else { (Map.empty, Map.empty) } @@ -274,8 +282,10 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector { op.getIcebergScan.getFileScanTasksCount == 0 && op.getIcebergScan.hasCommon - override def getKey(op: Operator): Option[String] = - Some(op.getIcebergScan.getCommon.getMetadataLocation) + override def getKey(op: Operator): Option[String] = { + val common = op.getIcebergScan.getCommon + Some(s"${common.getMetadataLocation}_${common.getScanHashCode}") + } override def inject( op: Operator, 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 601ce9a7e7..c8fed67cb2 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 @@ -21,10 +21,48 @@ package org.apache.spark.sql.comet import org.scalatest.funsuite.AnyFunSuite +import org.apache.comet.serde.OperatorOuterClass 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 = { + val common = OperatorOuterClass.IcebergScanCommon + .newBuilder() + .setMetadataLocation(metadataLocation) + .setScanHashCode(scanHashCode) + .build() + val icebergScan = OperatorOuterClass.IcebergScan.newBuilder().setCommon(common).build() + Operator.newBuilder().setIcebergScan(icebergScan).build() + } + + /** Serialized (commonBytes, partitionBytes) a real CometIcebergNativeScanExec would produce. */ + private def icebergPlanData( + metadataLocation: String, + scanHashCode: Int, + columnNames: Seq[String], + dataFilePath: String): (Array[Byte], Array[Byte]) = { + val commonBuilder = OperatorOuterClass.IcebergScanCommon + .newBuilder() + .setMetadataLocation(metadataLocation) + .setScanHashCode(scanHashCode) + columnNames.foreach { name => + commonBuilder.addRequiredSchema( + OperatorOuterClass.SparkStructField.newBuilder().setName(name).setNullable(true).build()) + } + val commonBytes = commonBuilder.build().toByteArray + + val partitionBytes = OperatorOuterClass.IcebergScan + .newBuilder() + .addFileScanTasks( + OperatorOuterClass.IcebergFileScanTask.newBuilder().setDataFilePath(dataFilePath).build()) + .build() + .toByteArray + + (commonBytes, partitionBytes) + } + test("injectPlanData leaves a non-scan operator tree unchanged") { // An operator with no injectable scan (here, an empty op_struct, but the same holds for // Filter/Projection/etc.) must pass through untouched. This exercises the O(1) @@ -50,4 +88,56 @@ class PlanDataInjectorSuite extends AnyFunSuite { assert(IcebergPlanDataInjector.opStructCase == Operator.OpStructCase.ICEBERG_SCAN) assert(NativeScanPlanDataInjector.opStructCase == Operator.OpStructCase.NATIVE_SCAN) } + + test("two Iceberg scans of the same table with different scan_hash_code get distinct keys") { + val targetOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111) + val sourceOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 222) + + 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) + + assert(IcebergPlanDataInjector.getKey(opA) == IcebergPlanDataInjector.getKey(opB)) + } + + test( + "self-join: scans sharing a metadataLocation but differing scan_hash_code inject their " + + "own data, not each other's") { + val targetOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 111) + val sourceOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 222) + + val (targetCommon, targetPartition) = + icebergPlanData( + "s3://table/metadata/v1.json", + scanHashCode = 111, + columnNames = Seq("id", "v", "_file", "_pos"), + dataFilePath = "target.parquet") + val (sourceCommon, sourcePartition) = + icebergPlanData( + "s3://table/metadata/v1.json", + scanHashCode = 222, + columnNames = Seq("id", "v"), + dataFilePath = "source.parquet") + + val targetKey = IcebergPlanDataInjector.getKey(targetOp).get + val sourceKey = IcebergPlanDataInjector.getKey(sourceOp).get + val commonByKey = Map(targetKey -> targetCommon, sourceKey -> sourceCommon) + val partitionByKey = Map(targetKey -> targetPartition, sourceKey -> sourcePartition) + + val injectedTarget = PlanDataInjector.injectPlanData(targetOp, commonByKey, partitionByKey) + val injectedSource = PlanDataInjector.injectPlanData(sourceOp, commonByKey, partitionByKey) + + assert( + injectedTarget.getIcebergScan.getCommon.getRequiredSchemaList + .get(0) + .getName == "id") + assert(injectedTarget.getIcebergScan.getCommon.getRequiredSchemaCount == 4) + assert(injectedTarget.getIcebergScan.getFileScanTasks(0).getDataFilePath == "target.parquet") + + assert(injectedSource.getIcebergScan.getCommon.getRequiredSchemaCount == 2) + assert(injectedSource.getIcebergScan.getFileScanTasks(0).getDataFilePath == "source.parquet") + } }