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 13f811a68c6..d64492c1658 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -194,6 +194,15 @@ For a custom S3-compatible endpoint, configure the catalog with the endpoint, pa These `s3.*` storage properties are not specific to the Hive catalog shown here. When `s3.access-key-id` / `s3.secret-access-key` are omitted, credentials come from the standard AWS chain (environment variables, instance profiles, and so on). The region is not auto-detected: when neither the catalog (`client.region` or `s3.region`) nor the executor environment (`AWS_REGION` or `AWS_DEFAULT_REGION`) supplies one, Comet uses `us-east-1`, so set it for AWS buckets in any other region. If your REST catalog vends temporary credentials, the native reader does not consume them automatically, and wiring that requires the credential provider bridge. See Iceberg's [S3 FileIO](https://iceberg.apache.org/docs/latest/aws/#s3-fileio) docs for the full property list, and [S3 Credential Providers](s3-credential-providers.md) for vended or per-request credentials. +### 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 that many partitions +use, such as an unpartitioned equality delete, 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 28d94e78bf3..7761ff838bf 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -388,7 +388,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 @@ -414,8 +419,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 @@ -432,7 +437,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 4070d785367..483de3f2b5a 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 @@ -32,7 +32,7 @@ import org.json4s.jackson.JsonMethods._ 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._ @@ -306,8 +306,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") @@ -315,13 +314,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 => @@ -334,7 +327,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. @@ -343,8 +338,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( @@ -353,7 +348,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 @@ -409,7 +403,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()) } /** @@ -427,6 +421,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") @@ -586,7 +581,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) @@ -950,18 +945,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 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 @@ -970,40 +971,75 @@ 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 + + // 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]() + val residuals = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Int]() + + 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 build(scan: OperatorOuterClass.IcebergScan.Builder): IcebergPartitionPlanData = { + if (partitionData.nonEmpty || deleteFileSets.nonEmpty || residuals.nonEmpty) { + scan.setCommon(builder) + } + new IcebergPartitionPlanData( + scan.build().toByteArray, + deleteFileBytes.toArray, + deleteFilePaths.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. @@ -1037,7 +1073,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 @@ -1090,6 +1126,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]) @@ -1150,25 +1187,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) @@ -1181,10 +1210,10 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit pageIndexUnsupportedColumns) residualExprOpt.foreach { residual => - val residualIdx = residualToPoolIndex.getOrElseUpdate( + val residualIdx = pools.residuals.getOrElseUpdate( residual, { - val idx = residualToPoolIndex.size - commonBuilder.addResidualPool(residual) + val idx = pools.residuals.size + pools.builder.addResidualPool(residual) idx }) taskBuilder.setResidualIdx(residualIdx) @@ -1197,7 +1226,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit taskBuilder, commonBuilder, partitionSpecToPoolIndex, - partitionDataToPoolIndex) + pools.builder, + pools.partitionData) metadata.nameMapping.foreach { nm => val nmIdx = nameMappingToPoolIndex.getOrElseUpdate( @@ -1214,7 +1244,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 @@ -1233,7 +1267,8 @@ 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) } case other => throw new IllegalStateException( @@ -1247,11 +1282,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" @@ -1293,30 +1327,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 @@ -1328,13 +1341,18 @@ 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( + " Delete files, each held once and shared by the partitions that use it: " + + s"${sharedDeleteFiles.size} files, ${sharedDeleteFilePaths.size} paths, " + + s"$sharedDeleteFileBytes bytes") + logInfo( + s" Per-partition messages, without delete files: ${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 31e518cfcf8..cd6ad426e30 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..2e28ec8163c --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala @@ -0,0 +1,104 @@ +/* + * 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, residuals and + * delete-file sets in `common`. Its delete files are kept apart, because a delete file is often + * used by many partitions: the driver serializes each delete file once, and every partition that + * uses it holds the same array and the same path string. A coalesced task then carries each + * delete file once, and the driver keeps one copy of it. + * + * `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. + */ +final class IcebergPartitionPlanData( + val scan: Array[Byte], + val deleteFiles: Array[Array[Byte]], + val deleteFilePaths: Array[String]) + extends PartitionPlanData { + + require( + deleteFiles.length == deleteFilePaths.length, + s"${deleteFiles.length} delete files but ${deleteFilePaths.length} paths") + + override def bytes: Array[Byte] = { + if (deleteFiles.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)) + } + 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 cf8b6ffed9b..2def9b81bc5 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 dd7c59f6556..98af5872315 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 b8835212129..5e3896f75d9 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -45,6 +45,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 @@ -773,10 +774,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 " + @@ -800,30 +801,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 @@ -873,12 +970,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). @@ -888,22 +983,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") @@ -1166,6 +1365,86 @@ 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") + checkSparkAnswer(spark.table("test_cat.db.global_eq_delete").coalesce(1)) + + 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) @@ -3742,10 +4021,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") } } } @@ -4888,9 +5165,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") @@ -5018,27 +5293,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") } @@ -5690,9 +5968,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, @@ -5959,9 +6235,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 @@ -5996,8 +6270,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") } @@ -6034,9 +6307,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 => @@ -6500,9 +6771,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 6ff5daf27fe..fffbd71dba4 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 @@ -242,13 +242,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 => @@ -262,11 +263,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..10253de1820 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,222 @@ 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`, as the driver builds it: the delete + * files 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])]): 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 => + scan.addFileScanTasks( + OperatorOuterClass.IcebergFileScanTask + .newBuilder() + .setDataFilePath(path) + .setDeleteFilesIdx(0)) + } + new IcebergPartitionPlanData( + scan.build().toByteArray, + deleteFiles.map(_._2).toArray, + deleteFiles.map(_._1).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 its partitions share once") { + // A coalesce without a shuffle writes all of its input partitions into one task stream. A + // delete file 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) + val partitions = (0 until 16).map { i => + new CometExecPartition( + i, + Array.empty, + Map("scan" -> sharedDeletePartition(Seq(s"data-$i.parquet"), Seq(shared)))) + } + 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 > shared._2.length, s"a partition should carry the delete file: $one bytes") + assert( + together < one + shared._2.length, + s"16 partitions took $together bytes; one took $one, the delete file ${shared._2.length}") + + 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)) + received.zip(partitions).foreach { case (got, sent) => + assert(got.bytes.sameElements(sent.planDataByKey("scan").bytes)) + } + } + 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) } }