diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala index bcd1d82f906..77b64128bb9 100644 --- a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala @@ -210,7 +210,7 @@ case class CometDeltaNativeScanExec( sparkContext, Seq.empty, Map(sourceKey -> commonData), - Map(sourceKey -> perPartitionData), + Map(sourceKey -> PartitionPlanData.wrap(perPartitionData)), serializedPlan, PlanDataInjector.planFingerprint(serializedPlan), perPartitionData.length, diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 5e11b460149..ea1f517a85a 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -234,6 +234,15 @@ system user must be that user. Otherwise the scan falls back with a reason that To read HDFS tables with Spark's Iceberg reader and keep the native scan for other storage, set `spark.comet.scan.icebergNative.hdfs.enabled=false`. Writes to HDFS use Iceberg's own writer. +### Plan data per task + +The native reader sends each task only the delete files, partition values and residual filters +that its own files need, the way Spark's Iceberg reader does. A delete file or residual filter +that many partitions use, such as an unpartitioned equality delete or a filter on a +non-partition column, is serialized once and shared by those partitions. The driver keeps one +copy of it, and a coalesce without a shuffle (`coalesce(n)`, or a write that coalesces its +input), which combines several input partitions into one task, sends that task one copy. + ### Current limitations The following scenarios will fall back to the JVM Iceberg reader: diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 90c01a7511d..7237e9547e0 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -389,7 +389,12 @@ message IcebergScanCommon { // Schema to read repeated SparkStructField required_schema = 3; - // Deduplication pools (must contain all entries for cross-partition deduplication) + // Deduplication pools. In the common block every task gets, the per-file pools + // (partition_data_pool, delete_files_pool, residual_pool, delete_file_pool, + // delete_file_path_pool) are empty: each partition's IcebergScan.common holds the entries its + // tasks use, and the JVM merges the two before the plan reaches native. The driver ships a + // partition's delete files and their paths apart from its message, shared with the other + // partitions that use them, and the executor JVM adds them to the partition's pools. repeated string schema_pool = 4; // Index-aligned with partition_spec_pool: entry i is the StructType JSON for the spec at // partition_spec_pool[i]. Native uses this pairing to recover each entry's spec_id when @@ -415,8 +420,8 @@ message IcebergScanCommon { // Flat pool of unique delete files, referenced by index from DeleteFileList. Iceberg's default // partition delete granularity has one delete file apply to many data files, so interning each - // delete file here (rather than embedding copies per delete_files_pool set) keeps the broadcast - // common message from growing with the number of references. + // delete file here (rather than embedding copies per delete_files_pool set) keeps a + // partition's pools from growing with the number of references. repeated IcebergDeleteFile delete_file_pool = 14; // Flat pool of unique delete-file paths, referenced by index from IcebergDeleteFile. A V3 @@ -433,7 +438,8 @@ message IcebergScanCommon { } message IcebergScan { - // Common data shared across partitions (pools, metadata, catalog props) + // Common data shared across partitions (pools, metadata, catalog props). In the per-partition + // message the driver ships, it holds only that partition's per-file pool entries. IcebergScanCommon common = 1; // Single partition's file scan tasks 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 ef70389caa1..8d08e06d334 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 @@ -36,7 +36,7 @@ import org.apache.hadoop.fs.Path import org.apache.hadoop.hdfs.client.HdfsClientConfigKeys import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeExec} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeExec, IcebergPartitionPlanData, PartitionPlanData} import org.apache.spark.sql.comet.shims.ShimDataSourceRDDPartition import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceRDD, DataSourceRDDPartition} import org.apache.spark.sql.types._ @@ -310,8 +310,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit task: Any, contentFileClass: Class[_], fileScanTaskClass: Class[_], - deleteFileClass: Class[_], - internPath: String => Int): Seq[OperatorOuterClass.IcebergDeleteFile] = { + deleteFileClass: Class[_]): Seq[(String, OperatorOuterClass.IcebergDeleteFile)] = { try { // keyMetadata() is declared on ContentFile; present across all supported Iceberg versions. val keyMetadataMethod = IcebergReflection.getMethod(contentFileClass, "keyMetadata") @@ -319,13 +318,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) deletes.asScala - .map( - serializeDeleteFile( - _, - contentFileClass, - deleteFileClass, - keyMetadataMethod, - internPath)) + .map(serializeDeleteFile(_, contentFileClass, deleteFileClass, keyMetadataMethod)) .toSeq } catch { case e: Exception => @@ -338,7 +331,9 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } /** - * Serializes a single Iceberg DeleteFile to protobuf. + * Serializes a single Iceberg DeleteFile to protobuf. Returns the file's path apart from the + * message, which leaves `file_path_idx` unset: the index depends on which partition holds the + * file, and is set when the partition's message is built (see IcebergPartitionPlanData). * * Required delete metadata must never fall back to guessed values after native execution has * been selected. Missing accessors and invocation failures therefore propagate. @@ -347,8 +342,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit deleteFile: Any, contentFileClass: Class[_], deleteFileClass: Class[_], - keyMetadataMethod: java.lang.reflect.Method, - internPath: String => Int): OperatorOuterClass.IcebergDeleteFile = { + keyMetadataMethod: java.lang.reflect.Method) + : (String, OperatorOuterClass.IcebergDeleteFile) = { val deletePath = IcebergReflection .extractFileLocation(contentFileClass, deleteFile) .getOrElse( @@ -357,7 +352,6 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit "ContentFile -- cannot extract delete file path from FileScanTask")) val deleteBuilder = OperatorOuterClass.IcebergDeleteFile.newBuilder() - deleteBuilder.setFilePathIdx(internPath(deletePath)) val contentMethod = IcebergReflection.getMethod(deleteFileClass, "content") val contentType = contentMethod.invoke(deleteFile).toString @@ -413,7 +407,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // Unencrypted delete files leave the field unset. keyMetadataBytes(keyMetadataMethod, deleteFile).foreach(deleteBuilder.setKeyMetadata) - deleteBuilder.build() + (deletePath, deleteBuilder.build()) } /** @@ -431,6 +425,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit taskBuilder: OperatorOuterClass.IcebergFileScanTask.Builder, commonBuilder: OperatorOuterClass.IcebergScanCommon.Builder, partitionSpecToPoolIndex: mutable.HashMap[String, Int], + partitionPools: OperatorOuterClass.IcebergScanCommon.Builder, partitionDataToPoolIndex: mutable.HashMap[String, Int]): Unit = { try { val specMethod = IcebergReflection.getMethod(fileScanTaskClass, "spec") @@ -590,7 +585,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val partitionDataIdx = partitionDataToPoolIndex.getOrElseUpdate( partitionDataKey, { val idx = partitionDataToPoolIndex.size - commonBuilder.addPartitionDataPool(partitionDataProto) + partitionPools.addPartitionDataPool(partitionDataProto) idx }) val _ = taskBuilder.setPartitionDataIdx(partitionDataIdx) @@ -1187,18 +1182,24 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * Called after doPrepare() has resolved DPP subqueries. Builds pools and per-partition data in * one pass from the DPP-filtered partitions. * - * The result splits into a common block plus one block per Spark partition. The common block - * holds everything shared across partitions and is broadcast once per executor in the stage - * task binary; each per-partition block ships with its task. Protobuf serializes by value and - * shares nothing between repeated fields, so shared values (schemas, specs, partition data, - * residuals, delete files) are interned into pools in the common block and referenced by index. - * Pooling also keeps a message under protobuf's 2 GiB limit (getSerializedSize returns int and - * wraps past it), which matters for anything that scales with the number of tasks. + * The result splits into a common block plus one block per Spark partition. Protobuf serializes + * by value and shares nothing between repeated fields, so repeated values are interned into + * pools and referenced by index. Pooling also keeps a message under protobuf's 2 GiB limit + * (getSerializedSize returns int and wraps past it), which matters for anything that scales + * with the number of tasks. + * + * Every task deserializes the common block, so it holds only what does not grow with the files: + * schemas, partition specs and types, name mappings and projected field ids. Partition values, + * residuals, delete files and their paths and sets are interned per partition, so each task + * gets only the entries its own files use. Keeping them in the common block made each task's + * plan as large as the whole scan's delete metadata. * * Delete files use two levels: a flat pool of unique files, plus a per-task list of indices * into it. One delete file applies to many data files under Iceberg's default partition delete - * granularity, so this stores each file once and references it many times, as Iceberg's own - * DeleteFileIndex does. + * granularity, so each partition stores it once and references it many times, as Iceberg's own + * DeleteFileIndex does. Across partitions, each delete file and residual is serialized once and + * shared by reference (see IcebergPartitionPlanData), so a coalesced task carries it once and + * the driver keeps one copy. * * @param scanExec * The BatchScanExec whose inputRDD contains the DPP-filtered partitions @@ -1207,40 +1208,94 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * @param metadata * Pre-extracted Iceberg metadata from CometScanRule * @return - * Tuple of (commonBytes, perPartitionBytes) for native execution + * Tuple of (commonBytes, perPartitionData) for native execution */ def serializePartitions( scanExec: BatchScanExec, output: Seq[Attribute], - metadata: CometIcebergNativeScanMetadata): (Array[Byte], Array[Array[Byte]]) = { + metadata: CometIcebergNativeScanMetadata): (Array[Byte], Array[PartitionPlanData]) = { val commonBuilder = OperatorOuterClass.IcebergScanCommon.newBuilder() - // Deduplication structures - map unique values to pool indices + // Deduplication structures for the pools in the common block - map unique values to indices val partitionSpecToPoolIndex = mutable.HashMap[String, Int]() val nameMappingToPoolIndex = mutable.HashMap[String, Int]() - val partitionDataToPoolIndex = mutable.HashMap[String, Int]() - // Individual delete files are interned into a flat pool; deleteFilesToPoolIndex then dedups - // the per-task sets as lists of indices into it, so a delete file that applies to many data - // files (Iceberg's default partition delete granularity) is serialized once rather than once - // per referencing FileScanTask. Keyed on the whole message rather than the path: V3 deletion - // vectors for different data files share one Puffin file and differ only by content offset, - // so a path key would collapse them and drop every vector but the first. - val deleteFileToPoolIndex = mutable.HashMap[OperatorOuterClass.IcebergDeleteFile, Int]() - val deleteFilesToPoolIndex = - mutable.HashMap[Seq[Int], Int]() - // Delete-file paths are interned separately from the delete files themselves: every deletion - // vector in a commit lives in one Puffin file, so one path is shared by as many pool entries - // as there are data files. - val deleteFilePathToPoolIndex = mutable.HashMap[String, Int]() - def internDeleteFilePath(path: String): Int = - deleteFilePathToPoolIndex.getOrElseUpdate( - path, { - val idx = deleteFilePathToPoolIndex.size - commonBuilder.addDeleteFilePathPool(path) - idx - }) - val residualToPoolIndex = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Int]() + + // A delete file is often used by many partitions, so it is serialized once for the whole + // scan, and every partition that uses it holds the same array and the same path string (see + // IcebergPartitionPlanData). Keyed on the path and the whole message: V3 deletion vectors for + // different data files share one Puffin file and differ only by content offset, so a path key + // would collapse them and drop every vector but the first. + val sharedDeleteFiles = + mutable.HashMap[(String, OperatorOuterClass.IcebergDeleteFile), (String, Array[Byte])]() + // Every deletion vector in a commit lives in one Puffin file, so one path string is shared by + // as many delete files as there are data files. + val sharedDeleteFilePaths = mutable.HashMap[String, String]() + var sharedDeleteFileBytes = 0L + // A residual is often the scan's whole filter, used by every partition, so it is shared the + // same way. + val sharedResiduals = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Array[Byte]]() + var sharedResidualBytes = 0L + + // The per-file pools of one partition, so a task never holds entries that only other + // partitions use. Within a partition each entry is stored once. + final class PartitionPools { + val builder: OperatorOuterClass.IcebergScanCommon.Builder = + OperatorOuterClass.IcebergScanCommon.newBuilder() + val partitionData = mutable.HashMap[String, Int]() + // The partition's delete files, in the order its delete-file sets index them. The sets are + // lists of indices into them, so a delete file that applies to many data files (Iceberg's + // default partition delete granularity) is stored once rather than once per FileScanTask. + val deleteFiles = mutable.HashMap[(String, OperatorOuterClass.IcebergDeleteFile), Int]() + val deleteFileBytes = mutable.ArrayBuffer[Array[Byte]]() + val deleteFilePaths = mutable.ArrayBuffer[String]() + val deleteFileSets = mutable.HashMap[Seq[Int], Int]() + // The partition's residuals, in the order its tasks index them. + val residuals = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Int]() + val residualBytes = mutable.ArrayBuffer[Array[Byte]]() + + def internDeleteFile(file: (String, OperatorOuterClass.IcebergDeleteFile)): Int = + deleteFiles.getOrElseUpdate( + file, { + val (path, bytes) = sharedDeleteFiles.getOrElseUpdate( + file, { + val bytes = file._2.toByteArray + sharedDeleteFileBytes += bytes.length + (sharedDeleteFilePaths.getOrElseUpdate(file._1, file._1), bytes) + }) + deleteFileBytes += bytes + deleteFilePaths += path + deleteFileBytes.size - 1 + }) + + def internResidual(residual: OperatorOuterClass.IcebergPredicate): Int = + residuals.getOrElseUpdate( + residual, { + residualBytes += sharedResiduals.getOrElseUpdate( + residual, { + val bytes = residual.toByteArray + sharedResidualBytes += bytes.length + bytes + }) + residualBytes.size - 1 + }) + + def build(scan: OperatorOuterClass.IcebergScan.Builder): IcebergPartitionPlanData = { + if (partitionData.nonEmpty || deleteFileSets.nonEmpty) { + scan.setCommon(builder) + } + new IcebergPartitionPlanData( + scan.build().toByteArray, + deleteFileBytes.toArray, + deleteFilePaths.toArray, + residualBytes.toArray) + } + } + // Entries in the per-file pools, summed over partitions, for the log. + var partitionDataEntries = 0 + var deleteFileEntries = 0 + var deleteFileSetEntries = 0 + var residualEntries = 0 // Columns whose Iceberg type iceberg-rust cannot use for page-index pruning; residual // predicates over them are dropped (see icebergExprToProto). Computed once from the full table // schema so a filter column projected out of the scan output is still recognized. @@ -1274,7 +1329,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } commonBuilder.addProjectFieldIdsPool(projectFieldIds.build()) - val perPartitionBuilders = mutable.ArrayBuffer[OperatorOuterClass.IcebergScan]() + val perPartition = mutable.ArrayBuffer[IcebergPartitionPlanData]() var totalTasks = 0 @@ -1327,6 +1382,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val partitions = rdd.partitions partitions.foreach { partition => val partitionBuilder = OperatorOuterClass.IcebergScan.newBuilder() + val pools = new PartitionPools val inputPartitions = ShimDataSourceRDDPartition .inputPartitions(partition.asInstanceOf[DataSourceRDDPartition]) @@ -1387,25 +1443,17 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit task, contentFileClass, fileScanTaskClass, - deleteFileClass, - internDeleteFilePath) + deleteFileClass) if (deleteFilesList.nonEmpty) { - // Intern each delete file into the flat pool, then dedup this task's set as the - // resulting list of pool indices. - val deleteFileIndices = deleteFilesList.map { df => - deleteFileToPoolIndex.getOrElseUpdate( - df, { - val idx = deleteFileToPoolIndex.size - commonBuilder.addDeleteFilePool(df) - idx - }) - } - val deleteFilesIdx = deleteFilesToPoolIndex.getOrElseUpdate( + // Intern each delete file into the partition's delete files, then dedup this + // task's set as the resulting list of indices. + val deleteFileIndices = deleteFilesList.map(pools.internDeleteFile) + val deleteFilesIdx = pools.deleteFileSets.getOrElseUpdate( deleteFileIndices, { - val idx = deleteFilesToPoolIndex.size + val idx = pools.deleteFileSets.size val listBuilder = OperatorOuterClass.DeleteFileList.newBuilder() deleteFileIndices.foreach(idx => listBuilder.addDeleteFileIndices(idx)) - commonBuilder.addDeleteFilesPool(listBuilder.build()) + pools.builder.addDeleteFilesPool(listBuilder.build()) idx }) taskBuilder.setDeleteFilesIdx(deleteFilesIdx) @@ -1418,13 +1466,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit pageIndexUnsupportedColumns) residualExprOpt.foreach { residual => - val residualIdx = residualToPoolIndex.getOrElseUpdate( - residual, { - val idx = residualToPoolIndex.size - commonBuilder.addResidualPool(residual) - idx - }) - taskBuilder.setResidualIdx(residualIdx) + taskBuilder.setResidualIdx(pools.internResidual(residual)) } serializePartitionData( @@ -1434,7 +1476,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit taskBuilder, commonBuilder, partitionSpecToPoolIndex, - partitionDataToPoolIndex) + pools.builder, + pools.partitionData) metadata.nameMapping.foreach { nm => val nmIdx = nameMappingToPoolIndex.getOrElseUpdate( @@ -1451,7 +1494,11 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } - perPartitionBuilders += partitionBuilder.build() + partitionDataEntries += pools.partitionData.size + deleteFileEntries += pools.deleteFiles.size + deleteFileSetEntries += pools.deleteFileSets.size + residualEntries += pools.residuals.size + perPartition += pools.build(partitionBuilder) } case other if other.getClass.getName == "org.apache.spark.rdd.ParallelCollectionRDD" => // Spark 3.4's BatchScanExec.inputRDD returns sparkContext.parallelize(empty, 1) when @@ -1470,7 +1517,12 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit "BatchScanExec.inputRDD is ParallelCollectionRDD (DPP pruned all partitions); " + "preserving empty execution partitions") other.partitions.foreach { _ => - perPartitionBuilders += OperatorOuterClass.IcebergScan.getDefaultInstance + perPartition += + new IcebergPartitionPlanData( + Array.emptyByteArray, + Array.empty, + Array.empty, + Array.empty) } case other => throw new IllegalStateException( @@ -1484,11 +1536,10 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit partitionSpecToPoolIndex.size, nameMappingToPoolIndex.size, commonBuilder.getProjectFieldIdsPoolCount, - partitionDataToPoolIndex.size, - deleteFileToPoolIndex.size, - deleteFilesToPoolIndex.size, - deleteFilePathToPoolIndex.size, - residualToPoolIndex.size) + partitionDataEntries, + deleteFileEntries, + deleteFileSetEntries, + residualEntries) val avgDedup = if (totalTasks == 0) { "0.0" @@ -1530,30 +1581,9 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit "project_field_ids", commonBuilder.getProjectFieldIdsPoolCount, sumSizes( - commonBuilder.getProjectFieldIdsPoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "partition_data", - commonBuilder.getPartitionDataPoolCount, - sumSizes( - commonBuilder.getPartitionDataPoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "delete_file", - commonBuilder.getDeleteFilePoolCount, - sumSizes(commonBuilder.getDeleteFilePoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "delete_file_path", - commonBuilder.getDeleteFilePathPoolCount, - sumStrBytes(commonBuilder.getDeleteFilePathPoolList.asScala)), - ( - "delete_files_set", - commonBuilder.getDeleteFilesPoolCount, - sumSizes(commonBuilder.getDeleteFilesPoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "residual", - commonBuilder.getResidualPoolCount, - sumSizes(commonBuilder.getResidualPoolList.asScala.iterator.map(_.getSerializedSize)))) + commonBuilder.getProjectFieldIdsPoolList.asScala.iterator.map(_.getSerializedSize)))) - val perPartitionSizes = perPartitionBuilders.map(_.getSerializedSize.toLong) + val perPartitionSizes = perPartition.map(_.scan.length.toLong) val perPartitionTotal = perPartitionSizes.sum val perPartitionMax = if (perPartitionSizes.isEmpty) 0L else perPartitionSizes.max @@ -1565,13 +1595,19 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit .mkString(", ") + s"; total=${commonPoolBytes.map(_._3).sum} bytes") logInfo( - s" Per-partition messages: ${perPartitionBuilders.size}, " + + " Per-partition pool entries, summed over partitions: " + + s"partition_data=$partitionDataEntries, delete_file=$deleteFileEntries, " + + s"delete_files_set=$deleteFileSetEntries, residual=$residualEntries") + logInfo( + " Held once and shared by the partitions that use them: " + + s"${sharedDeleteFiles.size} delete files, ${sharedDeleteFilePaths.size} paths, " + + s"$sharedDeleteFileBytes bytes; ${sharedResiduals.size} residuals, " + + s"$sharedResidualBytes bytes") + logInfo( + s" Per-partition messages, without delete files and residuals: ${perPartition.size}, " + s"total=$perPartitionTotal bytes, max=$perPartitionMax bytes") - val commonBytes = commonBuilder.build().toByteArray - val perPartitionBytes = perPartitionBuilders.map(_.toByteArray).toArray - - (commonBytes, perPartitionBytes) + (commonBuilder.build().toByteArray, perPartition.toArray[PartitionPlanData]) } override def createExec(nativeOp: Operator, op: CometBatchScanExec): CometNativeExec = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index c903b7f70b0..5ac61b4d7fb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -37,7 +37,7 @@ import org.apache.comet.{CometExecIterator, CometRuntimeException, CometShuffleB private[spark] class CometExecPartition( override val index: Int, val inputPartitions: Array[Partition], - val planDataByKey: Map[String, Array[Byte]], + val planDataByKey: Map[String, PartitionPlanData], val filePaths: Seq[String] = Seq.empty) extends Partition @@ -58,7 +58,7 @@ private[spark] class CometExecRDD( sc: SparkContext, var inputRDDs: Seq[RDD[_]], commonByKey: Map[String, Array[Byte]], - @transient perPartitionByKey: Map[String, Array[Array[Byte]]], + @transient perPartitionByKey: Map[String, Array[PartitionPlanData]], serializedPlan: Array[Byte], planFingerprint: Long, defaultNumPartitions: Int, @@ -134,8 +134,8 @@ private[spark] class CometExecRDD( // fingerprints the bytes once so the probe here does not rehash them. val actualPlan = if (commonByKey.nonEmpty) { val basePlan = PlanDataInjector.parseBasePlan(serializedPlan, planFingerprint) - val injected = - PlanDataInjector.injectPlanData(basePlan, commonByKey, partition.planDataByKey) + val partitionBytes = partition.planDataByKey.map { case (key, data) => key -> data.bytes } + val injected = PlanDataInjector.injectPlanData(basePlan, commonByKey, partitionBytes) PlanDataInjector.serializeOperator(injected) } else { serializedPlan @@ -236,7 +236,7 @@ object CometExecRDD { sc: SparkContext, inputRDDs: Seq[RDD[_]], commonByKey: Map[String, Array[Byte]], - perPartitionByKey: Map[String, Array[Array[Byte]]], + perPartitionByKey: Map[String, Array[PartitionPlanData]], serializedPlan: Array[Byte], planFingerprint: Long, numPartitions: Int, 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 e3466c71d27..2c067e7c46b 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 @@ -92,7 +92,7 @@ case class CometIcebergNativeScanExec( * by every construction site), so values resolved through `waitForSubqueries` are visible on * both sides. */ - @transient private lazy val serializedPartitionData: (Array[Byte], Array[Array[Byte]]) = { + @transient private lazy val serializedPartitionData: (Array[Byte], Array[PartitionPlanData]) = { // Canonicalized instances set originalPlan = null and are not meant to be executed. // If we ever reach this lazy val on a canonicalized form, fail loud rather than NPE // deep inside originalPlan.inputRDD. @@ -122,7 +122,7 @@ case class CometIcebergNativeScanExec( def commonData: Array[Byte] = serializedPartitionData._1 - def perPartitionData: Array[Array[Byte]] = serializedPartitionData._2 + def perPartitionData: Array[PartitionPlanData] = serializedPartitionData._2 // numPartitions for execution - derived from actual DPP-filtered partitions // Only accessed during execution, not planning 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 8c8456b074a..1134ab2bcd2 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 @@ -264,7 +264,7 @@ case class CometNativeScanExec( sparkContext, Seq.empty, Map(sourceKey -> commonData), - Map(sourceKey -> perPartitionData), + Map(sourceKey -> PartitionPlanData.wrap(perPartitionData)), serializedPlan, PlanDataInjector.planFingerprint(serializedPlan), perPartitionData.length, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala b/spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala new file mode 100644 index 00000000000..9bce29c3d0e --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import scala.collection.mutable + +import org.apache.comet.serde.OperatorOuterClass + +/** + * One partition's planning data for one scan. It rides in the task with its partition object. + * + * A coalesce packs several partitions into one task, and Java serialization writes an object they + * share once per task stream. So a scan can hold an entry several partitions use as one shared + * object, and a coalesced task carries it once, as Spark's Iceberg reader does with its shared + * DeleteFile objects. + */ +trait PartitionPlanData extends Serializable { + + /** The serialized scan message this scan's [[PlanDataInjector]] reads. */ + def bytes: Array[Byte] +} + +object PartitionPlanData { + + /** Plan data that is already one serialized message, with nothing shared. */ + def apply(bytes: Array[Byte]): PartitionPlanData = new Bytes(bytes) + + def wrap(perPartition: Array[Array[Byte]]): Array[PartitionPlanData] = + perPartition.map(apply) + + /** Whether two scans' plan data match, for [[PlanDataInjector.findAllPlanData]]. */ + def same(a: Array[PartitionPlanData], b: Array[PartitionPlanData]): Boolean = + a.length == b.length && a.indices.forall { i => + (a(i) eq b(i)) || java.util.Arrays.equals(a(i).bytes, b(i).bytes) + } + + private final class Bytes(override val bytes: Array[Byte]) extends PartitionPlanData +} + +/** + * One native Iceberg scan partition's plan data. + * + * `scan` is the partition's IcebergScan: its file tasks, plus its partition values and + * delete-file sets in `common`. Its delete files and residuals are kept apart, because many + * partitions often use the same ones, such as an unpartitioned equality delete or the scan's + * filter. The driver serializes each once, and every partition that uses it holds the same array. + * A coalesced task then carries each once, and the driver keeps one copy. + * + * `deleteFiles(i)` is the partition's delete file `i`, an IcebergDeleteFile without + * `file_path_idx`, and `deleteFilePaths(i)` is its path. The delete-file sets in `scan` index + * into `deleteFiles`. Each partition orders its delete files differently, so the index into the + * path pool is set here, when the partition's message is built. `residuals(i)` is the + * IcebergPredicate its tasks' `residual_idx = i` refers to. + */ +final class IcebergPartitionPlanData( + val scan: Array[Byte], + val deleteFiles: Array[Array[Byte]], + val deleteFilePaths: Array[String], + val residuals: Array[Array[Byte]]) + extends PartitionPlanData { + + require( + deleteFiles.length == deleteFilePaths.length, + s"${deleteFiles.length} delete files but ${deleteFilePaths.length} paths") + + override def bytes: Array[Byte] = { + if (deleteFiles.isEmpty && residuals.isEmpty) { + scan + } else { + val builder = OperatorOuterClass.IcebergScan.newBuilder().mergeFrom(scan) + val common = builder.getCommonBuilder + val pathIndex = mutable.HashMap[String, Int]() + deleteFiles.indices.foreach { i => + val path = deleteFilePaths(i) + val pathIdx = pathIndex.getOrElseUpdate( + path, { + common.addDeleteFilePathPool(path) + pathIndex.size + }) + common.addDeleteFilePool( + OperatorOuterClass.IcebergDeleteFile + .newBuilder() + .mergeFrom(deleteFiles(i)) + .setFilePathIdx(pathIdx)) + } + residuals.foreach(r => common.addResidualPoolBuilder().mergeFrom(r)) + builder.build().toByteArray + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala index 8770496bff2..b34a4c56f87 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ import org.apache.spark.rdd.{DeterministicLevel, RDD} -import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} +import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode, PartitionPlanData} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometShuffleBlockIterator @@ -44,7 +44,7 @@ private[shuffle] class CometNativeShuffleInputRDD( numPartitionsParam: Int, shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, - @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, + @transient perPartitionByKey: Map[String, Array[PartitionPlanData]] = Map.empty, positionalRoundRobin: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, @@ -144,7 +144,7 @@ private[shuffle] class CometNativeShuffleInputRDD( private[shuffle] class CometNativeShuffleInputPartition( override val index: Int, val inputPartitions: Array[Partition], - val planDataByKey: Map[String, Array[Byte]]) + val planDataByKey: Map[String, PartitionPlanData]) extends Partition /** @@ -159,7 +159,7 @@ private[shuffle] class CometNativeShuffleInputIterator( val partitionIndex: Int, val inputObjects: Array[Object], val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator], - val planDataByKey: Map[String, Array[Byte]]) + val planDataByKey: Map[String, PartitionPlanData]) extends Iterator[Product2[Int, ColumnarBatch]] { override def hasNext: Boolean = false diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index f9e93533275..69ec05a4ed7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -156,7 +156,7 @@ class CometNativeShuffleWriter[K, V]( shuffleId, unifiedPlan, ctx.commonByKey, - shuffleInputIter.planDataByKey) + shuffleInputIter.planDataByKey.map { case (key, data) => key -> data.bytes }) } else { unifiedPlan } 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 d07490ee914..2549fadb159 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 @@ -454,7 +454,7 @@ private[comet] object PlanDataInjector extends Logging { * (commonByKey, perPartitionByKey) - common data is shared, per-partition varies */ private[comet] def findAllPlanData( - plan: SparkPlan): (Map[String, Array[Byte]], Map[String, Array[Array[Byte]]]) = { + plan: SparkPlan): (Map[String, Array[Byte]], Map[String, Array[PartitionPlanData]]) = { plan match { case iceberg: CometIcebergNativeScanExec => // Trigger Spark's standard prepare -> waitForSubqueries lifecycle so DPP @@ -496,7 +496,9 @@ private[comet] object PlanDataInjector extends Logging { s.ensureSubqueriesResolved() s.sendDriverMetrics() if (s.commonData.nonEmpty && s.perPartitionData.nonEmpty) { - (Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.perPartitionData)) + ( + Map(s.sourceKey -> s.commonData), + Map(s.sourceKey -> PartitionPlanData.wrap(s.perPartitionData))) } else { (Map.empty, Map.empty) } @@ -518,13 +520,10 @@ private[comet] object PlanDataInjector extends Logging { val results = plan.children.map(findAllPlanData) ( mergeByKey(results.map(_._1))(java.util.Arrays.equals(_, _)), - mergeByKey(results.map(_._2))(samePartitionData)) + mergeByKey(results.map(_._2))(PartitionPlanData.same)) } } - 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 @@ -573,11 +572,18 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector { override def inject(op: Operator, common: Prepared, partitionBytes: Array[Byte]): Operator = { val scan = op.getIcebergScan - val tasksOnly = OperatorOuterClass.IcebergScan.parseFrom(partitionBytes) + val partition = OperatorOuterClass.IcebergScan.parseFrom(partitionBytes) val scanBuilder = scan.toBuilder - scanBuilder.setCommon(common) - scanBuilder.addAllFileScanTasks(tasksOnly.getFileScanTasksList) + // The partition carries the per-file pool entries its tasks use, which the common block does + // not hold. mergeFrom appends repeated fields, so the tasks' indices, which start at 0, stay + // valid. + if (partition.hasCommon) { + scanBuilder.setCommon(common.toBuilder.mergeFrom(partition.getCommon)) + } else { + scanBuilder.setCommon(common) + } + scanBuilder.addAllFileScanTasks(partition.getFileScanTasksList) op.toBuilder.setIcebergScan(scanBuilder).build() } @@ -887,7 +893,7 @@ private[comet] case class NativeExecContext( // Partition object - see CometNativeShuffleInputRDD / CometExecRDD); the executor reads its own // slice, never this map. Keeping it off the wire stops it from bloating the broadcast task // binary when this context rides on the non-transient CometShuffleDependency.nativeShuffleSpec. - @transient perPartitionByKey: Map[String, Array[Array[Byte]]], + @transient perPartitionByKey: Map[String, Array[PartitionPlanData]], shuffleScanIndices: Set[Int], hasScanInput: Boolean, // Like plan data, file paths are sliced onto each task's Partition on the driver. diff --git a/spark/src/test/resources/sql-tests/iceberg/merge_on_read_partition_layout.sql b/spark/src/test/resources/sql-tests/iceberg/merge_on_read_partition_layout.sql new file mode 100644 index 00000000000..f986a07a466 --- /dev/null +++ b/spark/src/test/resources/sql-tests/iceberg/merge_on_read_partition_layout.sql @@ -0,0 +1,156 @@ +-- 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. + +-- The native scan ships each partition value, delete file and residual with every Spark +-- partition whose tasks use it (see CometIcebergNativeScan.serializePartitions). Grouping the +-- read by partition key makes every partition value and delete file belong to one Spark +-- partition. Without grouping, the tiny split size puts each data file +-- in its own Spark partition, so the two data files of a partition both get its value and its +-- delete file. The matrix runs both layouts against Spark, with and without a coalesce, which +-- puts several partitions' plan data in one task. Identity partition values become column +-- constants in the native scan, so a value resolved from the wrong pool entry shows up in the +-- results. CometIcebergNativeSuite checks what each partition ships. + +-- No Iceberg spark-runtime is published for Spark 4.2 yet, and the 4.0 runtime the build reuses +-- is binary-incompatible with it. See https://github.com/apache/datafusion-comet/issues/4969. +-- MaxSparkVersion: 4.1 + +-- Config: spark.sql.catalog.test_cat=org.apache.iceberg.spark.SparkCatalog +-- Config: spark.sql.catalog.test_cat.type=hadoop +-- Config: spark.sql.catalog.test_cat.warehouse=/tmp/comet-iceberg-sql-test +-- Config: spark.comet.scan.icebergNative.enabled=true +-- Config: spark.sql.sources.v2.bucketing.enabled=true +-- Config: spark.sql.session.timeZone=America/Los_Angeles +-- ConfigMatrix: spark.sql.iceberg.planning.preserve-data-grouping=true,false + +statement +DROP TABLE IF EXISTS test_cat.db.mor_partition_layout + +-- Hash distribution writes one data file per partition for each insert, and one delete file per +-- partition for the delete. +statement +CREATE TABLE test_cat.db.mor_partition_layout ( + id INT, + region STRING, + d DATE, + amount DECIMAL(18, 4), + ts TIMESTAMP, + ts_ntz TIMESTAMP_NTZ, + payload STRING +) USING iceberg +PARTITIONED BY (region, d, amount, ts, ts_ntz) +TBLPROPERTIES ( + 'format-version' = '2', + 'write.delete.mode' = 'merge-on-read', + 'write.delete.granularity' = 'partition', + 'write.distribution-mode' = 'hash', + 'write.delete.distribution-mode' = 'hash', + 'read.split.target-size' = '1', + 'read.split.open-file-cost' = '1') + +-- Five partitions: non-ASCII strings, dates and timestamps before 1970, the epoch, a leap day, +-- both ends of the decimal range, timestamps in a DST gap and overlap of the session timezone, +-- and a partition where every value is null. Each insert writes three rows to each partition. +statement +INSERT INTO test_cat.db.mor_partition_layout +SELECT CAST(t.k * 3 + r.id + 1 AS INT), t.region, t.d, t.amount, t.ts, t.ts_ntz, + concat('p', CAST(t.k * 3 + r.id + 1 AS STRING)) +FROM VALUES + (0, '東京', DATE '1969-12-31', CAST(-99999999999999.9999 AS DECIMAL(18, 4)), + TIMESTAMP '1969-07-20 20:17:40', TIMESTAMP_NTZ '1900-01-01 00:00:00'), + (1, 'Zürich', DATE '1900-01-01', CAST(0.0001 AS DECIMAL(18, 4)), + TIMESTAMP '2024-03-10 02:30:00', TIMESTAMP_NTZ '2024-03-10 02:30:00'), + (2, 'São Paulo', DATE '2024-02-29', CAST(99999999999999.9999 AS DECIMAL(18, 4)), + TIMESTAMP '2024-11-03 01:30:00', TIMESTAMP_NTZ '2024-11-03 01:30:00'), + (3, 'eng', DATE '1970-01-01', CAST(0 AS DECIMAL(18, 4)), + TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00'), + (4, CAST(NULL AS STRING), CAST(NULL AS DATE), CAST(NULL AS DECIMAL(18, 4)), + CAST(NULL AS TIMESTAMP), CAST(NULL AS TIMESTAMP_NTZ)) + AS t(k, region, d, amount, ts, ts_ntz) +CROSS JOIN range(3) AS r + +statement +INSERT INTO test_cat.db.mor_partition_layout +SELECT CAST(15 + t.k * 3 + r.id + 1 AS INT), t.region, t.d, t.amount, t.ts, t.ts_ntz, + concat('p', CAST(15 + t.k * 3 + r.id + 1 AS STRING)) +FROM VALUES + (0, '東京', DATE '1969-12-31', CAST(-99999999999999.9999 AS DECIMAL(18, 4)), + TIMESTAMP '1969-07-20 20:17:40', TIMESTAMP_NTZ '1900-01-01 00:00:00'), + (1, 'Zürich', DATE '1900-01-01', CAST(0.0001 AS DECIMAL(18, 4)), + TIMESTAMP '2024-03-10 02:30:00', TIMESTAMP_NTZ '2024-03-10 02:30:00'), + (2, 'São Paulo', DATE '2024-02-29', CAST(99999999999999.9999 AS DECIMAL(18, 4)), + TIMESTAMP '2024-11-03 01:30:00', TIMESTAMP_NTZ '2024-11-03 01:30:00'), + (3, 'eng', DATE '1970-01-01', CAST(0 AS DECIMAL(18, 4)), + TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00'), + (4, CAST(NULL AS STRING), CAST(NULL AS DATE), CAST(NULL AS DECIMAL(18, 4)), + CAST(NULL AS TIMESTAMP), CAST(NULL AS TIMESTAMP_NTZ)) + AS t(k, region, d, amount, ts, ts_ntz) +CROSS JOIN range(3) AS r + +-- Removes one row from each data file, so each partition's delete file covers both of its data +-- files. +statement +DELETE FROM test_cat.db.mor_partition_layout WHERE id % 3 = 2 + +query +SELECT * FROM test_cat.db.mor_partition_layout ORDER BY id + +query +SELECT id, _partition FROM test_cat.db.mor_partition_layout ORDER BY id + +query +SELECT region, d, amount, ts, ts_ntz, count(*) AS n, sum(id) AS s +FROM test_cat.db.mor_partition_layout +GROUP BY region, d, amount, ts, ts_ntz +ORDER BY s + +query +SELECT count(*), sum(id) FROM test_cat.db.mor_partition_layout + +-- Every task carries the same residual on a data column. +query +SELECT id, region, amount FROM test_cat.db.mor_partition_layout +WHERE payload IN ('p1', 'p17', 'p29') +ORDER BY id + +-- The residual depends on the partition: true where region matches, id > 25 elsewhere. +query +SELECT id, region, d FROM test_cat.db.mor_partition_layout +WHERE region = '東京' OR id > 25 +ORDER BY id + +query +SELECT id, d, ts FROM test_cat.db.mor_partition_layout +WHERE d < DATE '1970-01-01' +ORDER BY id + +query +SELECT id, payload FROM test_cat.db.mor_partition_layout +WHERE region IS NULL +ORDER BY id + +-- A coalesce without a shuffle runs several partitions in one task. +query +SELECT /*+ COALESCE(1) */ * FROM test_cat.db.mor_partition_layout ORDER BY id + +query +SELECT /*+ COALESCE(2) */ id, _partition FROM test_cat.db.mor_partition_layout +WHERE region = '東京' OR id > 25 +ORDER BY id + +statement +DROP TABLE test_cat.db.mor_partition_layout diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala index ad9bd99b599..ef1f0d40253 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.functions.{col, lit} import org.apache.spark.sql.types._ import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.IcebergScanPlanData.distinctResiduals class CometFuzzIcebergSuite extends CometFuzzIcebergBase { @@ -177,7 +177,7 @@ class CometFuzzIcebergSuite extends CometFuzzIcebergBase { } // Asserts that predicate pushdown actually happens for the types that support it. For each - // primitive column it runs an equality filter and inspects the serialized IcebergScanCommon: the + // primitive column it runs an equality filter and inspects the scan's serialized pools: the // residual pool must be non-empty exactly when the column's type is one Comet maps to an Iceberg // literal (pushable), and empty for a gated type (binary). This works because the fuzz table is // unpartitioned, so Iceberg's per-file residual is the full filter rather than a partition-pruned @@ -226,8 +226,7 @@ class CometFuzzIcebergSuite extends CometFuzzIcebergBase { val (_, cometPlan) = checkSparkAnswer(df.where(col(name) === lit(v))) val scans = collectIcebergNativeScans(cometPlan) assert(scans.length == 1, s"expected one native scan for '$name'\n$cometPlan") - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scans.head.commonData) - val pushed = common.getResidualPoolCount > 0 + val pushed = distinctResiduals(scans.head).nonEmpty if (isPushable(field.dataType)) { assert( pushed, @@ -253,8 +252,8 @@ class CometFuzzIcebergSuite extends CometFuzzIcebergBase { assert(scans.length == 1, s"$cometPlan") // Older Iceberg leaks manifest streams when commonData forces planning. if (!isIcebergVersionLessThan("1.8.0")) { - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scans.head.commonData) - assert(common.getResidualPoolCount == 0, s"unexpected residual for $name: $predicate") + val residuals = distinctResiduals(scans.head) + assert(residuals.isEmpty, s"unexpected residual for $name: $predicate") } } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 3f2bd95b6be..85fd3ad3f43 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -34,7 +34,7 @@ import org.apache.iceberg.expressions.Expressions import org.apache.iceberg.spark.Spark3Util import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile -import org.apache.spark.CometListenerBusUtils +import org.apache.spark.{CometListenerBusUtils, SparkEnv} import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{Ascending, DynamicPruningExpression, SortOrder} @@ -49,6 +49,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, StringType, StructType, TimestampType} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus} +import org.apache.comet.IcebergScanPlanData.{fileScanTaskCount, partitionPools, partitionScans, storedPools} import org.apache.comet.cloud.s3.MinioLocationScopedCredentialProvider import org.apache.comet.iceberg.{IcebergReflection, RESTCatalogHelper} import org.apache.comet.rules.CometScanRule @@ -783,10 +784,10 @@ class CometIcebergNativeSuite checkSparkAnswerAndOperator("SELECT * FROM test_cat.db.dv_shared_puffin ORDER BY id") assertSingleNativeScan(cometPlan) - val commonBytes = collectIcebergNativeScans(cometPlan).head.commonData - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(commonBytes) + val scan = collectIcebergNativeScans(cometPlan).head + val stored = storedPools(scan) - val vectors = common.getDeleteFilePoolList.asScala.toSeq + val vectors = stored.flatMap(_.getDeleteFilePoolList.asScala) assert( vectors.length == dataFiles, s"expected $dataFiles pooled deletion vectors, one per data file, but got " + @@ -810,30 +811,126 @@ class CometIcebergNativeSuite vectors.forall(v => v.hasRecordCount && v.getRecordCount > 0), s"expected every vector to carry a record count, got ${vectors.map(_.getRecordCount)}") - // All of them point at the one interned path, which is serialized exactly once however - // many vectors reference it. - val pathPool = common.getDeleteFilePathPoolList.asScala.toSeq + // All of them point at the one interned path, which each partition serializes once + // however many of its vectors reference it. + val pathPool = stored.flatMap(_.getDeleteFilePathPoolList.asScala).distinct + val resolvedPaths = partitionPools(scan).flatMap { pools => + pools.getDeleteFilePoolList.asScala.map(v => + pools.getDeleteFilePathPool(v.getFilePathIdx)) + } assert( - pathPool.length == 1 && vectors.map(_.getFilePathIdx).distinct == Seq(0), + pathPool.length == 1 && resolvedPaths.distinct == pathPool, s"expected one interned Puffin path referenced by every vector, got pool $pathPool " + - s"and indices ${vectors.map(_.getFilePathIdx)}") + s"and resolved paths $resolvedPaths") + val occurrences = + scan.perPartitionData.toSeq + .map(data => countByteOccurrences(data.bytes, pathPool.head.getBytes(UTF_8))) + // Each partition that uses the path ships it once; the common block never holds it. + val inCommon = countByteOccurrences(scan.commonData, pathPool.head.getBytes(UTF_8)) assert( - countByteOccurrences(commonBytes, pathPool.head.getBytes(UTF_8)) == 1, - s"Puffin path ${pathPool.head} serialized more than once in IcebergScanCommon " + - s"(${commonBytes.length} bytes)") + inCommon == 0 && occurrences.forall(_ <= 1), + s"Puffin path ${pathPool.head} serialized more than once: common $inCommon, " + + s"partitions $occurrences") spark.sql("DROP TABLE test_cat.db.dv_shared_puffin") } } } + // The same shared-Puffin layout, but with every data file in its own Spark partition. Each + // partition ships its own vector and the Puffin path with it, and nothing goes in the common + // block. The rows must still match Spark. + test("V3 deletion vectors in different partitions keep their shared Puffin path") { + assume(icebergAvailable, "Iceberg not available in classpath") + assume(icebergVersionAtLeast(1, 11), "Iceberg V3 tables require Iceberg 1.11+") + + withHadoopCatalog("test_cat") { + withSQLConf( + // One shuffle partition routes every vector into one Puffin file (see the test above). + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + // The tiny split size puts every data file in its own task. + spark.sql(""" + CREATE TABLE test_cat.db.dv_split_puffin (id INT, name STRING) + USING iceberg + TBLPROPERTIES ( + 'format-version' = '3', + 'write.delete.mode' = 'merge-on-read', + 'write.merge.mode' = 'merge-on-read', + 'read.split.target-size' = '1', + 'read.split.open-file-cost' = '1' + ) + """) + val dataFiles = 3 + for (f <- 0 until dataFiles) { + val values = (1 to 10).map(i => s"(${f * 100 + i}, 'n${f * 100 + i}')").mkString(", ") + spark.sql(s"INSERT INTO test_cat.db.dv_split_puffin VALUES $values") + } + val deletedIds = (0 until dataFiles).map(f => f * 100 + 1) + spark.sql( + s"DELETE FROM test_cat.db.dv_split_puffin WHERE id IN (${deletedIds.mkString(", ")})") + + val puffinPaths = spark + .sql("SELECT file_path FROM test_cat.db.dv_split_puffin.files WHERE content = 1") + .collect() + .map(_.getString(0)) + .toSeq + assert( + puffinPaths.length == dataFiles && puffinPaths.distinct.length == 1, + s"expected $dataFiles vectors in one Puffin file, got: $puffinPaths") + + val (_, cometPlan) = + checkSparkAnswerAndOperator("SELECT * FROM test_cat.db.dv_split_puffin ORDER BY id") + val scan = collectIcebergNativeScans(cometPlan).head + // Each INSERT can write several data files; the split size gives each its own partition. + val dataFileCount = spark + .sql("SELECT count(*) FROM test_cat.db.dv_split_puffin.files WHERE content = 0") + .collect() + .head + .getLong(0) + assert( + scan.perPartitionData.length == dataFileCount, + "expected one partition per data file") + val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) + assert( + common.getDeleteFilePathPoolCount == 0 && common.getDeleteFilePoolCount == 0, + s"expected no delete files in the common block: $common") + val own = partitionScans(scan).map(_.getCommon) + // Each vector belongs to one data file, so exactly that file's partition ships it, + // together with one copy of the Puffin path. + val vectorCounts = own.map(_.getDeleteFilePoolCount) + assert( + vectorCounts.sum == dataFiles && vectorCounts.forall(_ <= 1), + s"vectors per partition: $vectorCounts") + assert( + own + .filter(_.getDeleteFilePoolCount == 1) + .forall(_.getDeleteFilePathPoolList.asScala.toSeq == puffinPaths.distinct), + "expected each partition with a vector to ship the Puffin path once") + // The partitions hold one path string between them, so a coalesced task carries it once. + val heldPaths = scan.perPartitionData.toSeq.flatMap { + case data: IcebergPartitionPlanData => data.deleteFilePaths.toSeq + case other => fail(s"unexpected Iceberg plan data: $other") + } + assert( + heldPaths.length == dataFiles && heldPaths.forall(_ eq heldPaths.head), + "partitions should share one copy of the Puffin path") + + spark.sql("DROP TABLE test_cat.db.dv_split_puffin") + } + } + } + // Under Iceberg's default partition delete granularity, one position-delete file applies to every // data file in the partition with a compatible sequence number (DeleteFileIndex.forDataFile). // Interleaving inserts and deletes staggers data-file sequence numbers, so each FileScanTask sees // a different subset of the same delete files. A shared delete file must be serialized once in - // the broadcast IcebergScanCommon, not once per referencing set: duplicating it is quadratic in - // the number of delete commits and can overflow protobuf's 2 GiB message limit - // (getSerializedSize() returns int, so it wraps to a negative array length). + // each partition's plan data, not once per referencing set: duplicating it is quadratic in the + // number of delete commits and can overflow protobuf's 2 GiB message limit (getSerializedSize() + // returns int, so it wraps to a negative array length). test("delete file pool does not duplicate shared delete files (serde size)") { assume(icebergAvailable, "Iceberg not available in classpath") // Reading scan.commonData forces Iceberg planning (ParallelIterable), which on older Iceberg @@ -883,12 +980,10 @@ class CometIcebergNativeSuite val scans = collectIcebergNativeScans(cometPlan) assert(scans.length == 1, s"expected one native scan, got ${scans.length}\n$cometPlan") - val commonBytes = scans.head.commonData - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(commonBytes) - - val distinctPaths = common.getDeleteFilePathPoolList.asScala.toSeq + val stored = storedPools(scans.head) + val distinctPaths = stored.flatMap(_.getDeleteFilePathPoolList.asScala).distinct val totalReferences = - common.getDeleteFilesPoolList.asScala.map(_.getDeleteFileIndicesCount).sum + stored.flatMap(_.getDeleteFilesPoolList.asScala).map(_.getDeleteFileIndicesCount).sum // Guard against a vacuous pass: we must actually exercise the shared-delete-file case // (more references than distinct files, i.e. at least one file shared across tasks). @@ -898,22 +993,126 @@ class CometIcebergNativeSuite s"$totalReferences references") // Count raw-byte occurrences of each path so the check is agnostic to how the pool is - // structured (set-of-copies before the fix, flat index pool after). - val duplicated = distinctPaths - .map(p => p -> countByteOccurrences(commonBytes, p.getBytes(UTF_8))) - .filter(_._2 > 1) + // structured (set-of-copies before the fix, flat index pool after). Each partition ships + // the delete files its tasks use, so a path may appear in several partitions, but only + // once in each. + val planData = scans.head.commonData +: scans.head.perPartitionData.toSeq.map(_.bytes) + val duplicated = for { + p <- distinctPaths + (bytes, i) <- planData.zipWithIndex + count = countByteOccurrences(bytes, p.getBytes(UTF_8)) + if count > 1 + } yield s"$p x$count in block $i" assert( duplicated.isEmpty, - "delete files serialized more than once in IcebergScanCommon " + - s"(${commonBytes.length} bytes): " + - duplicated.map { case (p, c) => s"$p x$c" }.mkString(", ")) + "delete files serialized more than once in one partition's plan data: " + + duplicated.mkString(", ")) spark.sql("DROP TABLE test_cat.db.delete_pool_test") } } } + // Every task receives the common block, so it must not grow with the table. Each delete file + // ships with every partition that uses it: one partition when the read is grouped by bucket, + // each of its data files' partitions when it is not. sql-tests/iceberg/merge_on_read_partition_layout.sql checks + // the rows both layouts read against Spark. + test("delete files ship with each partition that uses them") { + assume(icebergAvailable, "Iceberg not available in classpath") + // Reading scan.commonData forces Iceberg planning (ParallelIterable), which on older Iceberg + // leaves a prefetched manifest stream open that Spark's DebugFilesystem flags at teardown. + assume( + icebergVersionAtLeast(1, 8), + "ParallelIterable leaks manifest streams on older Iceberg") + + withHadoopCatalog("test_cat") { + withSQLConf( + 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 when the read is not grouped. + // Hash distribution writes one data file per bucket for each insert, and one delete file + // per bucket for the delete, whatever Iceberg's defaults are. + spark.sql(""" + CREATE TABLE test_cat.db.bucketed_mor (id INT, v STRING) USING iceberg + PARTITIONED BY (bucket(4, id)) + TBLPROPERTIES ( + 'format-version' = '2', + 'write.delete.mode' = 'merge-on-read', + 'write.delete.granularity' = 'partition', + 'write.distribution-mode' = 'hash', + 'write.delete.distribution-mode' = 'hash', + 'read.split.target-size' = '1', + 'read.split.open-file-cost' = '1') + """) + // Two inserts give each bucket two data files. The delete then writes one position delete + // file per bucket, covering rows in both. + spark.sql( + "INSERT INTO test_cat.db.bucketed_mor " + + "SELECT CAST(id AS INT), concat('a', id) FROM range(0, 40, 1, 1)") + spark.sql( + "INSERT INTO test_cat.db.bucketed_mor " + + "SELECT CAST(id AS INT), concat('b', id) FROM range(40, 80, 1, 1)") + spark.sql("DELETE FROM test_cat.db.bucketed_mor WHERE id % 5 = 0") + + val deleteFiles = spark + .sql("SELECT count(*) FROM test_cat.db.bucketed_mor.files WHERE content = 1") + .collect() + .head + .getLong(0) + assert( + deleteFiles >= 2, + s"expected position delete files in several buckets: $deleteFiles") + + def scanOf(grouped: Boolean): CometIcebergNativeScanExec = { + var scan: CometIcebergNativeScanExec = null + withSQLConf( + SQLConf.V2_BUCKETING_ENABLED.key -> "true", + "spark.sql.iceberg.planning.preserve-data-grouping" -> grouped.toString) { + val df = spark.sql("SELECT * FROM test_cat.db.bucketed_mor") + df.collect() + val plan = df.queryExecution.executedPlan + val scans = collectIcebergNativeScans(plan) + assert(scans.length == 1, s"expected one native scan\n$plan") + scan = scans.head + } + scan + } + + // The common block never holds per-file entries, whatever the layout. + def assertNoFileEntries(scan: CometIcebergNativeScanExec): Unit = { + val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) + assert( + common.getDeleteFilePoolCount == 0 && common.getDeleteFilePathPoolCount == 0 && + common.getDeleteFilesPoolCount == 0 && common.getPartitionDataPoolCount == 0, + s"per-file entries in the common block: $common") + } + def ownDeletes(scan: CometIcebergNativeScanExec): Seq[Int] = + partitionScans(scan).map(_.getCommon.getDeleteFilePoolCount) + + // One bucket per partition: each delete file goes to exactly one partition. + val grouped = scanOf(grouped = true) + assertNoFileEntries(grouped) + assert(ownDeletes(grouped).sum == deleteFiles, s"per partition: ${ownDeletes(grouped)}") + + // One data file per partition: a bucket's delete file goes to each of its data files' + // partitions, and no partition gets more than its own bucket's delete file. + val perFile = scanOf(grouped = false) + assertNoFileEntries(perFile) + assert( + perFile.perPartitionData.length > grouped.perPartitionData.length, + "expected the ungrouped read to put data files of one bucket in different partitions") + assert( + ownDeletes(perFile).sum > deleteFiles && ownDeletes(perFile).forall(_ <= 1), + s"expected shared delete files to ship with each partition: ${ownDeletes(perFile)}") + + spark.sql("DROP TABLE test_cat.db.bucketed_mor") + } + } + } + test("bytes_scanned includes delete file I/O") { assume(icebergAvailable, "Iceberg not available in classpath") @@ -1169,6 +1368,110 @@ class CometIcebergNativeSuite table.newRowDelta().addDeletes(deleteFile).commit() } + // An unpartitioned equality delete applies to every older data file, so each partition holding + // one ships it, as Spark's Iceberg reader ships each task's delete files. The rows must still + // match Spark, including a newer row the delete does not apply to. + test("a global equality delete ships with each partition that uses it") { + assume(icebergAvailable, "Iceberg not available in classpath") + assume( + icebergVersionAtLeast(1, 8), + "ParallelIterable leaks manifest streams on older Iceberg") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_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 puts every data file in its own task. + spark.sql(""" + CREATE TABLE test_cat.db.global_eq_delete (id INT, name STRING) USING iceberg + TBLPROPERTIES ( + 'format-version' = '2', + 'read.split.target-size' = '1', + 'read.split.open-file-cost' = '1') + """) + for (f <- 0 until 4) { + val values = (1 to 10).map(i => s"(${f * 100 + i}, 'n${f * 100 + i}')").mkString(", ") + spark.sql(s"INSERT INTO test_cat.db.global_eq_delete VALUES $values") + } + commitEqualityDeletes( + "test_cat", + "db", + "global_eq_delete", + "id", + Seq(2, 105, 210, 315), + warehouseDir) + // A newer append the delete does not apply to: its id 2 row must survive, and the + // partitions holding only this file must not get the delete. + spark.sql("INSERT INTO test_cat.db.global_eq_delete VALUES (2, 'newer')") + + val (_, cometPlan) = + checkSparkAnswerAndOperator("SELECT * FROM test_cat.db.global_eq_delete ORDER BY id") + val scan = collectIcebergNativeScans(cometPlan).head + val own = partitionScans(scan).map(_.getCommon) + // Each INSERT can write several data files; the split size gives each its own partition. + val dataFileCount = spark + .sql("SELECT count(*) FROM test_cat.db.global_eq_delete.files WHERE content = 0") + .collect() + .head + .getLong(0) + assert( + own.length == dataFileCount, + s"expected one partition per data file: ${own.length}") + // Older files' partitions ship the delete once each; the newer file's partition does not. + val copies = + own.map(_.getDeleteFilePoolList.asScala.count(_.getContentType == "EQUALITY_DELETES")) + assert( + copies.forall(_ <= 1) && copies.contains(1) && copies.contains(0), + s"equality delete copies per partition: $copies") + val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) + assert( + common.getDeleteFilePoolCount == 0, + s"the common block holds delete files: $common") + // The partitions hold one serialized copy between them, so the driver keeps one and a + // coalesced task carries one. + val heldCopies = scan.perPartitionData.toSeq.flatMap { + case data: IcebergPartitionPlanData => data.deleteFiles.toSeq + case other => fail(s"unexpected Iceberg plan data: $other") + } + assert( + heldCopies.length == copies.sum && heldCopies.forall(_ eq heldCopies.head), + "partitions should share one serialized copy of the equality delete") + val coalesced = spark.table("test_cat.db.global_eq_delete").coalesce(1) + checkSparkAnswer(coalesced) + + // A real coalesce runs every partition in one task. That task's serialized partition + // carries the equality delete and its path once, not once per input partition. + val coalescedPlan = coalesced.queryExecution.executedPlan + assert(collectIcebergNativeScans(coalescedPlan).nonEmpty, s"$coalescedPlan") + val coalescedRdd = coalescedPlan.execute() + assert(coalescedRdd.getNumPartitions == 1, s"$coalescedPlan") + val taskPartition = SparkEnv.get.closureSerializer + .newInstance() + .serialize(coalescedRdd.partitions.head) + val taskBytes = new Array[Byte](taskPartition.remaining()) + taskPartition.get(taskBytes) + val eqDeletePath = scan.perPartitionData.toSeq.collectFirst { + case data: IcebergPartitionPlanData if data.deleteFilePaths.nonEmpty => + data.deleteFilePaths.head + }.get + assert( + copies.count(_ == 1) > 1, + s"expected several partitions to use the delete: $copies") + assert( + countByteOccurrences(taskBytes, eqDeletePath.getBytes(UTF_8)) == 1 && + countByteOccurrences(taskBytes, heldCopies.head) == 1, + s"the coalesced task (${taskBytes.length} bytes) should carry one copy of the delete") + + spark.sql("DROP TABLE test_cat.db.global_eq_delete") + } + } + } + private def latestSnapshotId(table: String): Long = spark .sql(s"SELECT snapshot_id FROM $table.snapshots ORDER BY committed_at DESC LIMIT 1") .collect()(0) @@ -3902,10 +4205,8 @@ class CometIcebergNativeSuite assert(scans.length == 1, s"$cometPlan") // Planning commonData leaks manifest streams on Iceberg versions before 1.8.0. if (icebergVersionAtLeast(1, 8)) { - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scans.head.commonData) - assert( - common.getResidualPoolCount == 0, - s"unexpected complex-column residual: $query") + val residuals = storedPools(scans.head).map(_.getResidualPoolCount).sum + assert(residuals == 0, s"unexpected complex-column residual: $query") } } } @@ -5046,9 +5347,7 @@ class CometIcebergNativeSuite // count the planned file tasks instead: 1 of the 3 date partitions' files is left. The // num_splits metric is not used here because the ORDER BY's range partitioning runs the // scan once to sample bounds and again for the shuffle, so each split is read twice. - val plannedTasks = icebergScans.head.perPartitionData - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) - .sum + val plannedTasks = fileScanTaskCount(icebergScans.head) assert( plannedTasks == 1, s"Expected DPP to prune to 1 of 3 files but planned $plannedTasks tasks:\n$cometPlan") @@ -5176,27 +5475,30 @@ class CometIcebergNativeSuite val scans = collectIcebergNativeScans(cometPlan) assert(scans.length == 1, s"expected one native scan, got ${scans.length}\n$cometPlan") - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scans.head.commonData) - val residuals = common.getResidualPoolList.asScala.toSeq + // Each partition ships the residual its tasks use once, however many tasks carry it. + val residualsPerPartition = storedPools(scans.head).map(_.getResidualPoolCount) // Count how many FileScanTasks reference a residual across all partitions. Every task here // carries the same IN predicate, so a correct pool holds one entry shared by all of them. - val tasksWithResidual = scans.head.perPartitionData.map { bytes => - OperatorOuterClass.IcebergScan - .parseFrom(bytes) - .getFileScanTasksList - .asScala - .count(_.hasResidualIdx) - }.sum + val tasksWithResidualPerPartition = + partitionScans(scans.head).map { partition => + partition.getFileScanTasksList.asScala + .count(_.hasResidualIdx) + } + val tasksWithResidual = tasksWithResidualPerPartition.sum - // Guard against a vacuous pass: the dedup is only meaningful if many tasks share the entry. + // Guard against a vacuous pass: the dedup is only meaningful if several tasks in one + // partition share the entry. assert( - tasksWithResidual > 1, - s"expected multiple tasks to carry a residual, got $tasksWithResidual") + tasksWithResidualPerPartition.exists(_ > 1), + "expected a partition with several tasks carrying a residual: " + + tasksWithResidualPerPartition) + // The common block holds none, and each partition with tasks ships it once. + val expectedResiduals = 0 +: tasksWithResidualPerPartition.map(n => if (n > 0) 1 else 0) assert( - residuals.size == 1, - "expected the shared non-partition residual to dedup to 1 pool entry across " + - s"$tasksWithResidual task references, got ${residuals.size}") + residualsPerPartition == expectedResiduals, + "expected the shared non-partition residual to dedup to 1 pool entry per partition " + + s"across $tasksWithResidual task references, got $residualsPerPartition") spark.sql("DROP TABLE test_cat.db.residual_dedup_test") } @@ -5848,9 +6150,7 @@ class CometIcebergNativeSuite val icebergScans = collectIcebergNativeScans(cometPlan) assert(icebergScans.nonEmpty, "Expected CometIcebergNativeScanExec in plan") val scan = icebergScans.head - val plannedTasks = scan.perPartitionData - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) - .sum + val plannedTasks = fileScanTaskCount(scan) val numSplits = scan.metrics("num_splits").value assert( plannedTasks == 1 && numSplits == 1, @@ -6117,9 +6417,7 @@ class CometIcebergNativeSuite val unprunedScan = collectIcebergNativeScans(unprunedPlan).head assert(unprunedScan.numPartitions == 1, s"Expected one split:\n$unprunedPlan") assert( - unprunedScan.perPartitionData - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) - .sum == 1, + fileScanTaskCount(unprunedScan) == 1, s"Expected one file scan task before pruning:\n$unprunedPlan") spark @@ -6154,8 +6452,7 @@ class CometIcebergNativeSuite collectIcebergDPPSubqueries(plan).nonEmpty, s"Expected executable DPP on the native scan:\n$plan") assert( - scan.perPartitionData - .forall(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount == 0), + partitionScans(scan).forall(_.getFileScanTasksCount == 0), s"Expected DPP to prune all file scan tasks:\n$plan") assert(scan.metrics("num_splits").value == 0, "Pruned files must not be read") } @@ -6192,9 +6489,7 @@ class CometIcebergNativeSuite val factB = "aqe_cat.db.single_split_join_b" val tables = Seq(factA, factB) def fileTaskCount(scan: CometIcebergNativeScanExec): Int = - scan.perPartitionData - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) - .sum + fileScanTaskCount(scan) try { tables.foreach { table => @@ -6658,9 +6953,7 @@ class CometIcebergNativeSuite // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 data files. // Iceberg filters its already planned tasks at runtime, so count the serialized tasks // and the splits read rather than Iceberg's planning metrics. - val plannedTasks = scan.perPartitionData - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) - .sum + val plannedTasks = fileScanTaskCount(scan) val numSplits = scan.metrics("num_splits").value assert( plannedTasks == 1 && numSplits == 1, diff --git a/spark/src/test/scala/org/apache/comet/IcebergScanPlanData.scala b/spark/src/test/scala/org/apache/comet/IcebergScanPlanData.scala new file mode 100644 index 00000000000..83195d99f9a --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/IcebergScanPlanData.scala @@ -0,0 +1,55 @@ +/* + * 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 + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.comet.CometIcebergNativeScanExec + +import org.apache.comet.serde.OperatorOuterClass.{IcebergPredicate, IcebergScan, IcebergScanCommon} + +/** + * Reads the plan data a native Iceberg scan ships, for tests. The common block holds no per-file + * pool entries; each partition ships the ones its own tasks use. + */ +object IcebergScanPlanData { + + /** Each partition's message, as its task's injector reads it. */ + def partitionScans(scan: CometIcebergNativeScanExec): Seq[IcebergScan] = + scan.perPartitionData.toSeq.map(data => IcebergScan.parseFrom(data.bytes)) + + /** The pool entries a scan ships: the common block's, then each partition's own. */ + def storedPools(scan: CometIcebergNativeScanExec): Seq[IcebergScanCommon] = + IcebergScanCommon.parseFrom(scan.commonData) +: partitionScans(scan).map(_.getCommon) + + /** The pools each partition's tasks index into once the injector merges them. */ + def partitionPools(scan: CometIcebergNativeScanExec): Seq[IcebergScanCommon] = { + val common = IcebergScanCommon.parseFrom(scan.commonData) + partitionScans(scan).map(partition => common.toBuilder.mergeFrom(partition.getCommon).build()) + } + + /** The residuals a scan ships, each counted once however many partitions use it. */ + def distinctResiduals(scan: CometIcebergNativeScanExec): Set[IcebergPredicate] = + storedPools(scan).flatMap(_.getResidualPoolList.asScala).toSet + + /** The file scan tasks planned for the scan, over all partitions. */ + def fileScanTaskCount(scan: CometIcebergNativeScanExec): Int = + partitionScans(scan).map(_.getFileScanTasksCount).sum +} diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala index afb3de2a55c..978b1fb89df 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala @@ -449,13 +449,14 @@ class CometIcebergNativeScanSuite extends AnyFunSuite with Matchers { private def keyMetadataMethod(clazz: Class[_]) = clazz.getMethod("keyMetadata") - private def serializeDeleteFile(file: AnyRef) = + private def serializeDeleteFileWithPath(file: AnyRef) = CometIcebergNativeScan.serializeDeleteFile( file, file.getClass, file.getClass, - keyMetadataMethod(file.getClass), - _ => 0) + keyMetadataMethod(file.getClass)) + + private def serializeDeleteFile(file: AnyRef) = serializeDeleteFileWithPath(file)._2 test("required Iceberg delete-file accessors are present") { Seq("content", "specId", "equalityFieldIds").foreach { accessor => @@ -469,11 +470,12 @@ class CometIcebergNativeScanSuite extends AnyFunSuite with Matchers { } test("position-delete file with null equality ids serializes without equality ids") { - val proto = serializeDeleteFile(new PositionDeleteFile) + val (path, proto) = serializeDeleteFileWithPath(new PositionDeleteFile) + // The path ships apart from the message, which sets its pool index per partition. + path shouldBe "s3://bucket/pos-delete.parquet" proto.getContentType shouldBe "POSITION_DELETES" proto.getPartitionSpecId shouldBe 7 proto.getEqualityIdsCount shouldBe 0 - proto.getFilePathIdx shouldBe 0 } test("equality-delete file serializes declared equality ids") { 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 9bf49125d7d..c3633121621 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 @@ -120,7 +120,8 @@ class CometScanWithPlanDataSuite extends AnyFunSuite { assert(commonByKey.keySet == Set("contrib-key")) assert(commonByKey("contrib-key") sameElements common) assert(perPartitionByKey.keySet == Set("contrib-key")) - assert(perPartitionByKey("contrib-key").map(_.toSeq).toSeq == perPart.map(_.toSeq).toSeq) + assert( + perPartitionByKey("contrib-key").map(_.bytes.toSeq).toSeq == perPart.map(_.toSeq).toSeq) assert( scan.subqueriesResolved, "findAllPlanData must drive the scan's subquery-resolution lifecycle before reading data") @@ -145,7 +146,9 @@ class CometScanWithPlanDataSuite extends AnyFunSuite { 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))) + perPartitionByKey("shared-key").map(_.bytes.toSeq).toSeq == Seq( + Seq[Byte](10), + Seq[Byte](20))) } test("findAllPlanData rejects two scans with different data under one key") { 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 bb165f58346..19ac47f19f4 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 @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet +import scala.jdk.CollectionConverters._ + import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.SparkConf @@ -126,6 +128,247 @@ class PlanDataInjectorSuite extends AnyFunSuite { assert(result.getChildren(0).getPlanId == 2) } + /** + * One partition of a bucketed merge-on-read scan as the driver ships it: two tasks, and in its + * IcebergScan.common only the delete files, paths and partition value they use. + */ + private def bucketPartition(bucket: Int): OperatorOuterClass.IcebergScan = { + val pools = OperatorOuterClass.IcebergScanCommon.newBuilder() + (0 until 2).foreach { d => + pools.addDeleteFilePathPool(f"s3://bucket/table/data/b$bucket%08d-$d-deletes.parquet") + pools.addDeleteFilePool( + OperatorOuterClass.IcebergDeleteFile + .newBuilder() + .setFilePathIdx(d) + .setContentType("POSITION_DELETES") + .setFileFormat("PARQUET") + .setRecordCount(10)) + } + pools.addDeleteFilesPool( + OperatorOuterClass.DeleteFileList + .newBuilder() + .addDeleteFileIndices(0) + .addDeleteFileIndices(1)) + pools.addPartitionDataPool( + OperatorOuterClass.PartitionData + .newBuilder() + .addValues( + OperatorOuterClass.PartitionValue + .newBuilder() + .setFieldId(1000) + .setLiteral(OperatorOuterClass.IcebergLiteral.newBuilder().setIntVal(bucket)))) + val scan = OperatorOuterClass.IcebergScan.newBuilder().setCommon(pools) + (0 until 2).foreach { f => + scan.addFileScanTasks( + OperatorOuterClass.IcebergFileScanTask + .newBuilder() + .setDataFilePath(f"s3://bucket/table/data/b$bucket%08d-$f.parquet") + .setDeleteFilesIdx(0) + .setPartitionDataIdx(0)) + } + scan.build() + } + + test("an injected Iceberg plan holds the common block plus its own partition's pools") { + // The common block holds no per-file pools, so a task's plan carries only the delete files + // and partition values its own files use, however many other partitions the scan has. + val scanOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 7) + val key = IcebergPlanDataInjector.getKey(scanOp).get + val (commonBytes, _) = + icebergPlanData("s3://table/metadata/v1.json", 7, Seq("id", "v"), "unused.parquet") + val partition = bucketPartition(42) + val injected = PlanDataInjector + .injectPlanData( + parseBasePlan(scanOp.toByteArray), + Map(key -> commonBytes), + Map(key -> partition.toByteArray)) + .getIcebergScan + + val pools = injected.getCommon + assert(pools.getRequiredSchemaCount == 2, "the common block's fields must survive the merge") + assert(pools.getDeleteFilePoolCount == 2 && pools.getDeleteFilePathPoolCount == 2) + injected.getFileScanTasksList.asScala.foreach { task => + val deletes = pools.getDeleteFilesPool(task.getDeleteFilesIdx).getDeleteFileIndicesList + val paths = deletes.asScala.map(d => + pools.getDeleteFilePathPool(pools.getDeleteFilePool(d).getFilePathIdx)) + assert(paths == Seq(0, 1).map(d => f"s3://bucket/table/data/b00000042-$d-deletes.parquet")) + assert( + pools + .getPartitionDataPool(task.getPartitionDataIdx) + .getValues(0) + .getLiteral + .getIntVal == 42) + } + val planSize = + PlanDataInjector.serializeOperator(scanOp.toBuilder.setIcebergScan(injected).build()).length + assert( + planSize <= commonBytes.length + partition.getSerializedSize + 64, + s"the injected plan ($planSize bytes) should hold only the common block and this partition") + } + + /** + * One partition whose tasks each use all of `deleteFiles` and `residual`, as the driver builds + * it: they are shared arrays and path strings, and the partition's own message indexes into + * them. + */ + private def sharedDeletePartition( + dataFiles: Seq[String], + deleteFiles: Seq[(String, Array[Byte])], + residual: Option[Array[Byte]] = None): IcebergPartitionPlanData = { + val sets = OperatorOuterClass.DeleteFileList.newBuilder() + deleteFiles.indices.foreach(i => sets.addDeleteFileIndices(i)) + val scan = OperatorOuterClass.IcebergScan + .newBuilder() + .setCommon(OperatorOuterClass.IcebergScanCommon.newBuilder().addDeleteFilesPool(sets)) + dataFiles.foreach { path => + val task = OperatorOuterClass.IcebergFileScanTask + .newBuilder() + .setDataFilePath(path) + .setDeleteFilesIdx(0) + residual.foreach(_ => task.setResidualIdx(0)) + scan.addFileScanTasks(task) + } + new IcebergPartitionPlanData( + scan.build().toByteArray, + deleteFiles.map(_._2).toArray, + deleteFiles.map(_._1).toArray, + residual.toArray) + } + + private def deletionVector(dataFile: String, offset: Long): Array[Byte] = + OperatorOuterClass.IcebergDeleteFile + .newBuilder() + .setContentType("POSITION_DELETES") + .setFileFormat("PUFFIN") + .setReferencedDataFile(dataFile) + .setContentOffset(offset) + .setContentSizeInBytes(40) + .setRecordCount(3) + .build() + .toByteArray + + test("partitions sharing delete files resolve each path through their own pool") { + // Two deletion vectors share one Puffin file, and an equality delete has its own. The + // partitions list them in different orders, so each sets its own path indices. + val puffin = "s3://bucket/table/data/dv.puffin" + val eqPath = "s3://bucket/table/data/eq-deletes.parquet" + val dvA = (puffin, deletionVector("a.parquet", 4)) + val dvB = (puffin, deletionVector("b.parquet", 44)) + val eq = ( + eqPath, + OperatorOuterClass.IcebergDeleteFile + .newBuilder() + .setContentType("EQUALITY_DELETES") + .setFileFormat("PARQUET") + .addEqualityIds(1) + .setRecordCount(5) + .build() + .toByteArray) + val scanOp = icebergScanOp("s3://table/metadata/v1.json", scanHashCode = 7) + val key = IcebergPlanDataInjector.getKey(scanOp).get + val (commonBytes, _) = + icebergPlanData("s3://table/metadata/v1.json", 7, Seq("id"), "unused.parquet") + + Seq(Seq(eq, dvA), Seq(dvB, eq, dvA)).foreach { deletes => + val data = sharedDeletePartition(Seq("a.parquet"), deletes) + val injected = PlanDataInjector + .injectPlanData( + parseBasePlan(scanOp.toByteArray), + Map(key -> commonBytes), + Map(key -> data.bytes)) + .getIcebergScan + val pools = injected.getCommon + assert( + pools.getDeleteFilePathPoolList.asScala.toSet == deletes.map(_._1).toSet && + pools.getDeleteFilePathPoolCount == deletes.map(_._1).distinct.size, + s"each path should be pooled once: ${pools.getDeleteFilePathPoolList}") + val task = injected.getFileScanTasks(0) + val resolved = pools + .getDeleteFilesPool(task.getDeleteFilesIdx) + .getDeleteFileIndicesList + .asScala + .map { i => + val file = pools.getDeleteFilePool(i) + (pools.getDeleteFilePathPool(file.getFilePathIdx), file.getContentOffset) + } + val expected = deletes.map { case (path, bytes) => + (path, OperatorOuterClass.IcebergDeleteFile.parseFrom(bytes).getContentOffset) + } + assert(resolved == expected) + } + } + + test("a coalesced task carries a delete file and residual its partitions share once") { + // A coalesce without a shuffle writes all of its input partitions into one task stream. A + // delete file or residual those partitions share is one array, so Java serialization writes + // it once. + val keyMetadata = com.google.protobuf.ByteString.copyFrom(new Array[Byte](8 * 1024)) + val shared = ( + "s3://bucket/table/data/eq-deletes.parquet", + OperatorOuterClass.IcebergDeleteFile + .newBuilder() + .setContentType("EQUALITY_DELETES") + .setFileFormat("PARQUET") + .addEqualityIds(1) + .setKeyMetadata(keyMetadata) + .build() + .toByteArray) + // A large IN list, the residual of a filter every file in the scan gets. + val inList = OperatorOuterClass.IcebergSetPredicate + .newBuilder() + .setColumn("id") + .setOp(OperatorOuterClass.IcebergPredicateOperator.In) + (0 until 2000).foreach { v => + inList.addValues(OperatorOuterClass.IcebergLiteral.newBuilder().setLongVal(v)) + } + val residual = + OperatorOuterClass.IcebergPredicate.newBuilder().setSet(inList).build().toByteArray + val sharedBytes = shared._2.length + residual.length + val partitions = (0 until 16).map { i => + new CometExecPartition( + i, + Array.empty, + Map( + "scan" -> sharedDeletePartition(Seq(s"data-$i.parquet"), Seq(shared), Some(residual)))) + } + def serializedSize(obj: AnyRef): Int = { + val buffer = new java.io.ByteArrayOutputStream() + val out = new java.io.ObjectOutputStream(buffer) + out.writeObject(obj) + out.close() + buffer.size() + } + + val one = serializedSize(partitions.head) + val together = serializedSize(partitions.toArray) + assert(one > sharedBytes, s"a partition should carry the shared entries: $one bytes") + assert( + together < one + sharedBytes, + s"16 partitions took $together bytes; one took $one, the shared entries $sharedBytes") + + val buffer = new java.io.ByteArrayOutputStream() + val out = new java.io.ObjectOutputStream(buffer) + out.writeObject(partitions.toArray) + out.close() + val received = + new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(buffer.toByteArray)) + .readObject() + .asInstanceOf[Array[CometExecPartition]] + .map(_.planDataByKey("scan").asInstanceOf[IcebergPartitionPlanData]) + assert(received.forall(_.deleteFiles.head eq received.head.deleteFiles.head)) + assert(received.forall(_.residuals.head eq received.head.residuals.head)) + received.zip(partitions).foreach { case (got, sent) => + assert(got.bytes.sameElements(sent.planDataByKey("scan").bytes)) + } + val task = OperatorOuterClass.IcebergScan.parseFrom(received.head.bytes) + assert( + task.getCommon + .getResidualPool(task.getFileScanTasks(0).getResidualIdx) + .getSet + .getValuesCount == 2000, + "the task's residual should resolve to the shared IN list") + } + test("each registered injector is reachable by its opStructCase") { // The O(1) lookup keys injectors by opStructCase, so two injectors sharing a kind would // silently shadow one another in the map. Guard that every registered injector resolves back 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 6ca87258b92..756da7badf2 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 @@ -23,7 +23,7 @@ import org.apache.spark.{HashPartitioner, Partition, TaskContext} import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.serializer.JavaSerializer import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext, PartitionPlanData} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -112,7 +112,8 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { override def getPreferredLocations(split: Partition): Seq[String] = Seq(s"host-${split.index}") } - val planData = Map("scan-0" -> Array(Array[Byte](1), Array[Byte](2))) + val planData = + Map("scan-0" -> PartitionPlanData.wrap(Array(Array[Byte](1), Array[Byte](2)))) val remote = new CometNativeShuffleInputRDD( spark.sparkContext, Seq(upstream), @@ -128,7 +129,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { local.partitions.foreach { part => val input = part.asInstanceOf[CometNativeShuffleInputPartition] assert(input.inputPartitions.toSeq == Seq(upstream.partitions(part.index))) - assert(input.planDataByKey("scan-0").sameElements(planData("scan-0")(part.index))) + assert(input.planDataByKey("scan-0") eq planData("scan-0")(part.index)) assert(local.preferredLocations(part) == Seq(s"host-${part.index}")) val context = TaskContext.empty() @@ -136,7 +137,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { val iterator = local.iterator(part, context).asInstanceOf[CometNativeShuffleInputIterator] assert(iterator.partitionIndex == part.index) assert(iterator.inputObjects.toSeq == Seq(null)) - assert(iterator.planDataByKey("scan-0").sameElements(planData("scan-0")(part.index))) + assert(iterator.planDataByKey("scan-0") eq planData("scan-0")(part.index)) } finally { context.markTaskCompleted(None) } @@ -207,7 +208,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { CometNativeShuffleInputRDD, CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]) = { val perPartitionByKey = - Map("scan-0" -> Array.fill(numPartitions)(new Array[Byte](1024))) + Map("scan-0" -> Array.fill(numPartitions)(PartitionPlanData(new Array[Byte](1024)))) val childMetricNode = CometMetricNode(Map.empty) val writerMetrics = Map( "spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "disk spilled bytes"), @@ -261,6 +262,6 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { // Each task's Partition object still carries its own slice so the writer can inject plan data. val part = largeRdd.partitions(7).asInstanceOf[CometNativeShuffleInputPartition] assert(part.planDataByKey.keySet == Set("scan-0")) - assert(part.planDataByKey("scan-0").length == 1024) + assert(part.planDataByKey("scan-0").bytes.length == 1024) } }