From bd2f9de7d01db8aa0b013acaa05b4b25fc7122b7 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 4 Nov 2026 10:19:58 -0800 Subject: [PATCH 1/8] fix: ship Iceberg scan pool entries per partition so task memory does not grow with the table Native Iceberg scans put the whole scan's delete files, delete-file paths and sets, partition values and residuals in the common block. Every task deserialized it from the task binary, and the injector copied it into each partition's native plan, so tasks on large merge-on-read tables ran out of heap. Each partition now carries only the pool entries its tasks use, in IcebergScan.own_pools (a new IcebergFilePools message), and the common block keeps only the small pools and settings. The driver builds each partition's message lazily, while Spark serializes its task (PartitionPlanData / PlanDataSlice), so it never holds every partition's message at once. The executor merges the partition's pools into the common block before planning, and native rejects a plan that still carries own_pools. A coalesce without a shuffle repeats shared delete files once per input partition; this is documented and pinned by a test. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../sql/comet/CometDeltaNativeScanExec.scala | 2 +- docs/source/user-guide/latest/iceberg.md | 10 + native/core/src/execution/planner.rs | 26 ++ native/proto/src/proto/operator.proto | 34 +- .../operator/CometIcebergNativeScan.scala | 141 ++++--- .../serde/operator/IcebergScanPools.scala | 181 ++++++++ .../apache/spark/sql/comet/CometExecRDD.scala | 13 +- .../comet/CometIcebergNativeScanExec.scala | 22 +- .../spark/sql/comet/CometNativeScanExec.scala | 2 +- .../spark/sql/comet/PartitionPlanData.scala | 89 ++++ .../shuffle/CometNativeShuffleInputRDD.scala | 22 +- .../apache/spark/sql/comet/operators.scala | 34 +- .../merge_on_read_partition_layout.sql | 146 +++++++ .../apache/comet/CometFuzzIcebergSuite.scala | 11 +- .../comet/CometIcebergNativeSuite.scala | 377 ++++++++++++++--- .../operator/IcebergScanPoolsSuite.scala | 397 ++++++++++++++++++ .../operator/IcebergScanPoolsTestUtils.scala | 169 ++++++++ .../comet/CometScanWithPlanDataSuite.scala | 7 +- .../sql/comet/PlanDataInjectorSuite.scala | 82 ++++ .../CometNativeShuffleInputRDDSuite.scala | 57 ++- 22 files changed, 1662 insertions(+), 162 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala create mode 100644 spark/src/test/resources/sql-tests/iceberg/merge_on_read_partition_layout.sql create mode 100644 spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala create mode 100644 spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 478339062b0..69275ac3ae4 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -657,6 +657,7 @@ jobs: org.apache.comet.objectstore.AbfsAuthParitySuite org.apache.comet.objectstore.AbfsScanFallbackSuite org.apache.comet.serde.operator.CometIcebergNativeScanSuite + org.apache.comet.serde.operator.IcebergScanPoolsSuite org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite org.apache.spark.sql.CometCollationSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 3b4b49c728a..9dd44265b2d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -292,6 +292,7 @@ jobs: org.apache.comet.objectstore.AbfsAuthParitySuite org.apache.comet.objectstore.AbfsScanFallbackSuite org.apache.comet.serde.operator.CometIcebergNativeScanSuite + org.apache.comet.serde.operator.IcebergScanPoolsSuite org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite org.apache.spark.sql.CometCollationSuite 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..f2386a2eeb7 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), 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..8e241fe0a3f 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -194,6 +194,16 @@ 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 coalesce without a shuffle +(`coalesce(n)`, or a write that coalesces its input) combines several input partitions into one +task, and that task then carries a copy of a shared delete file for every input partition that +uses it. Scans where many partitions share delete files, such as tables with unpartitioned +equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning +instead of coalescing, avoids this. + ### Current limitations The following scenarios will fall back to the JVM Iceberg reader: diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 603bba872bc..d265309fbd5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1967,6 +1967,14 @@ impl PhysicalPlanner { OpStruct::IcebergScan(scan) => { // Extract common data and single partition's file tasks // Per-partition injection happens in Scala before sending to native + if scan.own_pools.is_some() { + // The tasks would index pools that are not in common. + return Err(GeneralError( + "IcebergScan still carries its per-partition pools: the JVM did not \ + merge them into common" + .into(), + )); + } let common = scan .common .as_ref() @@ -3606,6 +3614,24 @@ mod tests { } } + #[test] + fn iceberg_scan_rejects_unmerged_partition_pools() { + let op = spark_operator::Operator { + op_struct: Some(OpStruct::IcebergScan(spark_operator::IcebergScan { + common: Some(spark_operator::IcebergScanCommon::default()), + file_scan_tasks: vec![], + own_pools: Some(spark_operator::IcebergFilePools::default()), + })), + ..Default::default() + }; + let planner = PhysicalPlanner::default(); + let err = planner.create_plan(&op, &mut vec![], 1).unwrap_err(); + assert!( + format!("{err}").contains("did not merge them into common"), + "got: {err}" + ); + } + #[test] fn delta_scan_errors_without_delta_feature() { let op = delta_spark_envelope(spark_operator::DeltaSparkScan { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 28d94e78bf3..7547ba7222a 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -377,7 +377,8 @@ enum IcebergPredicateOperator { NotIn = 14; } -// Common data shared by all partitions in split mode (sent once, captured in closure) +// Common data shared by all partitions of a scan. Each task deserializes its own copy, so it +// holds no per-file pool entries (see IcebergScan.own_pools). message IcebergScanCommon { // Catalog-specific configuration for FileIO (credentials, S3/GCS config, etc.) map catalog_properties = 1; @@ -388,7 +389,10 @@ 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 the driver ships, the five per-file pools + // (partition_data_pool, delete_files_pool, residual_pool, delete_file_pool, + // delete_file_path_pool) are empty: each partition ships the entries its tasks use in + // IcebergScan.own_pools, and the JVM merges them in before the plan reaches native. 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 +418,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,11 +436,28 @@ message IcebergScanCommon { } message IcebergScan { - // Common data shared across partitions (pools, metadata, catalog props) + // Common data shared across partitions (pools, metadata, catalog props). Unset in the + // per-partition message the driver ships. In the plan that reaches native it holds everything + // the partition's tasks index into. IcebergScanCommon common = 1; // Single partition's file scan tasks repeated IcebergFileScanTask file_scan_tasks = 2; + + // Only in the per-partition message the driver ships: the per-file pool entries this + // partition's tasks use. The JVM merges them into common before the plan reaches native, and + // native rejects a plan that still has them (see IcebergScanPools). + IcebergFilePools own_pools = 3; +} + +// The five per-file pools of IcebergScanCommon, holding one partition's entries. Field meanings +// match the IcebergScanCommon fields of the same name. +message IcebergFilePools { + repeated PartitionData partition_data_pool = 1; + repeated DeleteFileList delete_files_pool = 2; + repeated IcebergPredicate residual_pool = 3; + repeated IcebergDeleteFile delete_file_pool = 4; + repeated string delete_file_path_pool = 5; } // Common data shared by all partitions for native Lance scans. @@ -496,7 +517,8 @@ message IcebergFileScanTask { // decodes this directly; no KMS unwrap is needed on the native side. optional bytes key_metadata = 6; - // Indices into IcebergScan deduplication pools + // Indices into the deduplication pools. The four non-file pools are in the common block; the + // per-file pools are in IcebergScan.own_pools until the JVM merges them (see IcebergScanPools). uint32 schema_idx = 15; optional uint32 partition_spec_idx = 17; optional uint32 name_mapping_idx = 18; 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..614cd1d9cdb 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, 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._ @@ -950,17 +950,21 @@ 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. + * 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 (schemas, specs, + * partition data, residuals, delete files) 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 has to stay small. The per-file pool entries + * (delete files, their paths and sets, partition values, residuals) ship with each partition + * whose tasks use them, and the common block holds none of them (see IcebergScanPools). Each + * partition's block is built when its task is serialized. Keeping every delete file in the + * common block made each partition'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 + * granularity, so each partition stores it once and references it many times, as Iceberg's own * DeleteFileIndex does. * * @param scanExec @@ -970,12 +974,12 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * @param metadata * Pre-extracted Iceberg metadata from CometScanRule * @return - * Tuple of (commonBytes, perPartitionBytes) for native execution + * The common block's bytes, and each partition's bytes built when its task is serialized */ def serializePartitions( scanExec: BatchScanExec, output: Seq[Attribute], - metadata: CometIcebergNativeScanMetadata): (Array[Byte], Array[Array[Byte]]) = { + metadata: CometIcebergNativeScanMetadata): (Array[Byte], PartitionPlanData) = { val commonBuilder = OperatorOuterClass.IcebergScanCommon.newBuilder() @@ -1265,76 +1269,75 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } - // Per-pool byte sizes to diagnose an oversized common message. Sizes sum as Long because a - // single pool at or past protobuf's 2 GiB message limit overflows the int getSerializedSize, - // and the logging runs before toByteArray so the breakdown survives even if that allocation - // fails. String pools carry JSON or file paths, whose serialized size is the UTF-8 length. + // The per-file pool entries ship with each partition that uses them, so the common block + // every task deserializes doesn't grow with the number of files. + val original = commonBuilder.build() + val common = IcebergScanPools.commonBlock(original) + + // Per-pool byte sizes of the whole scan, to diagnose large plan data. Sizes sum as Long because + // a single pool at or past protobuf's 2 GiB message limit overflows the int getSerializedSize, + // and they are logged before any toByteArray so the breakdown survives a failed allocation. + // String pools carry JSON or file paths, whose serialized size is the UTF-8 length. def sumSizes(sizes: Iterator[Int]): Long = sizes.map(_.toLong).sum def sumStrBytes(strings: mutable.Buffer[String]): Long = strings.iterator.map(_.getBytes(UTF_8).length.toLong).sum - val commonPoolBytes: Seq[(String, Int, Long)] = Seq( - ( - "schema", - commonBuilder.getSchemaPoolCount, - sumStrBytes(commonBuilder.getSchemaPoolList.asScala)), + val scanPoolBytes: Seq[(String, Int, Long)] = Seq( + ("schema", original.getSchemaPoolCount, sumStrBytes(original.getSchemaPoolList.asScala)), ( "partition_type", - commonBuilder.getPartitionTypePoolCount, - sumStrBytes(commonBuilder.getPartitionTypePoolList.asScala)), + original.getPartitionTypePoolCount, + sumStrBytes(original.getPartitionTypePoolList.asScala)), ( "partition_spec", - commonBuilder.getPartitionSpecPoolCount, - sumStrBytes(commonBuilder.getPartitionSpecPoolList.asScala)), + original.getPartitionSpecPoolCount, + sumStrBytes(original.getPartitionSpecPoolList.asScala)), ( "name_mapping", - commonBuilder.getNameMappingPoolCount, - sumStrBytes(commonBuilder.getNameMappingPoolList.asScala)), + original.getNameMappingPoolCount, + sumStrBytes(original.getNameMappingPoolList.asScala)), ( "project_field_ids", - commonBuilder.getProjectFieldIdsPoolCount, - sumSizes( - commonBuilder.getProjectFieldIdsPoolList.asScala.iterator.map(_.getSerializedSize))), + original.getProjectFieldIdsPoolCount, + sumSizes(original.getProjectFieldIdsPoolList.asScala.iterator.map(_.getSerializedSize))), ( "partition_data", - commonBuilder.getPartitionDataPoolCount, - sumSizes( - commonBuilder.getPartitionDataPoolList.asScala.iterator.map(_.getSerializedSize))), + original.getPartitionDataPoolCount, + sumSizes(original.getPartitionDataPoolList.asScala.iterator.map(_.getSerializedSize))), ( "delete_file", - commonBuilder.getDeleteFilePoolCount, - sumSizes(commonBuilder.getDeleteFilePoolList.asScala.iterator.map(_.getSerializedSize))), + original.getDeleteFilePoolCount, + sumSizes(original.getDeleteFilePoolList.asScala.iterator.map(_.getSerializedSize))), ( "delete_file_path", - commonBuilder.getDeleteFilePathPoolCount, - sumStrBytes(commonBuilder.getDeleteFilePathPoolList.asScala)), + original.getDeleteFilePathPoolCount, + sumStrBytes(original.getDeleteFilePathPoolList.asScala)), ( "delete_files_set", - commonBuilder.getDeleteFilesPoolCount, - sumSizes(commonBuilder.getDeleteFilesPoolList.asScala.iterator.map(_.getSerializedSize))), + original.getDeleteFilesPoolCount, + sumSizes(original.getDeleteFilesPoolList.asScala.iterator.map(_.getSerializedSize))), ( "residual", - commonBuilder.getResidualPoolCount, - sumSizes(commonBuilder.getResidualPoolList.asScala.iterator.map(_.getSerializedSize)))) - - val perPartitionSizes = perPartitionBuilders.map(_.getSerializedSize.toLong) - val perPartitionTotal = perPartitionSizes.sum - val perPartitionMax = if (perPartitionSizes.isEmpty) 0L else perPartitionSizes.max + original.getResidualPoolCount, + sumSizes(original.getResidualPoolList.asScala.iterator.map(_.getSerializedSize)))) logInfo(s"IcebergScan: $totalTasks tasks, ${allPoolSizes.size} pools ($avgDedup% avg dedup)") logInfo( - " Common pools (unique/bytes): " + - commonPoolBytes + " Scan pools (unique/bytes): " + + scanPoolBytes .map { case (name, count, bytes) => s"$name=$count/$bytes" } .mkString(", ") + - s"; total=${commonPoolBytes.map(_._3).sum} bytes") - logInfo( - s" Per-partition messages: ${perPartitionBuilders.size}, " + - s"total=$perPartitionTotal bytes, max=$perPartitionMax bytes") - - val commonBytes = commonBuilder.build().toByteArray - val perPartitionBytes = perPartitionBuilders.map(_.toByteArray).toArray - - (commonBytes, perPartitionBytes) + s"; total=${scanPoolBytes.map(_._3).sum} bytes") + + val commonBytes = common.toByteArray + // Tasks only, still indexing the full pools. Each partition's message is built from these + // when its task is serialized, so the driver holds every partition's tasks once and never + // every partition's message at the same time. + val taskBytes = perPartitionBuilders.map(_.toByteArray).toArray + logInfo(s" Common block: ${commonBytes.length} bytes. Per-partition messages: " + + s"${taskBytes.length}, task bytes total=${taskBytes.iterator.map(_.length.toLong).sum}, " + + s"max=${if (taskBytes.isEmpty) 0 else taskBytes.iterator.map(_.length).max}") + + (commonBytes, new IcebergPartitionPlanData(original, taskBytes)) } override def createExec(nativeOp: Operator, op: CometBatchScanExec): CometNativeExec = { @@ -1355,3 +1358,35 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit CometIcebergNativeScanExec(nativeOp, op.wrapped, op.session, metadataLocation, metadata) } } + +/** + * A scan's per-partition messages, each built when its task is serialized: the partition's tasks + * plus the per-file pool entries they use (see IcebergScanPools.partitionPayload). + */ +private[operator] final class IcebergPartitionPlanData( + private val original: OperatorOuterClass.IcebergScanCommon, + private val taskBytes: Array[Array[Byte]]) + extends PartitionPlanData { + + override def length: Int = taskBytes.length + + // Same table state read into the same partitions, compared without building any partition. + override def sameAs(other: PartitionPlanData): Boolean = other match { + case that: IcebergPartitionPlanData => + that.original == original && that.taskBytes.length == taskBytes.length && + taskBytes.indices.forall(i => java.util.Arrays.equals(taskBytes(i), that.taskBytes(i))) + case _ => false + } + + // With no per-file entries in the scan, the stored task bytes already are each message. + private val hasFileEntries = IcebergScanPools.hasFileEntries(original) + + override def apply(i: Int): Array[Byte] = + if (!hasFileEntries) { + taskBytes(i) + } else { + IcebergScanPools + .partitionPayload(original, OperatorOuterClass.IcebergScan.parseFrom(taskBytes(i))) + .toByteArray + } +} diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala new file mode 100644 index 00000000000..5dd73f95d1b --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala @@ -0,0 +1,181 @@ +/* + * 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.serde.operator + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergFilePools, IcebergFileScanTask, IcebergScan, IcebergScanCommon} + +/** + * Decides where an Iceberg scan's pool entries ship. + * + * Five pools grow with the number of files in the scan: delete file paths, delete files, delete + * file sets, partition values and residuals. Their entries ship with each Spark partition whose + * tasks use them, as Iceberg's own Spark reader ships each task's delete files with the task. So + * a task never holds entries that only other partitions use, and the common block every task + * deserializes holds only the small pools (schemas, specs, name mappings, field ids) and the + * scan's settings. + * + * On the driver, [[partitionPayload]] builds one partition's message from the full pools. On the + * executor, [[forPartition]] merges it with the common block into the scan the native plan runs. + * + * Known limit: a coalesce without a shuffle puts several partitions' messages in one task, and + * each carries its own copy of the entries they share. Delete files that many partitions use, + * such as unpartitioned equality deletes, are then repeated once per coalesced partition. + */ +private[apache] object IcebergScanPools { + + /** The common block every task gets: the scan with the five per-file pools emptied. */ + def commonBlock(original: IcebergScanCommon): IcebergScanCommon = + original.toBuilder + .clearDeleteFilePathPool() + .clearDeleteFilePool() + .clearDeleteFilesPool() + .clearPartitionDataPool() + .clearResidualPool() + .build() + + /** + * One partition's message, built on the driver. `tasks` indexes into `original`. The result + * carries the tasks, renumbered, and in `own_pools` only the per-file entries they use. + */ + def partitionPayload(original: IcebergScanCommon, tasks: IcebergScan): IcebergScan = { + val (pools, renumbered) = prune(original, tasks) + val result = tasks.toBuilder.clearFileScanTasks().addAllFileScanTasks(renumbered.asJava) + if (pools.getDeleteFilePathPoolCount > 0 || pools.getDeleteFilePoolCount > 0 || + pools.getDeleteFilesPoolCount > 0 || pools.getPartitionDataPoolCount > 0 || + pools.getResidualPoolCount > 0) { + result.setOwnPools(pools) + } + result.build() + } + + /** Whether `common` still holds per-file pool entries, which must ship with partitions. */ + def hasFileEntries(common: IcebergScanCommon): Boolean = + common.getDeleteFilePathPoolCount > 0 || common.getDeleteFilePoolCount > 0 || + common.getDeleteFilesPoolCount > 0 || common.getPartitionDataPoolCount > 0 || + common.getResidualPoolCount > 0 + + /** + * The scan one partition's native plan runs, built on the executor from the common block and + * the partition's message. The message already holds just the per-file entries its tasks use, + * numbered from 0, so this only adds them to the common block. + */ + def forPartition(common: IcebergScanCommon, partition: IcebergScan): IcebergScan = { + if (hasFileEntries(common)) { + throw new IllegalStateException( + "Iceberg common block holds per-file pool entries; they must ship with each partition") + } + val own = partition.getOwnPools + val merged = common.toBuilder + .addAllDeleteFilePathPool(own.getDeleteFilePathPoolList) + .addAllDeleteFilePool(own.getDeleteFilePoolList) + .addAllDeleteFilesPool(own.getDeleteFilesPoolList) + .addAllPartitionDataPool(own.getPartitionDataPoolList) + .addAllResidualPool(own.getResidualPoolList) + partition.toBuilder.clearOwnPools().setCommon(merged).build() + } + + /** Entry i of one of the scan's pools, failing with the pool's name when i is out of range. */ + private def entry[T](name: String, i: Int, size: Int)(get: Int => T): T = { + if (i < 0 || i >= size) { + throw new IllegalStateException( + s"Iceberg $name index $i is out of range: the pool has $size entries") + } + get(i) + } + + /** + * The per-file pool entries of `common` that `tasks` reach, and the tasks renumbered to index + * them. The work follows the tasks, not the size of the pools. + */ + private def prune( + common: IcebergScanCommon, + tasks: IcebergScan): (IcebergFilePools.Builder, Seq[IcebergFileScanTask]) = { + val pools = IcebergFilePools.newBuilder() + val pathIdx = mutable.HashMap.empty[Int, Int] + val deleteIdx = mutable.HashMap.empty[Int, Int] + val setIdx = mutable.HashMap.empty[Int, Int] + val partitionDataIdx = mutable.HashMap.empty[Int, Int] + val residualIdx = mutable.HashMap.empty[Int, Int] + + def path(i: Int): Int = + pathIdx.getOrElseUpdate( + i, { + pools.addDeleteFilePathPool( + entry("delete_file_path", i, common.getDeleteFilePathPoolCount)( + common.getDeleteFilePathPool)) + pools.getDeleteFilePathPoolCount - 1 + }) + + def delete(i: Int): Int = + deleteIdx.getOrElseUpdate( + i, { + val deleteFile = + entry("delete_file", i, common.getDeleteFilePoolCount)(common.getDeleteFilePool) + pools.addDeleteFilePool( + deleteFile.toBuilder.setFilePathIdx(path(deleteFile.getFilePathIdx))) + pools.getDeleteFilePoolCount - 1 + }) + + def set(i: Int): Int = + setIdx.getOrElseUpdate( + i, { + val set = + entry("delete_files_set", i, common.getDeleteFilesPoolCount)( + common.getDeleteFilesPool) + val remapped = DeleteFileList.newBuilder() + (0 until set.getDeleteFileIndicesCount).foreach { j => + remapped.addDeleteFileIndices(delete(set.getDeleteFileIndices(j))) + } + pools.addDeleteFilesPool(remapped) + pools.getDeleteFilesPoolCount - 1 + }) + + def partitionValue(i: Int): Int = + partitionDataIdx.getOrElseUpdate( + i, { + pools.addPartitionDataPool( + entry("partition_data", i, common.getPartitionDataPoolCount)( + common.getPartitionDataPool)) + pools.getPartitionDataPoolCount - 1 + }) + + def residual(i: Int): Int = + residualIdx.getOrElseUpdate( + i, { + pools.addResidualPool( + entry("residual", i, common.getResidualPoolCount)(common.getResidualPool)) + pools.getResidualPoolCount - 1 + }) + + val renumbered = tasks.getFileScanTasksList.asScala.toSeq.map { task => + val builder = task.toBuilder + if (task.hasResidualIdx) builder.setResidualIdx(residual(task.getResidualIdx)) + if (task.hasPartitionDataIdx) { + builder.setPartitionDataIdx(partitionValue(task.getPartitionDataIdx)) + } + if (task.hasDeleteFilesIdx) builder.setDeleteFilesIdx(set(task.getDeleteFilesIdx)) + builder.build() + } + (pools, renumbered) + } +} 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..3b6ef900360 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 @@ -33,13 +33,16 @@ import org.apache.comet.{CometExecIterator, CometRuntimeException, CometShuffleB /** * Partition that carries per-partition planning data, avoiding closure capture of all partitions. + * The data is built while the task is serialized (see [[PlanDataSlice]]). */ private[spark] class CometExecPartition( override val index: Int, val inputPartitions: Array[Partition], - val planDataByKey: Map[String, Array[Byte]], + planData: PlanDataSlice, val filePaths: Seq[String] = Seq.empty) - extends Partition + extends Partition { + def planDataByKey: Map[String, Array[Byte]] = planData.get +} /** * Unified RDD for Comet native execution. Non-shuffle input slots are `RDD[ArrowArrayStream]` @@ -58,7 +61,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, PartitionPlanData], serializedPlan: Array[Byte], planFingerprint: Long, defaultNumPartitions: Int, @@ -98,7 +101,7 @@ private[spark] class CometExecRDD( override protected def getPartitions: Array[Partition] = { (0 until numPartitions).map { idx => val inputParts = inputRDDs.map(_.partitions(idx)).toArray - val planData = perPartitionByKey.map { case (key, arr) => key -> arr(idx) } + val planData = new PlanDataSlice(perPartitionByKey, idx) val fp = if (perPartitionFilePaths.length > idx) perPartitionFilePaths(idx) else Seq.empty[String] new CometExecPartition(idx, inputParts, planData, fp) @@ -236,7 +239,7 @@ object CometExecRDD { sc: SparkContext, inputRDDs: Seq[RDD[_]], commonByKey: Map[String, Array[Byte]], - perPartitionByKey: Map[String, Array[Array[Byte]]], + perPartitionByKey: Map[String, 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..9e4afca90b6 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 @@ -79,7 +79,7 @@ case class CometIcebergNativeScanExec( /** * Lazy partition serialization, deferred until execution time. Triggered from `commonData` / - * `perPartitionData` (via `PlanDataInjector.findAllPlanData`) and from + * `partitionPlanData` (via `PlanDataInjector.findAllPlanData`) and from * `LazyIcebergMetric.value` (via Iceberg planning metrics). Lazy val semantics ensure single * evaluation across entry points. * @@ -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], 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,11 +122,19 @@ case class CometIcebergNativeScanExec( def commonData: Array[Byte] = serializedPartitionData._1 - def perPartitionData: Array[Array[Byte]] = serializedPartitionData._2 + /** Each partition's message, built when its task is serialized. */ + def partitionPlanData: PartitionPlanData = serializedPartitionData._2 + + /** + * Every partition's message, all built at once. For tests only: execution must use + * [[partitionPlanData]], which builds each one when its task is serialized. + */ + private[apache] def buildAllPartitionDataForTesting(): Array[Array[Byte]] = + partitionPlanData.toArray // numPartitions for execution - derived from actual DPP-filtered partitions // Only accessed during execution, not planning - def numPartitions: Int = perPartitionData.length + def numPartitions: Int = partitionPlanData.length // Like BatchScanExec on Spark 3.5 and later without a key-grouped partitioning, the scan does // not know its partition count while planning, because dynamic partition pruning decides which @@ -134,7 +142,7 @@ case class CometIcebergNativeScanExec( // optimizes a stage (ValidateRequirements after CoalesceShufflePartitions), which is before // CometPlanAdaptiveDynamicPruningFilters converts the DPP placeholders, so serializing the // tasks here would run a placeholder subquery. At execution, CometNativeExec sizes the native - // RDD from perPartitionData instead. + // RDD from partitionPlanData instead. override lazy val outputPartitioning: Partitioning = UnknownPartitioning(0) override lazy val outputOrdering: Seq[SortOrder] = Nil @@ -294,10 +302,10 @@ case class CometIcebergNativeScanExec( sparkContext, inputRDDs = Seq.empty, commonByKey = Map(injectorKey -> commonData), - perPartitionByKey = Map(injectorKey -> perPartitionData), + perPartitionByKey = Map(injectorKey -> partitionPlanData), serializedPlan = serializedPlan, planFingerprint = PlanDataInjector.planFingerprint(serializedPlan), - defaultNumPartitions = perPartitionData.length, + defaultNumPartitions = partitionPlanData.length, numOutputCols = output.length, nativeMetrics = nativeMetrics, subqueries = Seq.empty, 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..484d7919df7 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), 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..f18e19f2648 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala @@ -0,0 +1,89 @@ +/* + * 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 java.io.{ObjectInputStream, ObjectOutputStream} + +/** + * One scan's planning data for each partition. A scan can build a partition's bytes only when its + * task is serialized, the way Spark serializes each task's input partition, so the driver never + * holds every partition's bytes at once. + */ +private[apache] trait PartitionPlanData { + def length: Int + + /** Partition i's bytes. Must be safe to call from any thread, and more than once. */ + def apply(i: Int): Array[Byte] + + /** + * Whether `other` holds the same data, decided without building any partition. The default only + * knows identity. + */ + def sameAs(other: PartitionPlanData): Boolean = this eq other + + /** Every partition's bytes. For tests and diagnostics only: it builds them all. */ + def toArray: Array[Array[Byte]] = Array.tabulate(length)(apply) +} + +private[apache] object PartitionPlanData { + + private final class Built(val bytes: Array[Array[Byte]]) extends PartitionPlanData { + override def length: Int = bytes.length + override def apply(i: Int): Array[Byte] = bytes(i) + override def sameAs(other: PartitionPlanData): Boolean = other match { + case that: Built => + that.bytes.length == bytes.length && + bytes.indices.forall(i => java.util.Arrays.equals(bytes(i), that.bytes(i))) + case _ => false + } + } + + /** Data already built for every partition. */ + def apply(bytes: Array[Array[Byte]]): PartitionPlanData = new Built(bytes) + + /** Whether two scans carry the same data. */ + def same(a: PartitionPlanData, b: PartitionPlanData): Boolean = (a eq b) || a.sameAs(b) +} + +/** + * The planning data one task needs, keyed by scan. On the driver it holds only a reference to + * every scan's data and builds this partition's bytes while the task is serialized. On the + * executor it holds the bytes it received. + */ +private[comet] final class PlanDataSlice( + @transient private val source: Map[String, PartitionPlanData], + index: Int) + extends Serializable { + + @transient private var received: Map[String, Array[Byte]] = _ + + def get: Map[String, Array[Byte]] = + if (received != null) received else source.map { case (key, data) => key -> data(index) } + + private def writeObject(out: ObjectOutputStream): Unit = { + out.defaultWriteObject() + out.writeObject(get) + } + + private def readObject(in: ObjectInputStream): Unit = { + in.defaultReadObject() + received = in.readObject().asInstanceOf[Map[String, Array[Byte]]] + } +} 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..93468a7ba44 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, PlanDataSlice} 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, PartitionPlanData] = Map.empty, positionalRoundRobin: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, @@ -98,13 +98,11 @@ private[shuffle] class CometNativeShuffleInputRDD( // `leafRdd.partitions` on the executor, which would otherwise trigger getPartitions and // hit the @transient-null trap (e.g. CometExecRDD.perPartitionByKey). val inputParts = inputRDDs.map(_.partitions(i)).toArray - // Slice this partition's plan data off the @transient full map here on the driver. Carrying - // only the per-partition slice on the Partition object (serialized per task) keeps the full - // O(numPartitions) map out of the broadcast task binary, which otherwise blows the 2GB - // ByteArrayOutputStream limit on jobs with tens of millions of partitions. Mirrors - // CometExecRDD.getPartitions. - val planDataByKey = perPartitionByKey.map { case (key, arr) => key -> arr(i) } - new CometNativeShuffleInputPartition(i, inputParts, planDataByKey) + // Carry only this partition's slice of the plan data on the Partition object, built when the + // task is serialized. Keeping the full O(numPartitions) map out of the broadcast task binary + // avoids the 2GB ByteArrayOutputStream limit on jobs with tens of millions of partitions. + // Mirrors CometExecRDD.getPartitions. + new CometNativeShuffleInputPartition(i, inputParts, new PlanDataSlice(perPartitionByKey, i)) }.toArray override def compute( @@ -144,8 +142,10 @@ private[shuffle] class CometNativeShuffleInputRDD( private[shuffle] class CometNativeShuffleInputPartition( override val index: Int, val inputPartitions: Array[Partition], - val planDataByKey: Map[String, Array[Byte]]) - extends Partition + planData: PlanDataSlice) + extends Partition { + def planDataByKey: Map[String, Array[Byte]] = planData.get +} /** * Iterator handed to [[CometNativeShuffleWriter.write]] via Spark's ShuffleMapTask. Reports no 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..0ff0dbe86f7 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 @@ -66,7 +66,7 @@ import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatibl import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} -import org.apache.comet.serde.operator.CometSink +import org.apache.comet.serde.operator.{CometSink, IcebergScanPools} import org.apache.comet.shims.MergeRowsMetricsShim /** @@ -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, PartitionPlanData]) = { plan match { case iceberg: CometIcebergNativeScanExec => // Trigger Spark's standard prepare -> waitForSubqueries lifecycle so DPP @@ -469,7 +469,7 @@ private[comet] object PlanDataInjector extends Logging { // it. sendDriverMetrics guards on the execution id, so the standalone-scan path (which // posts from doExecuteColumnar) and this fused path never double-post. iceberg.sendDriverMetrics() - if (iceberg.commonData.nonEmpty && iceberg.perPartitionData.nonEmpty) { + if (iceberg.commonData.nonEmpty && iceberg.partitionPlanData.length > 0) { // A self-join/self-merge can put two scans of the same table (same metadata_location) // in one native plan, and they can match on scan_hash_code too while reading // different files, so the key also carries the op's plan_id (see withPlanId). @@ -477,7 +477,7 @@ private[comet] object PlanDataInjector extends Logging { // uses to look it up, keeps the two sides from drifting apart. IcebergPlanDataInjector.getKey(iceberg.nativeOp) match { case Some(key) => - (Map(key -> iceberg.commonData), Map(key -> iceberg.perPartitionData)) + (Map(key -> iceberg.commonData), Map(key -> iceberg.partitionPlanData)) case None => (Map.empty, Map.empty) } @@ -495,8 +495,8 @@ private[comet] object PlanDataInjector extends Logging { case s: CometLeafExec with CometScanWithPlanData => s.ensureSubqueriesResolved() s.sendDriverMetrics() - if (s.commonData.nonEmpty && s.perPartitionData.nonEmpty) { - (Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.perPartitionData)) + if (s.commonData.nonEmpty && s.partitionPlanData.length > 0) { + (Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.partitionPlanData)) } else { (Map.empty, Map.empty) } @@ -518,13 +518,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 +570,13 @@ 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) + // The common block plus the pool entries this partition's tasks use (see IcebergScanPools). + val pruned = IcebergScanPools.forPartition(common, partition) val scanBuilder = scan.toBuilder - scanBuilder.setCommon(common) - scanBuilder.addAllFileScanTasks(tasksOnly.getFileScanTasksList) + scanBuilder.setCommon(pruned.getCommon) + scanBuilder.addAllFileScanTasks(pruned.getFileScanTasksList) op.toBuilder.setIcebergScan(scanBuilder).build() } @@ -887,7 +886,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, PartitionPlanData], shuffleScanIndices: Set[Int], hasScanInput: Boolean, // Like plan data, file paths are sliced onto each task's Partition on the driver. @@ -1129,11 +1128,11 @@ abstract class CometNativeExec extends CometExec { // broadcast plan. val (firstNonBroadcastPlanRDD, firstNonBroadcastPlanNumPartitions) = firstNonBroadcastPlan.get._1 match { - // Plan-data scans, Iceberg included, are sized from perPartitionData, which + // Plan-data scans, Iceberg included, are sized from partitionPlanData, which // findAllPlanData above has resolved. A bucketed scan has one entry per bucket, after any // bucket coalescing, so the count still matches its HashPartitioning. case scan: CometScanWithPlanData => - (null.asInstanceOf[RDD[Any]], scan.perPartitionData.length) + (null.asInstanceOf[RDD[Any]], scan.partitionPlanData.length) case scan: CometIcebergNativeScanExec => (null.asInstanceOf[RDD[Any]], scan.numPartitions) case plan: CometNativeExec => @@ -1407,6 +1406,9 @@ trait CometScanWithPlanData { self: CometLeafExec => def commonData: Array[Byte] def perPartitionData: Array[Array[Byte]] + /** [[perPartitionData]] for execution. Override to build each partition's bytes on demand. */ + def partitionPlanData: PartitionPlanData = PartitionPlanData(perPartitionData) + // DPP / partition filters that may carry AQE SubqueryAdaptiveBroadcast // subqueries needing rewrite by CometPlanAdaptiveDynamicPruningFilters. // Default empty: scans with dedicated handling (CometNativeScanExec, 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..348838c4f8d --- /dev/null +++ b/spark/src/test/resources/sql-tests/iceberg/merge_on_read_partition_layout.sql @@ -0,0 +1,146 @@ +-- 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 IcebergScanPools). 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. 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 + +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..087a946983b 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.serde.operator.IcebergScanPoolsTestUtils.storedPools 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 = storedPools(scans.head).exists(_.getResidualPoolCount > 0) 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 = storedPools(scans.head).map(_.getResidualPoolCount).sum + assert(residuals == 0, 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..86f53a7d2d5 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -49,6 +49,7 @@ import org.apache.comet.cloud.s3.MinioLocationScopedCredentialProvider import org.apache.comet.iceberg.{IcebergReflection, RESTCatalogHelper} import org.apache.comet.rules.CometScanRule import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde} +import org.apache.comet.serde.operator.IcebergScanPoolsTestUtils.{partitionScans, storedPools} import org.apache.comet.testing.{FuzzDataGenerator, SchemaGenOptions} /** @@ -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,123 @@ 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 = partitionScans(scan).flatMap { injected => + val pools = injected.getCommon + 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 + .buildAllPartitionDataForTesting() + .toSeq + .map(countByteOccurrences(_, 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)") + occurrences.forall(_ <= 1) && + countByteOccurrences(scan.commonData, pathPool.head.getBytes(UTF_8)) == 0, + s"Puffin path ${pathPool.head} serialized more than once in a partition: $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.buildAllPartitionDataForTesting().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 = + scan + .buildAllPartitionDataForTesting() + .toSeq + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getOwnPools) + // 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") + + 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 +967,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 +980,131 @@ 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.buildAllPartitionDataForTesting().toSeq + 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] = + scan + .buildAllPartitionDataForTesting() + .toSeq + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getOwnPools.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 + .buildAllPartitionDataForTesting() + .length > grouped.buildAllPartitionDataForTesting().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 +1367,71 @@ class CometIcebergNativeSuite table.newRowDelta().addDeletes(deleteFile).commit() } + // An unpartitioned equality delete applies to every data file, so with each data file in its + // own Spark partition every partition ships it. The rows must still match Spark. + test("a global equality delete ships with every partition and reads the right rows") { + 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) + + val (_, cometPlan) = + checkSparkAnswerAndOperator("SELECT * FROM test_cat.db.global_eq_delete ORDER BY id") + val scan = collectIcebergNativeScans(cometPlan).head + val own = scan + .buildAllPartitionDataForTesting() + .toSeq + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getOwnPools) + // 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}") + assert( + own.forall( + _.getDeleteFilePoolList.asScala.exists(_.getContentType == "EQUALITY_DELETES")), + "expected every partition to ship the equality delete") + val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) + assert(common.getDeleteFilePoolCount == 0, "the common block holds the delete file") + + 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 +4008,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,7 +5152,8 @@ 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 + val plannedTasks = icebergScans.head + .buildAllPartitionDataForTesting() .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum assert( @@ -5018,27 +5283,32 @@ 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 = + scans.head.buildAllPartitionDataForTesting().toSeq.map { bytes => + OperatorOuterClass.IcebergScan + .parseFrom(bytes) + .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), + s"expected a partition with several tasks carrying a residual: " + + tasksWithResidualPerPartition) + 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,7 +5960,8 @@ class CometIcebergNativeSuite val icebergScans = collectIcebergNativeScans(cometPlan) assert(icebergScans.nonEmpty, "Expected CometIcebergNativeScanExec in plan") val scan = icebergScans.head - val plannedTasks = scan.perPartitionData + val plannedTasks = scan + .buildAllPartitionDataForTesting() .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum val numSplits = scan.metrics("num_splits").value @@ -5959,7 +6230,8 @@ class CometIcebergNativeSuite val unprunedScan = collectIcebergNativeScans(unprunedPlan).head assert(unprunedScan.numPartitions == 1, s"Expected one split:\n$unprunedPlan") assert( - unprunedScan.perPartitionData + unprunedScan + .buildAllPartitionDataForTesting() .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum == 1, s"Expected one file scan task before pruning:\n$unprunedPlan") @@ -5996,7 +6268,8 @@ class CometIcebergNativeSuite collectIcebergDPPSubqueries(plan).nonEmpty, s"Expected executable DPP on the native scan:\n$plan") assert( - scan.perPartitionData + scan + .buildAllPartitionDataForTesting() .forall(OperatorOuterClass.IcebergScan.parseFrom(_).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,7 +6307,8 @@ class CometIcebergNativeSuite val factB = "aqe_cat.db.single_split_join_b" val tables = Seq(factA, factB) def fileTaskCount(scan: CometIcebergNativeScanExec): Int = - scan.perPartitionData + scan + .buildAllPartitionDataForTesting() .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum @@ -6500,7 +6774,8 @@ 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 + val plannedTasks = scan + .buildAllPartitionDataForTesting() .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum val numSplits = scan.metrics("num_splits").value diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala new file mode 100644 index 00000000000..aa2c565eb3b --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala @@ -0,0 +1,397 @@ +/* + * 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.serde.operator + +import scala.jdk.CollectionConverters._ +import scala.util.Random + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteFile, IcebergFilePools, IcebergFileScanTask, IcebergPredicate, IcebergScan, IcebergScanCommon, PartitionData} +import org.apache.comet.serde.operator.IcebergScanPoolsTestUtils._ + +class IcebergScanPoolsSuite extends AnyFunSuite { + + // A task with every pool index replaced by the entry it points at, so tasks compare equal no + // matter how their pools are numbered. + private case class ResolvedTask( + task: IcebergFileScanTask, + deletes: Option[Seq[(IcebergDeleteFile, String)]], + partitionData: Option[PartitionData], + residual: Option[IcebergPredicate]) + + private def resolveDelete(pools: IcebergScanCommon, idx: Int): (IcebergDeleteFile, String) = { + val deleteFile = pools.getDeleteFilePool(idx) + ( + deleteFile.toBuilder.clearFilePathIdx().build(), + pools.getDeleteFilePathPool(deleteFile.getFilePathIdx)) + } + + private def resolveSet(pools: IcebergScanCommon, idx: Int): Seq[(IcebergDeleteFile, String)] = + pools + .getDeleteFilesPool(idx) + .getDeleteFileIndicesList + .asScala + .toSeq + .map(d => resolveDelete(pools, d)) + + private def resolve(pools: IcebergScanCommon, task: IcebergFileScanTask): ResolvedTask = + ResolvedTask( + task.toBuilder.clearDeleteFilesIdx().clearPartitionDataIdx().clearResidualIdx().build(), + if (task.hasDeleteFilesIdx) Some(resolveSet(pools, task.getDeleteFilesIdx)) else None, + if (task.hasPartitionDataIdx) Some(pools.getPartitionDataPool(task.getPartitionDataIdx)) + else None, + if (task.hasResidualIdx) Some(pools.getResidualPool(task.getResidualIdx)) else None) + + // Each partition's message as the driver ships it, and the scan the executor builds from it. + private def shipAndInject(common: IcebergScanCommon, partitions: IndexedSeq[IcebergScan]) + : (IcebergScanCommon, IndexedSeq[(IcebergScan, IcebergScan)]) = { + val block = IcebergScanPools.commonBlock(common) + ( + block, + partitions.map { p => + val payload = IcebergScanPools.partitionPayload(common, p) + payload -> IcebergScanPools.forPartition(block, payload) + }) + } + + /** Fails unless every task the executor runs resolves to the same entries as the original. */ + private def assertSameTasks( + common: IcebergScanCommon, + partitions: IndexedSeq[IcebergScan]): Unit = { + val (_, shipped) = shipAndInject(common, partitions) + partitions.indices.foreach { p => + val injected = shipped(p)._2 + val expected = partitions(p).getFileScanTasksList.asScala.map(resolve(common, _)) + val actual = injected.getFileScanTasksList.asScala.map(resolve(injected.getCommon, _)) + assert(actual == expected, s"partition $p resolves differently once shipped") + } + } + + // A pool entry, named by its pool and its index in the original pools. + private type Entry = (String, Int) + + // Which partitions use each entry: directly from a task, or through a set or delete file that + // holds it. + private def users( + common: IcebergScanCommon, + partitions: IndexedSeq[IcebergScan]): Map[Entry, Set[Int]] = { + val uses = for { + p <- partitions.indices + task <- partitions(p).getFileScanTasksList.asScala.toSeq + entry <- { + val set = if (task.hasDeleteFilesIdx) { + val deletes = common + .getDeleteFilesPool(task.getDeleteFilesIdx) + .getDeleteFileIndicesList + .asScala + .toSeq + .map(_.intValue) + Seq("set" -> task.getDeleteFilesIdx) ++ deletes.map("delete" -> _) ++ + deletes.map(d => "path" -> common.getDeleteFilePool(d).getFilePathIdx) + } else { + Seq.empty + } + val partitionData = + if (task.hasPartitionDataIdx) Seq("partition" -> task.getPartitionDataIdx) + else Seq.empty + val residual = + if (task.hasResidualIdx) Seq("residual" -> task.getResidualIdx) else Seq.empty + set ++ partitionData ++ residual + } + } yield entry -> p + uses.groupBy(_._1).map { case (entry, ps) => entry -> ps.map(_._2).toSet } + } + + // An entry's content, so entries can be matched across differently numbered pools. + private def content(pools: IcebergScanCommon, entry: Entry): Any = entry match { + case ("path", i) => pools.getDeleteFilePathPool(i) + case ("delete", i) => resolveDelete(pools, i) + case ("set", i) => resolveSet(pools, i) + case ("partition", i) => pools.getPartitionDataPool(i) + case ("residual", i) => pools.getResidualPool(i) + case other => throw new IllegalArgumentException(s"unknown pool entry $other") + } + + // Every entry of the five split pools, from position `from` on. + private def entries(pools: IcebergScanCommon, from: IcebergScanCommon): Seq[Entry] = { + def range(name: String, count: IcebergScanCommon => Int): Seq[Entry] = + (count(from) until count(pools)).map(name -> _) + range("path", _.getDeleteFilePathPoolCount) ++ + range("delete", _.getDeleteFilePoolCount) ++ + range("set", _.getDeleteFilesPoolCount) ++ + range("partition", _.getPartitionDataPoolCount) ++ + range("residual", _.getResidualPoolCount) + } + + /** + * A random scan. Entries are distinct by content, as the driver's interning makes them. Some + * delete files share a path, like deletion vectors in one Puffin file, and some entries are + * used by no task. + */ + private def randomScan(rnd: Random): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { + val paths = 1 + rnd.nextInt(30) + val deletes = 1 + rnd.nextInt(40) + val partitionValues = 1 + rnd.nextInt(10) + val residuals = 1 + rnd.nextInt(5) + val common = withOtherFields(IcebergScanCommon.newBuilder()) + (0 until paths).foreach(i => common.addDeleteFilePathPool(deleteFilePath(i))) + (0 until deletes).foreach { i => + common.addDeleteFilePool(positionDelete(rnd.nextInt(paths), recordCount = i + 1L)) + } + val sets = Seq + .fill(1 + rnd.nextInt(20))(Seq.fill(1 + rnd.nextInt(4))(rnd.nextInt(deletes)).distinct) + .distinct + sets.foreach { set => + val list = DeleteFileList.newBuilder() + set.foreach(d => list.addDeleteFileIndices(d)) + common.addDeleteFilesPool(list) + } + (0 until partitionValues).foreach(i => common.addPartitionDataPool(partitionValue(i))) + (0 until residuals).foreach(i => common.addResidualPool(notNull(s"c$i"))) + + def sometimes(percent: Int, bound: Int): Option[Int] = + if (rnd.nextInt(100) < percent) Some(rnd.nextInt(bound)) else None + var files = 0 + val partitions = (0 until 1 + rnd.nextInt(8)).map { _ => + val scan = IcebergScan.newBuilder() + (0 until rnd.nextInt(6)).foreach { _ => + files += 1 + scan.addFileScanTasks( + fileTask( + s"data-$files.parquet", + sometimes(70, sets.size), + sometimes(80, partitionValues), + sometimes(50, residuals))) + } + scan.build() + } + (common.build(), partitions) + } + + test("every task resolves to the same entries once shipped and injected") { + (0 until 200).foreach { seed => + val (common, partitions) = randomScan(new Random(seed)) + withClue(s"seed $seed: ") { + assertSameTasks(common, partitions) + } + } + } + + test("a partition ships exactly the entries its tasks use, each once") { + (0 until 200).foreach { seed => + val (common, partitions) = randomScan(new Random(seed)) + val (block, shipped) = shipAndInject(common, partitions) + val usedBy = users(common, partitions) + val empty = IcebergScanCommon.getDefaultInstance + withClue(s"seed $seed: ") { + assert(entries(block, empty).isEmpty, "the common block holds per-file entries") + partitions.indices.foreach { p => + val expected = usedBy.collect { case (e, ps) if ps.contains(p) => content(common, e) } + val own = asCommon(shipped(p)._1.getOwnPools) + val sent = entries(own, empty).map(content(own, _)) + assert(sent.size == expected.size, s"partition $p ships an entry twice or extra") + assert(sent.toSet == expected.toSet, s"partition $p ships the wrong entries") + val injected = shipped(p)._2.getCommon + val kept = entries(injected, empty).map(content(injected, _)) + assert(kept.toSet == expected.toSet && kept.size == expected.size) + } + } + } + } + + test("the common block and each partition's plan do not grow with the scan") { + // Partition 0's data is the same at every scan size, so its message and plan must be too. + // Shipping the whole scan's pools in the common block made both grow with the scan. + val sizes = Seq(2, 100, 10000).map { buckets => + val (common, partitions) = bucketedScan(buckets) + val (block, shipped) = shipAndInject(common, partitions) + ( + common.getSerializedSize, + block.getSerializedSize, + shipped.head._1.getSerializedSize, + shipped.head._2.getSerializedSize) + } + val unsplit = sizes.map(_._1) + assert(unsplit.last > 100 * unsplit.head, s"the test scan does not grow: $unsplit") + assert(sizes.map(_._2).distinct.size == 1, s"common block grows with the scan: $sizes") + assert(sizes.map(_._3).distinct.size == 1, s"partition message grows with the scan: $sizes") + assert(sizes.map(_._4).distinct.size == 1, s"injected plan grows with the scan: $sizes") + } + + test("a partition's message does not grow when partitions share more delete files") { + // Each delete file applies to two neighbouring partitions. Partition 0 must still get only + // the two it uses, however many other partitions there are. + def shippedSize(partitionCount: Int): (Int, Int) = { + val common = IcebergScanCommon.newBuilder().setMetadataLocation(metadataLocation) + (0 until partitionCount).foreach { d => + common.addDeleteFilePathPool(deleteFilePath(d)) + common.addDeleteFilePool(positionDelete(d, recordCount = 1)) + } + val partitions = (0 until partitionCount).map { p => + val set = common.getDeleteFilesPoolCount + common.addDeleteFilesPool( + DeleteFileList + .newBuilder() + .addDeleteFileIndices(p) + .addDeleteFileIndices((p + 1) % partitionCount)) + IcebergScan + .newBuilder() + .addFileScanTasks(fileTask(f"data-$p%08d.parquet", Some(set), None, None)) + .build() + } + val (block, shipped) = shipAndInject(common.build(), partitions) + (block.getSerializedSize, shipped.head._1.getSerializedSize) + } + val sizes = Seq(3, 100, 10000).map(shippedSize) + assert(sizes.distinct.size == 1, s"shipped sizes grow with shared delete files: $sizes") + } + + test("a delete file every partition uses ships with each of them") { + // An unpartitioned equality delete applies to every data file. + val global = IcebergDeleteFile + .newBuilder() + .setFilePathIdx(0) + .setContentType("EQUALITY_DELETES") + .setFileFormat("PARQUET") + .addEqualityIds(1) + .setRecordCount(5) + .build() + val common = IcebergScanCommon + .newBuilder() + .setMetadataLocation(metadataLocation) + .addDeleteFilePathPool("s3://bucket/table/data/global-eq-deletes.parquet") + .addDeleteFilePool(global) + .addDeleteFilesPool(DeleteFileList.newBuilder().addDeleteFileIndices(0)) + .build() + val partitions = (0 until 6).map { b => + IcebergScan + .newBuilder() + .addFileScanTasks(fileTask(s"a-$b.parquet", Some(0), None, None)) + .build() + } + val (_, shipped) = shipAndInject(common, partitions) + shipped.foreach { case (payload, _) => + assert(payload.getOwnPools.getDeleteFilePoolList.asScala.toSeq == Seq(global)) + assert(payload.getOwnPools.getDeleteFilePathPoolCount == 1) + } + assertSameTasks(common, partitions) + } + + test("a partition with no tasks ships an empty message") { + val (common, partitions) = bucketedScan(3) + val withEmpty = partitions :+ IcebergScan.getDefaultInstance + val payload = IcebergScanPools.partitionPayload(common, IcebergScan.getDefaultInstance) + assert(!payload.hasOwnPools) + assert(payload.getSerializedSize == 0) + assertSameTasks(common, withEmpty) + } + + test("an index past the end of its pool fails with the pool's name") { + val (common, partitions) = bucketedScan(2) + val bad = partitions.head.toBuilder + .setFileScanTasks(0, partitions.head.getFileScanTasks(0).toBuilder.setDeleteFilesIdx(99)) + .build() + val error = intercept[IllegalStateException] { + IcebergScanPools.partitionPayload(common, bad) + } + assert(error.getMessage.contains("delete_files_set index 99 is out of range")) + } + + // The common block with the five per-file pools emptied: everything shipping must leave alone. + private def otherFields(common: IcebergScanCommon): IcebergScanCommon = + common.toBuilder + .clearDeleteFilePathPool() + .clearDeleteFilePool() + .clearDeleteFilesPool() + .clearPartitionDataPool() + .clearResidualPool() + .build() + + test("shipping keeps every field outside the five pools") { + (0 until 50).foreach { seed => + val (common, partitions) = randomScan(new Random(seed)) + val (block, shipped) = shipAndInject(common, partitions) + withClue(s"seed $seed: ") { + assert(block == otherFields(common)) + shipped.foreach { case (_, injected) => + assert(otherFields(injected.getCommon) == otherFields(common)) + } + } + } + } + + test("every pool index field is one the pruning knows about") { + // A new index into a pool must be handled by IcebergScanPools, or tasks keep indices into + // pools that no longer match. This fails when one is added without updating it. + def indexFields(descriptor: com.google.protobuf.Descriptors.Descriptor): Set[String] = + descriptor.getFields.asScala + .map(_.getName) + .filter(n => n.endsWith("_idx") || n.endsWith("_indices")) + .toSet + assert( + indexFields(IcebergFileScanTask.getDescriptor) == Set( + "schema_idx", + "partition_spec_idx", + "name_mapping_idx", + "project_field_ids_idx", + "partition_data_idx", + "delete_files_idx", + "residual_idx")) + assert(indexFields(IcebergDeleteFile.getDescriptor) == Set("file_path_idx")) + assert(indexFields(DeleteFileList.getDescriptor) == Set("delete_file_indices")) + // A per-partition field added to IcebergScan must be carried by partitionPayload. + assert( + IcebergScan.getDescriptor.getFields.asScala.map(_.getName).toSet == + Set("common", "file_scan_tasks", "own_pools")) + // A partition's own pools must be exactly the common block's per-file pools. + assert( + IcebergFilePools.getDescriptor.getFields.asScala.map(_.getName).toSet == Set( + "partition_data_pool", + "delete_files_pool", + "residual_pool", + "delete_file_pool", + "delete_file_path_pool")) + assert( + IcebergScanCommon.getDescriptor.getFields.asScala + .map(_.getName) + .filter(_.endsWith("_pool")) + .toSet == + Set( + "schema_pool", + "partition_type_pool", + "partition_spec_pool", + "name_mapping_pool", + "project_field_ids_pool", + "partition_data_pool", + "delete_files_pool", + "residual_pool", + "delete_file_pool", + "delete_file_path_pool")) + } + + test("lazily built scan data compares by content without building partitions") { + val (common, partitions) = bucketedScan(3) + def data(c: IcebergScanCommon) = + new IcebergPartitionPlanData(c, partitions.map(_.toByteArray).toArray) + assert(data(common).sameAs(data(common))) + assert(!data(common).sameAs(data(common.toBuilder.setScanHashCode(7).build()))) + } +} diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala new file mode 100644 index 00000000000..cbe5bd98184 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala @@ -0,0 +1,169 @@ +/* + * 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.serde.operator + +import org.apache.spark.sql.comet.CometIcebergNativeScanExec + +import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteFile, IcebergFilePools, IcebergFileScanTask, IcebergLiteral, IcebergPredicate, IcebergPredicateOperator, IcebergScan, IcebergScanCommon, IcebergUnaryPredicate, PartitionData, PartitionValue, ProjectFieldIdList, SparkStructField} + +/** Builders and views for tests of how a native Iceberg scan splits its pools. */ +object IcebergScanPoolsTestUtils { + + val metadataLocation = "s3://bucket/table/metadata/v1.metadata.json" + + def deleteFilePath(i: Int): String = f"s3://bucket/table/data/part-$i%08d-deletes.parquet" + + def positionDelete(pathIdx: Int, recordCount: Long): IcebergDeleteFile = + IcebergDeleteFile + .newBuilder() + .setFilePathIdx(pathIdx) + .setContentType("POSITION_DELETES") + .setFileFormat("PARQUET") + .setRecordCount(recordCount) + .build() + + def partitionValue(value: Int): PartitionData = + PartitionData + .newBuilder() + .addValues( + PartitionValue + .newBuilder() + .setFieldId(1000) + .setLiteral(IcebergLiteral.newBuilder().setIntVal(value))) + .build() + + def notNull(column: String): IcebergPredicate = + IcebergPredicate + .newBuilder() + .setUnary( + IcebergUnaryPredicate + .newBuilder() + .setColumn(column) + .setOp(IcebergPredicateOperator.NotNull)) + .build() + + def fileTask( + dataFile: String, + deleteSet: Option[Int], + partitionData: Option[Int], + residual: Option[Int]): IcebergFileScanTask = { + val task = IcebergFileScanTask + .newBuilder() + .setDataFilePath(dataFile) + .setLength(1024) + .setFileSizeInBytes(1024) + deleteSet.foreach(task.setDeleteFilesIdx) + partitionData.foreach(task.setPartitionDataIdx) + residual.foreach(task.setResidualIdx) + task.build() + } + + /** Sets the common-block fields the split leaves alone, so tests can check they survive. */ + def withOtherFields(common: IcebergScanCommon.Builder): IcebergScanCommon.Builder = + common + .setMetadataLocation(metadataLocation) + .setScanHashCode(42) + .setCatalogName("test_cat") + .setDataFileConcurrencyLimit(4) + .putCatalogProperties("s3.region", "us-west-2") + .addRequiredSchema(SparkStructField.newBuilder().setName("id").setNullable(true)) + .addSchemaPool("{\"type\":\"struct\",\"schema-id\":0,\"fields\":[]}") + .addPartitionSpecPool("{\"spec-id\":0,\"fields\":[]}") + .addPartitionTypePool("{\"type\":\"struct\",\"fields\":[]}") + .addNameMappingPool("[]") + .addProjectFieldIdsPool(ProjectFieldIdList.newBuilder().addFieldIds(1)) + + /** + * A scan like a bucketed merge-on-read table read one bucket per partition. Partition b holds + * four data files, three delete files only it uses, and its own partition value. Every task + * carries the same residual. + */ + def bucketedScan(buckets: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { + val common = IcebergScanCommon + .newBuilder() + .setMetadataLocation(metadataLocation) + .addResidualPool(notNull("payload")) + val partitions = (0 until buckets).map { b => + val set = DeleteFileList.newBuilder() + (0 until 3).foreach { d => + val idx = common.getDeleteFilePoolCount + common.addDeleteFilePathPool(deleteFilePath(b * 3 + d)) + common.addDeleteFilePool(positionDelete(idx, recordCount = 10)) + set.addDeleteFileIndices(idx) + } + val setIdx = common.getDeleteFilesPoolCount + common.addDeleteFilesPool(set) + val partitionIdx = common.getPartitionDataPoolCount + common.addPartitionDataPool(partitionValue(b)) + val scan = IcebergScan.newBuilder() + (0 until 4).foreach { f => + scan.addFileScanTasks( + fileTask(f"data-$b%08d-$f.parquet", Some(setIdx), Some(partitionIdx), Some(0))) + } + scan.build() + } + (common.build(), partitions) + } + + /** A scan whose partitions each read one data file and share one unpartitioned delete file. */ + def globalDeleteScan(partitionCount: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { + val common = IcebergScanCommon + .newBuilder() + .setMetadataLocation(metadataLocation) + .addDeleteFilePathPool(deleteFilePath(0)) + .addDeleteFilePool(positionDelete(0, recordCount = 100)) + .addDeleteFilesPool(DeleteFileList.newBuilder().addDeleteFileIndices(0)) + .build() + val partitions = (0 until partitionCount).map { p => + IcebergScan + .newBuilder() + .addFileScanTasks(fileTask(f"data-$p%08d.parquet", Some(0), None, None)) + .build() + } + (common, partitions) + } + + /** A partition's own pools in a common block, so tests can read both the same way. */ + def asCommon(pools: IcebergFilePools): IcebergScanCommon = + IcebergScanCommon + .newBuilder() + .addAllDeleteFilePathPool(pools.getDeleteFilePathPoolList) + .addAllDeleteFilePool(pools.getDeleteFilePoolList) + .addAllDeleteFilesPool(pools.getDeleteFilesPoolList) + .addAllPartitionDataPool(pools.getPartitionDataPoolList) + .addAllResidualPool(pools.getResidualPoolList) + .build() + + /** 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) +: + scan + .buildAllPartitionDataForTesting() + .toSeq + .map(bytes => asCommon(IcebergScan.parseFrom(bytes).getOwnPools)) + + /** The scan each partition's native plan runs, as the injector builds it. */ + def partitionScans(scan: CometIcebergNativeScanExec): Seq[IcebergScan] = { + val common = IcebergScanCommon.parseFrom(scan.commonData) + scan.buildAllPartitionDataForTesting().toSeq.map { bytes => + IcebergScanPools.forPartition(common, IcebergScan.parseFrom(bytes)) + } + } +} 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..af78973665a 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").toArray.map(_.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").toArray.map(_.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..4361a737036 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 @@ -26,6 +28,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager import org.apache.comet.serde.OperatorOuterClass import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.{IcebergScanPools, IcebergScanPoolsTestUtils} class PlanDataInjectorSuite extends AnyFunSuite { @@ -165,6 +168,85 @@ class PlanDataInjectorSuite extends AnyFunSuite { assert(IcebergPlanDataInjector.getKey(opA) != IcebergPlanDataInjector.getKey(opB)) } + test("an injected Iceberg plan does not grow with the rest of the scan") { + // Partition 0's data is the same at every scan size, while every other partition adds its + // own delete files. Tasks allocate the injected plan in serializeOperator, so its size must + // not follow the whole scan. + val scanOp = icebergScanOp(IcebergScanPoolsTestUtils.metadataLocation, scanHashCode = 7) + val key = IcebergPlanDataInjector.getKey(scanOp).get + val injected = Seq(2, 100, 10000).map { buckets => + val (common, partitions) = IcebergScanPoolsTestUtils.bucketedScan(buckets) + PlanDataInjector.injectPlanData( + parseBasePlan(scanOp.toByteArray), + Map(key -> IcebergScanPools.commonBlock(common).toByteArray), + Map(key -> IcebergScanPools.partitionPayload(common, partitions.head).toByteArray)) + } + val sizes = injected.map(PlanDataInjector.serializeOperator(_).length) + assert(sizes.distinct.size == 1, s"injected plan grows with the scan: $sizes") + + // The partition's own delete files are reachable from its tasks after injection. + val scan = injected.last.getIcebergScan + val pools = scan.getCommon + val deletes = pools.getDeleteFilesPool(scan.getFileScanTasks(0).getDeleteFilesIdx) + val paths = deletes.getDeleteFileIndicesList.asScala.map { d => + pools.getDeleteFilePathPool(pools.getDeleteFilePool(d).getFilePathIdx) + } + assert(paths == (0 until 3).map(IcebergScanPoolsTestUtils.deleteFilePath)) + } + + test("a coalesced task carries a copy of a shared delete file per input partition") { + // Pins a documented limit: without a shuffle, a coalesced task serializes all its input + // partitions, and each brings its own copy of the delete files it uses. Spark's Iceberg + // reader writes a shared delete file once per task. Update this test if that changes. + val (common, partitions) = IcebergScanPoolsTestUtils.globalDeleteScan(8) + val data = new PartitionPlanData { + override def length: Int = partitions.length + override def apply(i: Int): Array[Byte] = + IcebergScanPools.partitionPayload(common, partitions(i)).toByteArray + } + def taskSize(inputs: Int): Int = { + val parents = (0 until inputs).map { i => + new CometExecPartition(i, Array.empty, new PlanDataSlice(Map("k" -> data), i)) + }.toArray + val bytes = new java.io.ByteArrayOutputStream() + val out = new java.io.ObjectOutputStream(bytes) + out.writeObject(parents) + out.close() + bytes.size() + } + val perPartition = taskSize(2) - taskSize(1) + assert(taskSize(8) - taskSize(1) == 7 * perPartition, "each input partition adds the same") + assert( + perPartition > data(0).length, + "each input partition carries its own copy of the shared delete file") + } + + test("a task's plan data is built only when its partition is serialized") { + // The driver holds every partition object for the job, so a partition must not keep its + // bytes. They are built while Spark serializes the task, and the executor gets them intact. + val built = new java.util.concurrent.atomic.AtomicInteger() + val data = new PartitionPlanData { + override def length: Int = 3 + override def apply(i: Int): Array[Byte] = { + built.incrementAndGet() + Array.fill(4)(i.toByte) + } + } + val partition = new CometExecPartition(1, Array.empty, new PlanDataSlice(Map("k" -> data), 1)) + assert(built.get == 0, "plan data built before the task was serialized") + + val bytes = new java.io.ByteArrayOutputStream() + val out = new java.io.ObjectOutputStream(bytes) + out.writeObject(partition) + out.close() + assert(built.get == 1) + + val in = new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(bytes.toByteArray)) + val received = in.readObject().asInstanceOf[CometExecPartition] + assert(received.planDataByKey("k").toSeq == Seq.fill(4)(1.toByte)) + assert(built.get == 1, "the executor side rebuilt the plan data") + } + /** * Builds an un-injected NativeScan operator the way the driver ships it: hasCommon, no * file_partition, source_key_hash embedded (see CometNativeScanExec.apply). 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..c0ca8290715 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, CometExecRDD, CometMetricNode, NativeExecContext, PartitionPlanData, PlanDataSlice} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -119,7 +119,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 2, Set.empty, CometMetricNode(Map.empty), - planData) + planData.map { case (key, bytes) => key -> PartitionPlanData(bytes) }) val local = remote.copyForLocalShuffle() assert(local.id != remote.id) @@ -207,7 +207,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" -> PartitionPlanData(Array.fill(numPartitions)(new Array[Byte](1024)))) val childMetricNode = CometMetricNode(Map.empty) val writerMetrics = Map( "spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "disk spilled bytes"), @@ -263,4 +263,55 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { assert(part.planDataByKey.keySet == Set("scan-0")) assert(part.planDataByKey("scan-0").length == 1024) } + + test("a shuffle map task's plan data is built only when its partition is serialized") { + val built = new java.util.concurrent.atomic.AtomicInteger() + val data = new PartitionPlanData { + override def length: Int = 2 + override def apply(i: Int): Array[Byte] = { + built.incrementAndGet() + Array.fill(3)(i.toByte) + } + } + val partition = + new CometNativeShuffleInputPartition(1, Array.empty, new PlanDataSlice(Map("k" -> data), 1)) + assert(built.get == 0, "plan data built before the task was serialized") + + val bytes = new java.io.ByteArrayOutputStream() + val out = new java.io.ObjectOutputStream(bytes) + out.writeObject(partition) + out.close() + assert(built.get == 1) + + val in = new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(bytes.toByteArray)) + val received = in.readObject().asInstanceOf[CometNativeShuffleInputPartition] + assert(received.planDataByKey("k").toSeq == Seq.fill(3)(1.toByte)) + assert(built.get == 1, "the executor side rebuilt the plan data") + } + + test("an RDD's partitions don't build any partition's plan data") { + // Spark keeps every partition object for the job's lifetime, so building the partitions must + // not build the plan data. Only serializing a task does. + val built = new java.util.concurrent.atomic.AtomicInteger() + val data = new PartitionPlanData { + override def length: Int = 1000 + override def apply(i: Int): Array[Byte] = { + built.incrementAndGet() + Array(i.toByte) + } + } + val rdd = CometExecRDD( + spark.sparkContext, + inputRDDs = Seq.empty, + commonByKey = Map("k" -> Array[Byte](1)), + perPartitionByKey = Map("k" -> data), + serializedPlan = Array.empty, + planFingerprint = 0L, + numPartitions = 1000, + numOutputCols = 1, + nativeMetrics = CometMetricNode(Map.empty), + subqueries = Seq.empty) + assert(rdd.partitions.length == 1000) + assert(built.get == 0, s"building the partitions built ${built.get} partitions' plan data") + } } From 20e0df2d1b76296e3097f3fe6bdc3be9126a33f6 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 4 Nov 2026 14:06:42 -0800 Subject: [PATCH 2/8] fix: keep Iceberg pool entries every partition uses in the common block, and hold driver pools as bytes Entries that every non-empty partition uses, such as an unpartitioned equality delete, now ship once in the common block instead of in every task. Every task needs them anyway, so executors pay nothing extra, and tasks no longer grow past spark.rpc.message.maxSize when many shared delete files exist. Partition messages index the common entries first, then their own. After planning, the driver keeps the full pools as the serialized bytes plus an offset per entry, and parses only the entries a task uses when that task is serialized. The parsed message took about 3.9 times its serialized size; bytes plus the index take about 1.1 times. Adds coalesce queries to the merge-on-read SQL test. --- docs/source/user-guide/latest/iceberg.md | 13 +- .../operator/CometIcebergNativeScan.scala | 24 +- .../serde/operator/IcebergScanPools.scala | 463 +++++++++++++----- .../merge_on_read_partition_layout.sql | 26 +- .../comet/CometIcebergNativeSuite.scala | 30 +- .../operator/IcebergScanPoolsSuite.scala | 86 +++- .../operator/IcebergScanPoolsTestUtils.scala | 10 +- .../sql/comet/PlanDataInjectorSuite.scala | 25 +- 8 files changed, 494 insertions(+), 183 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 8e241fe0a3f..b016f254a22 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -197,12 +197,13 @@ These `s3.*` storage properties are not specific to the Hive catalog shown here. ### 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 coalesce without a shuffle -(`coalesce(n)`, or a write that coalesces its input) combines several input partitions into one -task, and that task then carries a copy of a shared delete file for every input partition that -uses it. Scans where many partitions share delete files, such as tables with unpartitioned -equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning -instead of coalescing, avoids this. +that its own files need, the way Spark's Iceberg reader does. Entries that every partition needs, +such as unpartitioned equality deletes, are sent once per executor instead. A coalesce without a +shuffle (`coalesce(n)`, or a write that coalesces its input) combines several input partitions +into one task, and that task then carries a copy of a shared delete file for every input +partition that uses it. Scans where many, but not all, partitions share delete files can make +coalesced tasks large. Compacting delete files, or repartitioning instead of coalescing, avoids +this. ### Current limitations 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 614cd1d9cdb..5c0035de783 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 @@ -1272,7 +1272,10 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // The per-file pool entries ship with each partition that uses them, so the common block // every task deserializes doesn't grow with the number of files. val original = commonBuilder.build() - val common = IcebergScanPools.commonBlock(original) + // Entries every partition uses stay in the common block once; the rest ship per partition. + val layout = + IcebergScanPools.Layout(new IcebergScanPools.MessagePools(original), perPartitionBuilders) + val common = IcebergScanPools.commonBlock(original, layout) // Per-pool byte sizes of the whole scan, to diagnose large plan data. Sizes sum as Long because // a single pool at or past protobuf's 2 GiB message limit overflows the int getSerializedSize, @@ -1337,7 +1340,9 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit s"${taskBytes.length}, task bytes total=${taskBytes.iterator.map(_.length.toLong).sum}, " + s"max=${if (taskBytes.isEmpty) 0 else taskBytes.iterator.map(_.length).max}") - (commonBytes, new IcebergPartitionPlanData(original, taskBytes)) + // Keep the full pools as bytes, not as the parsed message, which takes several times the heap. + val pools = new IcebergScanPools.SerializedPools(original.toByteArray) + (commonBytes, new IcebergPartitionPlanData(pools, layout, taskBytes)) } override def createExec(nativeOp: Operator, op: CometBatchScanExec): CometNativeExec = { @@ -1364,29 +1369,30 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * plus the per-file pool entries they use (see IcebergScanPools.partitionPayload). */ private[operator] final class IcebergPartitionPlanData( - private val original: OperatorOuterClass.IcebergScanCommon, + private val pools: IcebergScanPools.SerializedPools, + layout: IcebergScanPools.Layout, private val taskBytes: Array[Array[Byte]]) extends PartitionPlanData { override def length: Int = taskBytes.length // Same table state read into the same partitions, compared without building any partition. + // The layout follows from the pools and the tasks. override def sameAs(other: PartitionPlanData): Boolean = other match { case that: IcebergPartitionPlanData => - that.original == original && that.taskBytes.length == taskBytes.length && + java.util.Arrays.equals(that.pools.bytes, pools.bytes) && + that.taskBytes.length == taskBytes.length && taskBytes.indices.forall(i => java.util.Arrays.equals(taskBytes(i), that.taskBytes(i))) case _ => false } - // With no per-file entries in the scan, the stored task bytes already are each message. - private val hasFileEntries = IcebergScanPools.hasFileEntries(original) - override def apply(i: Int): Array[Byte] = - if (!hasFileEntries) { + if (pools.isEmpty) { + // With no per-file entries in the scan, the stored task bytes already are each message. taskBytes(i) } else { IcebergScanPools - .partitionPayload(original, OperatorOuterClass.IcebergScan.parseFrom(taskBytes(i))) + .partitionPayload(pools, layout, OperatorOuterClass.IcebergScan.parseFrom(taskBytes(i))) .toByteArray } } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala index 5dd73f95d1b..73bf6f8d421 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala @@ -19,155 +19,358 @@ package org.apache.comet.serde.operator +import java.nio.charset.StandardCharsets.UTF_8 + import scala.collection.mutable import scala.jdk.CollectionConverters._ -import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergFilePools, IcebergFileScanTask, IcebergScan, IcebergScanCommon} +import com.google.protobuf.CodedInputStream + +import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteFile, IcebergFilePools, IcebergPredicate, IcebergScan, IcebergScanCommon, PartitionData} /** * Decides where an Iceberg scan's pool entries ship. * * Five pools grow with the number of files in the scan: delete file paths, delete files, delete - * file sets, partition values and residuals. Their entries ship with each Spark partition whose - * tasks use them, as Iceberg's own Spark reader ships each task's delete files with the task. So - * a task never holds entries that only other partitions use, and the common block every task - * deserializes holds only the small pools (schemas, specs, name mappings, field ids) and the - * scan's settings. + * file sets, partition values and residuals. An entry that every partition uses, such as an + * unpartitioned equality delete, ships once in the common block, since every task needs it + * anyway. Every other entry ships with each partition whose tasks use it, as Iceberg's own Spark + * reader ships each task's delete files with the task. So a task never holds entries that only + * other partitions use. * - * On the driver, [[partitionPayload]] builds one partition's message from the full pools. On the - * executor, [[forPartition]] merges it with the common block into the scan the native plan runs. + * A partition's tasks index the common entries first, then the partition's own. On the driver, + * [[partitionPayload]] builds one partition's message. On the executor, [[forPartition]] appends + * its own entries to the common block. * * Known limit: a coalesce without a shuffle puts several partitions' messages in one task, and - * each carries its own copy of the entries they share. Delete files that many partitions use, - * such as unpartitioned equality deletes, are then repeated once per coalesced partition. + * each carries its own copy of the entries they share. Entries that many but not all partitions + * use, such as the partition-scoped deletes of a large partition, are then repeated once per + * coalesced partition. */ private[apache] object IcebergScanPools { - /** The common block every task gets: the scan with the five per-file pools emptied. */ - def commonBlock(original: IcebergScanCommon): IcebergScanCommon = - original.toBuilder - .clearDeleteFilePathPool() - .clearDeleteFilePool() - .clearDeleteFilesPool() - .clearPartitionDataPool() - .clearResidualPool() - .build() + /** Read access to a scan's five per-file pools. */ + trait FilePools { + def pathCount: Int + def path(i: Int): String + def deleteCount: Int + def delete(i: Int): IcebergDeleteFile + def setCount: Int + def set(i: Int): DeleteFileList + def partitionValueCount: Int + def partitionValue(i: Int): PartitionData + def residualCount: Int + def residual(i: Int): IcebergPredicate - /** - * One partition's message, built on the driver. `tasks` indexes into `original`. The result - * carries the tasks, renumbered, and in `own_pools` only the per-file entries they use. - */ - def partitionPayload(original: IcebergScanCommon, tasks: IcebergScan): IcebergScan = { - val (pools, renumbered) = prune(original, tasks) - val result = tasks.toBuilder.clearFileScanTasks().addAllFileScanTasks(renumbered.asJava) - if (pools.getDeleteFilePathPoolCount > 0 || pools.getDeleteFilePoolCount > 0 || - pools.getDeleteFilesPoolCount > 0 || pools.getPartitionDataPoolCount > 0 || - pools.getResidualPoolCount > 0) { - result.setOwnPools(pools) - } - result.build() + def isEmpty: Boolean = + pathCount == 0 && deleteCount == 0 && setCount == 0 && partitionValueCount == 0 && + residualCount == 0 } - /** Whether `common` still holds per-file pool entries, which must ship with partitions. */ - def hasFileEntries(common: IcebergScanCommon): Boolean = - common.getDeleteFilePathPoolCount > 0 || common.getDeleteFilePoolCount > 0 || - common.getDeleteFilesPoolCount > 0 || common.getPartitionDataPoolCount > 0 || - common.getResidualPoolCount > 0 + /** The pools of a built message. */ + final class MessagePools(common: IcebergScanCommon) extends FilePools { + override def pathCount: Int = common.getDeleteFilePathPoolCount + override def path(i: Int): String = common.getDeleteFilePathPool(i) + override def deleteCount: Int = common.getDeleteFilePoolCount + override def delete(i: Int): IcebergDeleteFile = common.getDeleteFilePool(i) + override def setCount: Int = common.getDeleteFilesPoolCount + override def set(i: Int): DeleteFileList = common.getDeleteFilesPool(i) + override def partitionValueCount: Int = common.getPartitionDataPoolCount + override def partitionValue(i: Int): PartitionData = common.getPartitionDataPool(i) + override def residualCount: Int = common.getResidualPoolCount + override def residual(i: Int): IcebergPredicate = common.getResidualPool(i) + } /** - * The scan one partition's native plan runs, built on the executor from the common block and - * the partition's message. The message already holds just the per-file entries its tasks use, - * numbered from 0, so this only adds them to the common block. + * The pools of a serialized `IcebergScanCommon`, parsing only the entries asked for. The bytes + * plus one offset per entry take a fraction of the heap of the parsed message, which is several + * times its serialized size. */ - def forPartition(common: IcebergScanCommon, partition: IcebergScan): IcebergScan = { - if (hasFileEntries(common)) { - throw new IllegalStateException( - "Iceberg common block holds per-file pool entries; they must ship with each partition") + final class SerializedPools(val bytes: Array[Byte]) extends FilePools { + // Each entry's offset (high 32 bits) and length (low 32 bits), per pool, in wire order. + private val located: Array[Array[Long]] = { + val entries = Array.fill(5)(mutable.ArrayBuilder.make[Long]) + val in = CodedInputStream.newInstance(bytes) + var tag = in.readTag() + while (tag != 0) { + val pool = SerializedPools.poolOf(tag >>> 3) + if (pool >= 0) { + val length = in.readRawVarint32() + entries(pool) += (in.getTotalBytesRead.toLong << 32) | length.toLong + in.skipRawBytes(length) + } else { + val _ = in.skipField(tag) + } + tag = in.readTag() + } + entries.map(_.result()) } - val own = partition.getOwnPools - val merged = common.toBuilder - .addAllDeleteFilePathPool(own.getDeleteFilePathPoolList) - .addAllDeleteFilePool(own.getDeleteFilePoolList) - .addAllDeleteFilesPool(own.getDeleteFilesPoolList) - .addAllPartitionDataPool(own.getPartitionDataPoolList) - .addAllResidualPool(own.getResidualPoolList) - partition.toBuilder.clearOwnPools().setCommon(merged).build() + + private def offset(pool: Int, i: Int): Int = (located(pool)(i) >>> 32).toInt + private def length(pool: Int, i: Int): Int = (located(pool)(i) & 0xffffffffL).toInt + + override def pathCount: Int = located(SerializedPools.Paths).length + override def path(i: Int): String = + new String(bytes, offset(SerializedPools.Paths, i), length(SerializedPools.Paths, i), UTF_8) + override def deleteCount: Int = located(SerializedPools.Deletes).length + override def delete(i: Int): IcebergDeleteFile = + IcebergDeleteFile + .parser() + .parseFrom(bytes, offset(SerializedPools.Deletes, i), length(SerializedPools.Deletes, i)) + override def setCount: Int = located(SerializedPools.Sets).length + override def set(i: Int): DeleteFileList = + DeleteFileList + .parser() + .parseFrom(bytes, offset(SerializedPools.Sets, i), length(SerializedPools.Sets, i)) + override def partitionValueCount: Int = located(SerializedPools.PartitionValues).length + override def partitionValue(i: Int): PartitionData = + PartitionData + .parser() + .parseFrom( + bytes, + offset(SerializedPools.PartitionValues, i), + length(SerializedPools.PartitionValues, i)) + override def residualCount: Int = located(SerializedPools.Residuals).length + override def residual(i: Int): IcebergPredicate = + IcebergPredicate + .parser() + .parseFrom( + bytes, + offset(SerializedPools.Residuals, i), + length(SerializedPools.Residuals, i)) + } + + private object SerializedPools { + val Paths = 0 + val Deletes = 1 + val Sets = 2 + val PartitionValues = 3 + val Residuals = 4 + + def poolOf(field: Int): Int = + if (field == IcebergScanCommon.DELETE_FILE_PATH_POOL_FIELD_NUMBER) Paths + else if (field == IcebergScanCommon.DELETE_FILE_POOL_FIELD_NUMBER) Deletes + else if (field == IcebergScanCommon.DELETE_FILES_POOL_FIELD_NUMBER) Sets + else if (field == IcebergScanCommon.PARTITION_DATA_POOL_FIELD_NUMBER) PartitionValues + else if (field == IcebergScanCommon.RESIDUAL_POOL_FIELD_NUMBER) Residuals + else -1 } - /** Entry i of one of the scan's pools, failing with the pool's name when i is out of range. */ - private def entry[T](name: String, i: Int, size: Int)(get: Int => T): T = { + private val Unused = -1 + + private def checkIndex(name: String, i: Int, size: Int): Unit = if (i < 0 || i >= size) { throw new IllegalStateException( s"Iceberg $name index $i is out of range: the pool has $size entries") } - get(i) - } /** - * The per-file pool entries of `common` that `tasks` reach, and the tasks renumbered to index - * them. The work follows the tasks, not the size of the pools. + * Which entries every non-empty partition uses, and each one's index in the common block. With + * fewer than two non-empty partitions there are none. */ - private def prune( - common: IcebergScanCommon, - tasks: IcebergScan): (IcebergFilePools.Builder, Seq[IcebergFileScanTask]) = { - val pools = IcebergFilePools.newBuilder() - val pathIdx = mutable.HashMap.empty[Int, Int] - val deleteIdx = mutable.HashMap.empty[Int, Int] - val setIdx = mutable.HashMap.empty[Int, Int] - val partitionDataIdx = mutable.HashMap.empty[Int, Int] - val residualIdx = mutable.HashMap.empty[Int, Int] - - def path(i: Int): Int = - pathIdx.getOrElseUpdate( - i, { - pools.addDeleteFilePathPool( - entry("delete_file_path", i, common.getDeleteFilePathPoolCount)( - common.getDeleteFilePathPool)) - pools.getDeleteFilePathPoolCount - 1 - }) - - def delete(i: Int): Int = - deleteIdx.getOrElseUpdate( - i, { - val deleteFile = - entry("delete_file", i, common.getDeleteFilePoolCount)(common.getDeleteFilePool) - pools.addDeleteFilePool( - deleteFile.toBuilder.setFilePathIdx(path(deleteFile.getFilePathIdx))) - pools.getDeleteFilePoolCount - 1 - }) - - def set(i: Int): Int = - setIdx.getOrElseUpdate( - i, { - val set = - entry("delete_files_set", i, common.getDeleteFilesPoolCount)( - common.getDeleteFilesPool) - val remapped = DeleteFileList.newBuilder() - (0 until set.getDeleteFileIndicesCount).foreach { j => - remapped.addDeleteFileIndices(delete(set.getDeleteFileIndices(j))) + final class Layout private[IcebergScanPools] ( + val path: Array[Int], + val delete: Array[Int], + val set: Array[Int], + val partitionValue: Array[Int], + val residual: Array[Int]) { + + val pathCount: Int = path.count(_ != Unused) + val deleteCount: Int = delete.count(_ != Unused) + val setCount: Int = set.count(_ != Unused) + val partitionValueCount: Int = partitionValue.count(_ != Unused) + val residualCount: Int = residual.count(_ != Unused) + } + + object Layout { + + /** Finds the entries of `pools` that every non-empty partition among `partitions` uses. */ + def apply(pools: FilePools, partitions: collection.IndexedSeq[IcebergScan]): Layout = { + val paths = new Users("delete_file_path", pools.pathCount) + val deletes = new Users("delete_file", pools.deleteCount) + val sets = new Users("delete_files_set", pools.setCount) + val partitionValues = new Users("partition_data", pools.partitionValueCount) + val residuals = new Users("residual", pools.residualCount) + var nonEmpty = 0 + partitions.indices.foreach { p => + val tasks = partitions(p).getFileScanTasksList + if (!tasks.isEmpty) nonEmpty += 1 + tasks.asScala.foreach { task => + if (task.hasResidualIdx) residuals.use(task.getResidualIdx, p) + if (task.hasPartitionDataIdx) partitionValues.use(task.getPartitionDataIdx, p) + if (task.hasDeleteFilesIdx && sets.use(task.getDeleteFilesIdx, p)) { + val set = pools.set(task.getDeleteFilesIdx) + (0 until set.getDeleteFileIndicesCount).foreach { j => + val d = set.getDeleteFileIndices(j) + if (deletes.use(d, p)) paths.use(pools.delete(d).getFilePathIdx, p) + } } - pools.addDeleteFilesPool(remapped) - pools.getDeleteFilesPoolCount - 1 - }) - - def partitionValue(i: Int): Int = - partitionDataIdx.getOrElseUpdate( - i, { - pools.addPartitionDataPool( - entry("partition_data", i, common.getPartitionDataPoolCount)( - common.getPartitionDataPool)) - pools.getPartitionDataPoolCount - 1 - }) - - def residual(i: Int): Int = - residualIdx.getOrElseUpdate( - i, { - pools.addResidualPool( - entry("residual", i, common.getResidualPoolCount)(common.getResidualPool)) - pools.getResidualPoolCount - 1 - }) - - val renumbered = tasks.getFileScanTasksList.asScala.toSeq.map { task => + } + } + val everyone = if (nonEmpty >= 2) nonEmpty else Int.MaxValue + new Layout( + paths.usedBy(everyone), + deletes.usedBy(everyone), + sets.usedBy(everyone), + partitionValues.usedBy(everyone), + residuals.usedBy(everyone)) + } + + /** How many partitions use each entry of one pool. */ + private final class Users(name: String, size: Int) { + private val lastUser = Array.fill(size)(Unused) + private val users = new Array[Int](size) + + // Records that partition p uses entry i. Returns true the first time p does, so the caller + // records the entries it points to as well. + def use(i: Int, p: Int): Boolean = { + checkIndex(name, i, size) + if (lastUser(i) == p) { + false + } else { + lastUser(i) = p + users(i) += 1 + true + } + } + + // Each entry's index among the entries `partitions` partitions use, or Unused. + def usedBy(partitions: Int): Array[Int] = { + var next = 0 + users.map { n => + if (n == partitions) { + next += 1 + next - 1 + } else { + Unused + } + } + } + } + } + + /** + * The common block every task gets: the scan without its per-file pools, plus the entries every + * partition uses. + */ + def commonBlock(original: IcebergScanCommon, layout: Layout): IcebergScanCommon = { + val builder = original.toBuilder + .clearDeleteFilePathPool() + .clearDeleteFilePool() + .clearDeleteFilesPool() + .clearPartitionDataPool() + .clearResidualPool() + def shared(positions: Array[Int]): Iterator[Int] = + positions.indices.iterator.filter(positions(_) != Unused) + shared(layout.path).foreach(i => + builder.addDeleteFilePathPool(original.getDeleteFilePathPool(i))) + // Every partition that uses a delete file uses its path, and every one that uses a set uses + // its delete files, so those are in the common block too. + shared(layout.delete).foreach { i => + val deleteFile = original.getDeleteFilePool(i) + builder.addDeleteFilePool( + deleteFile.toBuilder.setFilePathIdx(layout.path(deleteFile.getFilePathIdx))) + } + shared(layout.set).foreach { i => + val set = original.getDeleteFilesPool(i) + val remapped = DeleteFileList.newBuilder() + (0 until set.getDeleteFileIndicesCount).foreach { j => + remapped.addDeleteFileIndices(layout.delete(set.getDeleteFileIndices(j))) + } + builder.addDeleteFilesPool(remapped) + } + shared(layout.partitionValue).foreach(i => + builder.addPartitionDataPool(original.getPartitionDataPool(i))) + shared(layout.residual).foreach(i => builder.addResidualPool(original.getResidualPool(i))) + builder.build() + } + + /** + * One partition's message, built on the driver. `tasks` indexes into `pools`. The result + * carries the tasks, renumbered to index the common block's entries first, and in `own_pools` + * the other entries they use. + */ + def partitionPayload(pools: FilePools, layout: Layout, tasks: IcebergScan): IcebergScan = { + val own = IcebergFilePools.newBuilder() + val ownPath = mutable.HashMap.empty[Int, Int] + val ownDelete = mutable.HashMap.empty[Int, Int] + val ownSet = mutable.HashMap.empty[Int, Int] + val ownPartitionValue = mutable.HashMap.empty[Int, Int] + val ownResidual = mutable.HashMap.empty[Int, Int] + + def path(i: Int): Int = { + checkIndex("delete_file_path", i, pools.pathCount) + if (layout.path(i) != Unused) { + layout.path(i) + } else { + ownPath.getOrElseUpdate( + i, { + own.addDeleteFilePathPool(pools.path(i)) + layout.pathCount + own.getDeleteFilePathPoolCount - 1 + }) + } + } + + def delete(i: Int): Int = { + checkIndex("delete_file", i, pools.deleteCount) + if (layout.delete(i) != Unused) { + layout.delete(i) + } else { + ownDelete.getOrElseUpdate( + i, { + val deleteFile = pools.delete(i) + val pathIdx = path(deleteFile.getFilePathIdx) + own.addDeleteFilePool(deleteFile.toBuilder.setFilePathIdx(pathIdx)) + layout.deleteCount + own.getDeleteFilePoolCount - 1 + }) + } + } + + def set(i: Int): Int = { + checkIndex("delete_files_set", i, pools.setCount) + if (layout.set(i) != Unused) { + layout.set(i) + } else { + ownSet.getOrElseUpdate( + i, { + val set = pools.set(i) + val remapped = DeleteFileList.newBuilder() + (0 until set.getDeleteFileIndicesCount).foreach { j => + remapped.addDeleteFileIndices(delete(set.getDeleteFileIndices(j))) + } + own.addDeleteFilesPool(remapped) + layout.setCount + own.getDeleteFilesPoolCount - 1 + }) + } + } + + def partitionValue(i: Int): Int = { + checkIndex("partition_data", i, pools.partitionValueCount) + if (layout.partitionValue(i) != Unused) { + layout.partitionValue(i) + } else { + ownPartitionValue.getOrElseUpdate( + i, { + own.addPartitionDataPool(pools.partitionValue(i)) + layout.partitionValueCount + own.getPartitionDataPoolCount - 1 + }) + } + } + + def residual(i: Int): Int = { + checkIndex("residual", i, pools.residualCount) + if (layout.residual(i) != Unused) { + layout.residual(i) + } else { + ownResidual.getOrElseUpdate( + i, { + own.addResidualPool(pools.residual(i)) + layout.residualCount + own.getResidualPoolCount - 1 + }) + } + } + + val renumbered = tasks.getFileScanTasksList.asScala.map { task => val builder = task.toBuilder if (task.hasResidualIdx) builder.setResidualIdx(residual(task.getResidualIdx)) if (task.hasPartitionDataIdx) { @@ -176,6 +379,28 @@ private[apache] object IcebergScanPools { if (task.hasDeleteFilesIdx) builder.setDeleteFilesIdx(set(task.getDeleteFilesIdx)) builder.build() } - (pools, renumbered) + val result = tasks.toBuilder.clearFileScanTasks().addAllFileScanTasks(renumbered.asJava) + if (own.getDeleteFilePathPoolCount > 0 || own.getDeleteFilePoolCount > 0 || + own.getDeleteFilesPoolCount > 0 || own.getPartitionDataPoolCount > 0 || + own.getResidualPoolCount > 0) { + result.setOwnPools(own) + } + result.build() + } + + /** + * The scan one partition's native plan runs, built on the executor from the common block and + * the partition's message. Its tasks already index the common entries first, then its own, so + * this only appends its own entries to the common block. + */ + def forPartition(common: IcebergScanCommon, partition: IcebergScan): IcebergScan = { + val own = partition.getOwnPools + val merged = common.toBuilder + .addAllDeleteFilePathPool(own.getDeleteFilePathPoolList) + .addAllDeleteFilePool(own.getDeleteFilePoolList) + .addAllDeleteFilesPool(own.getDeleteFilesPoolList) + .addAllPartitionDataPool(own.getPartitionDataPoolList) + .addAllResidualPool(own.getResidualPoolList) + partition.toBuilder.clearOwnPools().setCommon(merged).build() } } 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 index 348838c4f8d..341a4340823 100644 --- 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 @@ -15,14 +15,15 @@ -- 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 IcebergScanPools). 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. 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. +-- The native scan ships a partition value, delete file or residual once in the common block when +-- every Spark partition uses it, and otherwise with each Spark partition whose tasks use it (see +-- IcebergScanPools). 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. @@ -142,5 +143,14 @@ 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/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 86f53a7d2d5..5e636afb738 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -818,10 +818,13 @@ class CometIcebergNativeSuite .buildAllPartitionDataForTesting() .toSeq .map(countByteOccurrences(_, pathPool.head.getBytes(UTF_8))) + // The path ships once in the common block when every partition uses it, otherwise once + // in each partition that does. + val inCommon = countByteOccurrences(scan.commonData, pathPool.head.getBytes(UTF_8)) assert( - occurrences.forall(_ <= 1) && - countByteOccurrences(scan.commonData, pathPool.head.getBytes(UTF_8)) == 0, - s"Puffin path ${pathPool.head} serialized more than once in a partition: $occurrences") + inCommon <= 1 && occurrences.forall(_ <= 1) && (inCommon == 0 || occurrences.sum == 0), + s"Puffin path ${pathPool.head} serialized more than once: common $inCommon, " + + s"partitions $occurrences") spark.sql("DROP TABLE test_cat.db.dv_shared_puffin") } @@ -1367,9 +1370,9 @@ class CometIcebergNativeSuite table.newRowDelta().addDeletes(deleteFile).commit() } - // An unpartitioned equality delete applies to every data file, so with each data file in its - // own Spark partition every partition ships it. The rows must still match Spark. - test("a global equality delete ships with every partition and reads the right rows") { + // An unpartitioned equality delete applies to every data file, so every partition uses it and it + // ships once in the common block instead of in each task. The rows must still match Spark. + test("a global equality delete ships once in the common block and reads the right rows") { assume(icebergAvailable, "Iceberg not available in classpath") assume( icebergVersionAtLeast(1, 8), @@ -1422,10 +1425,12 @@ class CometIcebergNativeSuite s"expected one partition per data file: ${own.length}") assert( own.forall( - _.getDeleteFilePoolList.asScala.exists(_.getContentType == "EQUALITY_DELETES")), - "expected every partition to ship the equality delete") + !_.getDeleteFilePoolList.asScala.exists(_.getContentType == "EQUALITY_DELETES")), + "a partition ships the equality delete itself") val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) - assert(common.getDeleteFilePoolCount == 0, "the common block holds the delete file") + assert( + common.getDeleteFilePoolList.asScala.count(_.getContentType == "EQUALITY_DELETES") == 1, + s"expected the equality delete once in the common block: $common") spark.sql("DROP TABLE test_cat.db.global_eq_delete") } @@ -5304,10 +5309,11 @@ class CometIcebergNativeSuite tasksWithResidualPerPartition.exists(_ > 1), s"expected a partition with several tasks carrying a residual: " + tasksWithResidualPerPartition) - val expectedResiduals = 0 +: tasksWithResidualPerPartition.map(n => if (n > 0) 1 else 0) + // Every task uses it, so it ships once: in the common block when several partitions have + // tasks, otherwise with the one partition that does. assert( - residualsPerPartition == expectedResiduals, - "expected the shared non-partition residual to dedup to 1 pool entry per partition " + + residualsPerPartition.sum == 1, + "expected the shared non-partition residual to dedup to 1 pool entry " + s"across $tasksWithResidual task references, got $residualsPerPartition") spark.sql("DROP TABLE test_cat.db.residual_dedup_test") diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala index aa2c565eb3b..15710e82d3a 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala @@ -63,11 +63,14 @@ class IcebergScanPoolsSuite extends AnyFunSuite { // Each partition's message as the driver ships it, and the scan the executor builds from it. private def shipAndInject(common: IcebergScanCommon, partitions: IndexedSeq[IcebergScan]) : (IcebergScanCommon, IndexedSeq[(IcebergScan, IcebergScan)]) = { - val block = IcebergScanPools.commonBlock(common) + // Read the pools the way the driver does after planning: from the serialized scan. + val pools = new IcebergScanPools.SerializedPools(common.toByteArray) + val layout = IcebergScanPools.Layout(pools, partitions) + val block = IcebergScanPools.commonBlock(common, layout) ( block, partitions.map { p => - val payload = IcebergScanPools.partitionPayload(common, p) + val payload = IcebergScanPools.partitionPayload(pools, layout, p) payload -> IcebergScanPools.forPartition(block, payload) }) } @@ -201,17 +204,26 @@ class IcebergScanPoolsSuite extends AnyFunSuite { val (block, shipped) = shipAndInject(common, partitions) val usedBy = users(common, partitions) val empty = IcebergScanCommon.getDefaultInstance + // Entries every non-empty partition uses ship once in the common block. + val nonEmpty = partitions.indices.filter(partitions(_).getFileScanTasksCount > 0).toSet + // toSeq first: Map.collect into (pool, index) tuples would build a Map and merge them. + val everywhere = usedBy.toSeq.collect { + case (e, ps) if nonEmpty.size >= 2 && ps == nonEmpty => e + }.toSet withClue(s"seed $seed: ") { - assert(entries(block, empty).isEmpty, "the common block holds per-file entries") - partitions.indices.foreach { p => - val expected = usedBy.collect { case (e, ps) if ps.contains(p) => content(common, e) } - val own = asCommon(shipped(p)._1.getOwnPools) - val sent = entries(own, empty).map(content(own, _)) - assert(sent.size == expected.size, s"partition $p ships an entry twice or extra") - assert(sent.toSet == expected.toSet, s"partition $p ships the wrong entries") + val inCommon = entries(block, empty).map(content(block, _)) + assert(inCommon.size == everywhere.size, "the common block holds an entry twice or extra") + assert(inCommon.toSet == everywhere.map(content(common, _))) + nonEmpty.foreach { p => + val used = usedBy.toSeq.collect { case (e, ps) if ps.contains(p) => e }.toSet + val expected = (used -- everywhere).map(content(common, _)) + // Own entries follow the common ones and may point at them, so resolve them merged. val injected = shipped(p)._2.getCommon + val sent = entries(injected, block).map(content(injected, _)) + assert(sent.size == expected.size, s"partition $p ships an entry twice or extra") + assert(sent.toSet == expected, s"partition $p ships the wrong entries") val kept = entries(injected, empty).map(content(injected, _)) - assert(kept.toSet == expected.toSet && kept.size == expected.size) + assert(kept.size == used.size && kept.toSet == used.map(content(common, _))) } } } @@ -264,7 +276,7 @@ class IcebergScanPoolsSuite extends AnyFunSuite { assert(sizes.distinct.size == 1, s"shipped sizes grow with shared delete files: $sizes") } - test("a delete file every partition uses ships with each of them") { + test("a delete file every partition uses ships once in the common block") { // An unpartitioned equality delete applies to every data file. val global = IcebergDeleteFile .newBuilder() @@ -287,18 +299,21 @@ class IcebergScanPoolsSuite extends AnyFunSuite { .addFileScanTasks(fileTask(s"a-$b.parquet", Some(0), None, None)) .build() } - val (_, shipped) = shipAndInject(common, partitions) - shipped.foreach { case (payload, _) => - assert(payload.getOwnPools.getDeleteFilePoolList.asScala.toSeq == Seq(global)) - assert(payload.getOwnPools.getDeleteFilePathPoolCount == 1) - } + val (block, shipped) = shipAndInject(common, partitions) + assert(block.getDeleteFilePoolList.asScala.toSeq == Seq(global)) + assert(block.getDeleteFilePathPoolCount == 1 && block.getDeleteFilesPoolCount == 1) + shipped.foreach { case (payload, _) => assert(!payload.hasOwnPools) } assertSameTasks(common, partitions) } test("a partition with no tasks ships an empty message") { val (common, partitions) = bucketedScan(3) val withEmpty = partitions :+ IcebergScan.getDefaultInstance - val payload = IcebergScanPools.partitionPayload(common, IcebergScan.getDefaultInstance) + val pools = new IcebergScanPools.MessagePools(common) + val payload = IcebergScanPools.partitionPayload( + pools, + IcebergScanPools.Layout(pools, withEmpty), + IcebergScan.getDefaultInstance) assert(!payload.hasOwnPools) assert(payload.getSerializedSize == 0) assertSameTasks(common, withEmpty) @@ -309,8 +324,10 @@ class IcebergScanPoolsSuite extends AnyFunSuite { val bad = partitions.head.toBuilder .setFileScanTasks(0, partitions.head.getFileScanTasks(0).toBuilder.setDeleteFilesIdx(99)) .build() + val pools = new IcebergScanPools.MessagePools(common) + val layout = IcebergScanPools.Layout(pools, partitions) val error = intercept[IllegalStateException] { - IcebergScanPools.partitionPayload(common, bad) + IcebergScanPools.partitionPayload(pools, layout, bad) } assert(error.getMessage.contains("delete_files_set index 99 is out of range")) } @@ -330,7 +347,7 @@ class IcebergScanPoolsSuite extends AnyFunSuite { val (common, partitions) = randomScan(new Random(seed)) val (block, shipped) = shipAndInject(common, partitions) withClue(s"seed $seed: ") { - assert(block == otherFields(common)) + assert(otherFields(block) == otherFields(common)) shipped.foreach { case (_, injected) => assert(otherFields(injected.getCommon) == otherFields(common)) } @@ -389,9 +406,36 @@ class IcebergScanPoolsSuite extends AnyFunSuite { test("lazily built scan data compares by content without building partitions") { val (common, partitions) = bucketedScan(3) - def data(c: IcebergScanCommon) = - new IcebergPartitionPlanData(c, partitions.map(_.toByteArray).toArray) + def data(c: IcebergScanCommon) = { + val pools = new IcebergScanPools.SerializedPools(c.toByteArray) + new IcebergPartitionPlanData( + pools, + IcebergScanPools.Layout(pools, partitions), + partitions.map(_.toByteArray).toArray) + } assert(data(common).sameAs(data(common))) assert(!data(common).sameAs(data(common.toBuilder.setScanHashCode(7).build()))) } + + test("reading the serialized pools gives the same entries as the parsed message") { + (0 until 50).foreach { seed => + val (common, _) = randomScan(new Random(seed)) + val parsed = new IcebergScanPools.MessagePools(common) + val serialized = new IcebergScanPools.SerializedPools(common.toByteArray) + withClue(s"seed $seed: ") { + assert(serialized.pathCount == parsed.pathCount) + assert((0 until parsed.pathCount).forall(i => serialized.path(i) == parsed.path(i))) + assert(serialized.deleteCount == parsed.deleteCount) + assert((0 until parsed.deleteCount).forall(i => serialized.delete(i) == parsed.delete(i))) + assert(serialized.setCount == parsed.setCount) + assert((0 until parsed.setCount).forall(i => serialized.set(i) == parsed.set(i))) + assert(serialized.partitionValueCount == parsed.partitionValueCount) + assert((0 until parsed.partitionValueCount).forall(i => + serialized.partitionValue(i) == parsed.partitionValue(i))) + assert(serialized.residualCount == parsed.residualCount) + assert((0 until parsed.residualCount).forall(i => + serialized.residual(i) == parsed.residual(i))) + } + } + } } diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala index cbe5bd98184..dbf1ce68283 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala @@ -122,8 +122,11 @@ object IcebergScanPoolsTestUtils { (common.build(), partitions) } - /** A scan whose partitions each read one data file and share one unpartitioned delete file. */ - def globalDeleteScan(partitionCount: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { + /** + * A scan where every partition but the last reads one data file with the same delete file. The + * delete file is shared, but not by every partition, so each partition that uses it ships it. + */ + def sharedDeleteScan(partitionCount: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { val common = IcebergScanCommon .newBuilder() .setMetadataLocation(metadataLocation) @@ -132,9 +135,10 @@ object IcebergScanPoolsTestUtils { .addDeleteFilesPool(DeleteFileList.newBuilder().addDeleteFileIndices(0)) .build() val partitions = (0 until partitionCount).map { p => + val deletes = if (p < partitionCount - 1) Some(0) else None IcebergScan .newBuilder() - .addFileScanTasks(fileTask(f"data-$p%08d.parquet", Some(0), None, None)) + .addFileScanTasks(fileTask(f"data-$p%08d.parquet", deletes, None, None)) .build() } (common, partitions) 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 4361a737036..99338c35f70 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 @@ -168,6 +168,11 @@ class PlanDataInjectorSuite extends AnyFunSuite { assert(IcebergPlanDataInjector.getKey(opA) != IcebergPlanDataInjector.getKey(opB)) } + private def layoutOf( + common: OperatorOuterClass.IcebergScanCommon, + partitions: IndexedSeq[OperatorOuterClass.IcebergScan]): IcebergScanPools.Layout = + IcebergScanPools.Layout(new IcebergScanPools.MessagePools(common), partitions) + test("an injected Iceberg plan does not grow with the rest of the scan") { // Partition 0's data is the same at every scan size, while every other partition adds its // own delete files. Tasks allocate the injected plan in serializeOperator, so its size must @@ -178,8 +183,15 @@ class PlanDataInjectorSuite extends AnyFunSuite { val (common, partitions) = IcebergScanPoolsTestUtils.bucketedScan(buckets) PlanDataInjector.injectPlanData( parseBasePlan(scanOp.toByteArray), - Map(key -> IcebergScanPools.commonBlock(common).toByteArray), - Map(key -> IcebergScanPools.partitionPayload(common, partitions.head).toByteArray)) + Map( + key -> IcebergScanPools.commonBlock(common, layoutOf(common, partitions)).toByteArray), + Map( + key -> IcebergScanPools + .partitionPayload( + new IcebergScanPools.MessagePools(common), + layoutOf(common, partitions), + partitions.head) + .toByteArray)) } val sizes = injected.map(PlanDataInjector.serializeOperator(_).length) assert(sizes.distinct.size == 1, s"injected plan grows with the scan: $sizes") @@ -198,11 +210,14 @@ class PlanDataInjectorSuite extends AnyFunSuite { // Pins a documented limit: without a shuffle, a coalesced task serializes all its input // partitions, and each brings its own copy of the delete files it uses. Spark's Iceberg // reader writes a shared delete file once per task. Update this test if that changes. - val (common, partitions) = IcebergScanPoolsTestUtils.globalDeleteScan(8) + val (common, partitions) = IcebergScanPoolsTestUtils.sharedDeleteScan(8) + val layout = layoutOf(common, partitions) val data = new PartitionPlanData { override def length: Int = partitions.length override def apply(i: Int): Array[Byte] = - IcebergScanPools.partitionPayload(common, partitions(i)).toByteArray + IcebergScanPools + .partitionPayload(new IcebergScanPools.MessagePools(common), layout, partitions(i)) + .toByteArray } def taskSize(inputs: Int): Int = { val parents = (0 until inputs).map { i => @@ -215,7 +230,7 @@ class PlanDataInjectorSuite extends AnyFunSuite { bytes.size() } val perPartition = taskSize(2) - taskSize(1) - assert(taskSize(8) - taskSize(1) == 7 * perPartition, "each input partition adds the same") + assert(taskSize(7) - taskSize(1) == 6 * perPartition, "each input partition adds the same") assert( perPartition > data(0).length, "each input partition carries its own copy of the shared delete file") From cb7d76fe6f163d00ceb424ecfcba919e3ffd49e3 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 4 Nov 2026 14:24:40 -0800 Subject: [PATCH 3/8] fix: keep Iceberg pool entries that at least half of the partitions use in the common block An equality delete applies only to older data files, so newer appends, or a bucket with no deletes, leave some partitions without it. Requiring every partition to use an entry sent such deletes back into every task. Entries used by at least half of the non-empty partitions now ship once in the common block. Every entry there is needed by at least half of the tasks, so the common block is at most twice what an average partition's tasks need. --- docs/source/user-guide/latest/iceberg.md | 12 ++--- .../operator/CometIcebergNativeScan.scala | 3 +- .../serde/operator/IcebergScanPools.scala | 44 ++++++++++--------- .../merge_on_read_partition_layout.sql | 6 +-- .../comet/CometIcebergNativeSuite.scala | 8 +++- .../operator/IcebergScanPoolsSuite.scala | 34 +++++++++++--- .../operator/IcebergScanPoolsTestUtils.scala | 10 +++-- .../sql/comet/PlanDataInjectorSuite.scala | 8 ++-- 8 files changed, 80 insertions(+), 45 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index b016f254a22..b582777b4c9 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -197,12 +197,12 @@ These `s3.*` storage properties are not specific to the Hive catalog shown here. ### 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. Entries that every partition needs, -such as unpartitioned equality deletes, are sent once per executor instead. A coalesce without a -shuffle (`coalesce(n)`, or a write that coalesces its input) combines several input partitions -into one task, and that task then carries a copy of a shared delete file for every input -partition that uses it. Scans where many, but not all, partitions share delete files can make -coalesced tasks large. Compacting delete files, or repartitioning instead of coalescing, avoids +that its own files need, the way Spark's Iceberg reader does. Entries that at least half of the +partitions need, such as unpartitioned equality deletes, are sent once per executor instead. A +coalesce without a shuffle (`coalesce(n)`, or a write that coalesces its input) combines several +input partitions into one task, and that task then carries a copy of a shared delete file for +every input partition that uses it. Scans where delete files are shared by many partitions, but +fewer than half, can make coalesced tasks large. Compacting delete files, or repartitioning instead of coalescing, avoids this. ### Current limitations 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 5c0035de783..9608f3b5124 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 @@ -1272,7 +1272,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // The per-file pool entries ship with each partition that uses them, so the common block // every task deserializes doesn't grow with the number of files. val original = commonBuilder.build() - // Entries every partition uses stay in the common block once; the rest ship per partition. + // Entries at least half of the partitions use stay in the common block; the rest ship with + // each partition that uses them. val layout = IcebergScanPools.Layout(new IcebergScanPools.MessagePools(original), perPartitionBuilders) val common = IcebergScanPools.commonBlock(original, layout) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala index 73bf6f8d421..5110c74cab6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala @@ -32,20 +32,20 @@ import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteF * Decides where an Iceberg scan's pool entries ship. * * Five pools grow with the number of files in the scan: delete file paths, delete files, delete - * file sets, partition values and residuals. An entry that every partition uses, such as an - * unpartitioned equality delete, ships once in the common block, since every task needs it - * anyway. Every other entry ships with each partition whose tasks use it, as Iceberg's own Spark - * reader ships each task's delete files with the task. So a task never holds entries that only - * other partitions use. + * file sets, partition values and residuals. An entry that at least half of the partitions use, + * such as an unpartitioned equality delete, ships once in the common block. Every other entry + * ships with each partition whose tasks use it, as Iceberg's own Spark reader ships each task's + * delete files with the task. Since each entry in the common block is needed by at least half of + * the tasks, the common block is at most twice what an average partition's tasks need. * * A partition's tasks index the common entries first, then the partition's own. On the driver, * [[partitionPayload]] builds one partition's message. On the executor, [[forPartition]] appends * its own entries to the common block. * * Known limit: a coalesce without a shuffle puts several partitions' messages in one task, and - * each carries its own copy of the entries they share. Entries that many but not all partitions - * use, such as the partition-scoped deletes of a large partition, are then repeated once per - * coalesced partition. + * each carries its own copy of the entries they share. Entries that fewer than half of the + * partitions use, such as the partition-scoped deletes of a large partition, are then repeated + * once per coalesced partition. */ private[apache] object IcebergScanPools { @@ -165,8 +165,8 @@ private[apache] object IcebergScanPools { } /** - * Which entries every non-empty partition uses, and each one's index in the common block. With - * fewer than two non-empty partitions there are none. + * Which entries at least half of the non-empty partitions use, and each one's index in the + * common block. With fewer than two non-empty partitions there are none. */ final class Layout private[IcebergScanPools] ( val path: Array[Int], @@ -184,7 +184,7 @@ private[apache] object IcebergScanPools { object Layout { - /** Finds the entries of `pools` that every non-empty partition among `partitions` uses. */ + /** Finds the entries of `pools` that at least half of the non-empty `partitions` use. */ def apply(pools: FilePools, partitions: collection.IndexedSeq[IcebergScan]): Layout = { val paths = new Users("delete_file_path", pools.pathCount) val deletes = new Users("delete_file", pools.deleteCount) @@ -207,13 +207,15 @@ private[apache] object IcebergScanPools { } } } - val everyone = if (nonEmpty >= 2) nonEmpty else Int.MaxValue + // A set's delete files and a delete file's path have at least the set's users, so an entry + // in the common block only points at entries that are there too. + val half = if (nonEmpty >= 2) (nonEmpty + 1) / 2 else Int.MaxValue new Layout( - paths.usedBy(everyone), - deletes.usedBy(everyone), - sets.usedBy(everyone), - partitionValues.usedBy(everyone), - residuals.usedBy(everyone)) + paths.usedBy(half), + deletes.usedBy(half), + sets.usedBy(half), + partitionValues.usedBy(half), + residuals.usedBy(half)) } /** How many partitions use each entry of one pool. */ @@ -234,11 +236,11 @@ private[apache] object IcebergScanPools { } } - // Each entry's index among the entries `partitions` partitions use, or Unused. + // Each entry's index among the entries at least `partitions` partitions use, or Unused. def usedBy(partitions: Int): Array[Int] = { var next = 0 users.map { n => - if (n == partitions) { + if (n >= partitions) { next += 1 next - 1 } else { @@ -250,8 +252,8 @@ private[apache] object IcebergScanPools { } /** - * The common block every task gets: the scan without its per-file pools, plus the entries every - * partition uses. + * The common block every task gets: the scan without its per-file pools, plus the entries at + * least half of the partitions use. */ def commonBlock(original: IcebergScanCommon, layout: Layout): IcebergScanCommon = { val builder = original.toBuilder 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 index 341a4340823..46cd0615d68 100644 --- 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 @@ -16,9 +16,9 @@ -- under the License. -- The native scan ships a partition value, delete file or residual once in the common block when --- every Spark partition uses it, and otherwise with each Spark partition whose tasks use it (see --- IcebergScanPools). 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 +-- at least half of the Spark partitions use it, and otherwise with each Spark partition whose +-- tasks use it (see IcebergScanPools). 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 diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 5e636afb738..1b72492f2cc 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -1370,8 +1370,9 @@ class CometIcebergNativeSuite table.newRowDelta().addDeletes(deleteFile).commit() } - // An unpartitioned equality delete applies to every data file, so every partition uses it and it - // ships once in the common block instead of in each task. The rows must still match Spark. + // An unpartitioned equality delete applies to every older data file, so most partitions use it + // and it ships once in the common block instead of in each task. The rows must still match + // Spark, including a newer row the delete does not apply to. test("a global equality delete ships once in the common block and reads the right rows") { assume(icebergAvailable, "Iceberg not available in classpath") assume( @@ -1406,6 +1407,9 @@ class CometIcebergNativeSuite "id", Seq(2, 105, 210, 315), warehouseDir) + // A newer append the delete does not apply to: its id 2 row must survive. Partitions holding + // only this file do not use the delete, which still ships once while most partitions do. + 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") diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala index 15710e82d3a..2e2cf57682b 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala @@ -204,11 +204,12 @@ class IcebergScanPoolsSuite extends AnyFunSuite { val (block, shipped) = shipAndInject(common, partitions) val usedBy = users(common, partitions) val empty = IcebergScanCommon.getDefaultInstance - // Entries every non-empty partition uses ship once in the common block. + // Entries at least half of the non-empty partitions use ship once in the common block. val nonEmpty = partitions.indices.filter(partitions(_).getFileScanTasksCount > 0).toSet + val half = (nonEmpty.size + 1) / 2 // toSeq first: Map.collect into (pool, index) tuples would build a Map and merge them. val everywhere = usedBy.toSeq.collect { - case (e, ps) if nonEmpty.size >= 2 && ps == nonEmpty => e + case (e, ps) if nonEmpty.size >= 2 && ps.size >= half => e }.toSet withClue(s"seed $seed: ") { val inCommon = entries(block, empty).map(content(block, _)) @@ -222,8 +223,10 @@ class IcebergScanPoolsSuite extends AnyFunSuite { val sent = entries(injected, block).map(content(injected, _)) assert(sent.size == expected.size, s"partition $p ships an entry twice or extra") assert(sent.toSet == expected, s"partition $p ships the wrong entries") + // The plan holds the common entries, which some partitions skip, and its own. val kept = entries(injected, empty).map(content(injected, _)) - assert(kept.size == used.size && kept.toSet == used.map(content(common, _))) + val visible = used ++ everywhere + assert(kept.size == visible.size && kept.toSet == visible.map(content(common, _))) } } } @@ -232,7 +235,8 @@ class IcebergScanPoolsSuite extends AnyFunSuite { test("the common block and each partition's plan do not grow with the scan") { // Partition 0's data is the same at every scan size, so its message and plan must be too. // Shipping the whole scan's pools in the common block made both grow with the scan. - val sizes = Seq(2, 100, 10000).map { buckets => + // From four partitions up, a bucket's own entries are used by fewer than half of them. + val sizes = Seq(4, 100, 10000).map { buckets => val (common, partitions) = bucketedScan(buckets) val (block, shipped) = shipAndInject(common, partitions) ( @@ -272,7 +276,8 @@ class IcebergScanPoolsSuite extends AnyFunSuite { val (block, shipped) = shipAndInject(common.build(), partitions) (block.getSerializedSize, shipped.head._1.getSerializedSize) } - val sizes = Seq(3, 100, 10000).map(shippedSize) + // From five partitions up, a delete file two neighbours share is used by fewer than half. + val sizes = Seq(5, 100, 10000).map(shippedSize) assert(sizes.distinct.size == 1, s"shipped sizes grow with shared delete files: $sizes") } @@ -438,4 +443,23 @@ class IcebergScanPoolsSuite extends AnyFunSuite { } } } + + test("a delete file most partitions use ships once, one fewer than half use ships per task") { + // Newer appends that a delete does not apply to, or a bucket with no deletes, leave some + // partitions without it. It still ships once while at least half of the partitions use it. + def layout(users: Int): (IcebergScanCommon, IndexedSeq[(IcebergScan, IcebergScan)]) = { + val (common, partitions) = sharedDeleteScan(8, users) + assertSameTasks(common, partitions) + shipAndInject(common, partitions) + } + val (most, mostShipped) = layout(users = 7) + assert(most.getDeleteFilePoolCount == 1) + assert(mostShipped.forall(!_._1.hasOwnPools)) + val (half, _) = layout(users = 4) + assert(half.getDeleteFilePoolCount == 1) + val (few, fewShipped) = layout(users = 3) + assert(few.getDeleteFilePoolCount == 0) + assert(fewShipped.take(3).forall(_._1.getOwnPools.getDeleteFilePoolCount == 1)) + assert(fewShipped.drop(3).forall(!_._1.hasOwnPools)) + } } diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala index dbf1ce68283..308cfaa0157 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala @@ -123,10 +123,12 @@ object IcebergScanPoolsTestUtils { } /** - * A scan where every partition but the last reads one data file with the same delete file. The - * delete file is shared, but not by every partition, so each partition that uses it ships it. + * A scan where each partition reads one data file, and the first `users` of them share one + * delete file. Whether it ships in the common block depends on how many partitions use it. */ - def sharedDeleteScan(partitionCount: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { + def sharedDeleteScan( + partitionCount: Int, + users: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { val common = IcebergScanCommon .newBuilder() .setMetadataLocation(metadataLocation) @@ -135,7 +137,7 @@ object IcebergScanPoolsTestUtils { .addDeleteFilesPool(DeleteFileList.newBuilder().addDeleteFileIndices(0)) .build() val partitions = (0 until partitionCount).map { p => - val deletes = if (p < partitionCount - 1) Some(0) else None + val deletes = if (p < users) Some(0) else None IcebergScan .newBuilder() .addFileScanTasks(fileTask(f"data-$p%08d.parquet", deletes, None, None)) 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 99338c35f70..4ab3ee594c3 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 @@ -179,7 +179,8 @@ class PlanDataInjectorSuite extends AnyFunSuite { // not follow the whole scan. val scanOp = icebergScanOp(IcebergScanPoolsTestUtils.metadataLocation, scanHashCode = 7) val key = IcebergPlanDataInjector.getKey(scanOp).get - val injected = Seq(2, 100, 10000).map { buckets => + // From four partitions up, a bucket's own entries are used by fewer than half of them. + val injected = Seq(4, 100, 10000).map { buckets => val (common, partitions) = IcebergScanPoolsTestUtils.bucketedScan(buckets) PlanDataInjector.injectPlanData( parseBasePlan(scanOp.toByteArray), @@ -210,7 +211,8 @@ class PlanDataInjectorSuite extends AnyFunSuite { // Pins a documented limit: without a shuffle, a coalesced task serializes all its input // partitions, and each brings its own copy of the delete files it uses. Spark's Iceberg // reader writes a shared delete file once per task. Update this test if that changes. - val (common, partitions) = IcebergScanPoolsTestUtils.sharedDeleteScan(8) + // Three of eight partitions use the delete file: too few for the common block. + val (common, partitions) = IcebergScanPoolsTestUtils.sharedDeleteScan(8, users = 3) val layout = layoutOf(common, partitions) val data = new PartitionPlanData { override def length: Int = partitions.length @@ -230,7 +232,7 @@ class PlanDataInjectorSuite extends AnyFunSuite { bytes.size() } val perPartition = taskSize(2) - taskSize(1) - assert(taskSize(7) - taskSize(1) == 6 * perPartition, "each input partition adds the same") + assert(taskSize(3) - taskSize(1) == 2 * perPartition, "each input partition adds the same") assert( perPartition > data(0).length, "each input partition carries its own copy of the shared delete file") From b1fc2b80b6241fe589a321d47e0b77f67c7aeadf Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 4 Nov 2026 17:29:46 -0800 Subject: [PATCH 4/8] style: drop an interpolator with no placeholders --- .../test/scala/org/apache/comet/CometIcebergNativeSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 1b72492f2cc..9c993d0f8fb 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -5311,7 +5311,7 @@ class CometIcebergNativeSuite // partition share the entry. assert( tasksWithResidualPerPartition.exists(_ > 1), - s"expected a partition with several tasks carrying a residual: " + + "expected a partition with several tasks carrying a residual: " + tasksWithResidualPerPartition) // Every task uses it, so it ships once: in the common block when several partitions have // tasks, otherwise with the one partition that does. From 40c7c3e50cb6626dcb86e26368d89af922a1277f Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Thu, 5 Nov 2026 10:16:02 -0800 Subject: [PATCH 5/8] refactor: intern Iceberg per-file pools per partition in the existing serde pass Review feedback: build each partition's pools in the pass that already interns them, and ship them in the partition's existing IcebergScan.common, so the injector merges the two with mergeFrom. This drops the global-then-split step (IcebergScanPools, Layout, SerializedPools), the IcebergFilePools message and own_pools field, and the half-of-partitions threshold: the common block now holds none of the per-file pools. The lazy driver-side build (PartitionPlanData, PlanDataSlice) moves out of this PR, so CometExecRDD, CometNativeShuffleInputRDD, NativeExecContext, CometNativeScanExec and the Delta contrib scan are unchanged from main. --- .github/workflows/pr_build_linux.yml | 1 - .github/workflows/pr_build_macos.yml | 1 - .../sql/comet/CometDeltaNativeScanExec.scala | 2 +- docs/source/user-guide/latest/iceberg.md | 13 +- native/core/src/execution/planner.rs | 26 - native/proto/src/proto/operator.proto | 32 +- .../operator/CometIcebergNativeScan.scala | 261 +++++----- .../serde/operator/IcebergScanPools.scala | 408 --------------- .../apache/spark/sql/comet/CometExecRDD.scala | 13 +- .../comet/CometIcebergNativeScanExec.scala | 22 +- .../spark/sql/comet/CometNativeScanExec.scala | 2 +- .../spark/sql/comet/PartitionPlanData.scala | 89 ---- .../shuffle/CometNativeShuffleInputRDD.scala | 22 +- .../apache/spark/sql/comet/operators.scala | 39 +- .../merge_on_read_partition_layout.sql | 8 +- .../apache/comet/CometFuzzIcebergSuite.scala | 9 +- .../comet/CometIcebergNativeSuite.scala | 108 ++-- .../operator/IcebergScanPoolsSuite.scala | 465 ------------------ .../operator/IcebergScanPoolsTestUtils.scala | 175 ------- .../comet/CometScanWithPlanDataSuite.scala | 7 +- .../sql/comet/PlanDataInjectorSuite.scala | 175 +++---- .../CometNativeShuffleInputRDDSuite.scala | 57 +-- 22 files changed, 326 insertions(+), 1609 deletions(-) delete mode 100644 spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala delete mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala delete mode 100644 spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala delete mode 100644 spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 69275ac3ae4..478339062b0 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -657,7 +657,6 @@ jobs: org.apache.comet.objectstore.AbfsAuthParitySuite org.apache.comet.objectstore.AbfsScanFallbackSuite org.apache.comet.serde.operator.CometIcebergNativeScanSuite - org.apache.comet.serde.operator.IcebergScanPoolsSuite org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite org.apache.spark.sql.CometCollationSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 9dd44265b2d..3b4b49c728a 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -292,7 +292,6 @@ jobs: org.apache.comet.objectstore.AbfsAuthParitySuite org.apache.comet.objectstore.AbfsScanFallbackSuite org.apache.comet.serde.operator.CometIcebergNativeScanSuite - org.apache.comet.serde.operator.IcebergScanPoolsSuite org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite org.apache.spark.sql.CometCollationSuite 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 f2386a2eeb7..bcd1d82f906 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 -> partitionPlanData), + Map(sourceKey -> 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 b582777b4c9..8e241fe0a3f 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -197,13 +197,12 @@ These `s3.*` storage properties are not specific to the Hive catalog shown here. ### 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. Entries that at least half of the -partitions need, such as unpartitioned equality deletes, are sent once per executor instead. A -coalesce without a shuffle (`coalesce(n)`, or a write that coalesces its input) combines several -input partitions into one task, and that task then carries a copy of a shared delete file for -every input partition that uses it. Scans where delete files are shared by many partitions, but -fewer than half, can make coalesced tasks large. Compacting delete files, or repartitioning instead of coalescing, avoids -this. +that its own files need, the way Spark's Iceberg reader does. A coalesce without a shuffle +(`coalesce(n)`, or a write that coalesces its input) combines several input partitions into one +task, and that task then carries a copy of a shared delete file for every input partition that +uses it. Scans where many partitions share delete files, such as tables with unpartitioned +equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning +instead of coalescing, avoids this. ### Current limitations diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d265309fbd5..603bba872bc 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1967,14 +1967,6 @@ impl PhysicalPlanner { OpStruct::IcebergScan(scan) => { // Extract common data and single partition's file tasks // Per-partition injection happens in Scala before sending to native - if scan.own_pools.is_some() { - // The tasks would index pools that are not in common. - return Err(GeneralError( - "IcebergScan still carries its per-partition pools: the JVM did not \ - merge them into common" - .into(), - )); - } let common = scan .common .as_ref() @@ -3614,24 +3606,6 @@ mod tests { } } - #[test] - fn iceberg_scan_rejects_unmerged_partition_pools() { - let op = spark_operator::Operator { - op_struct: Some(OpStruct::IcebergScan(spark_operator::IcebergScan { - common: Some(spark_operator::IcebergScanCommon::default()), - file_scan_tasks: vec![], - own_pools: Some(spark_operator::IcebergFilePools::default()), - })), - ..Default::default() - }; - let planner = PhysicalPlanner::default(); - let err = planner.create_plan(&op, &mut vec![], 1).unwrap_err(); - assert!( - format!("{err}").contains("did not merge them into common"), - "got: {err}" - ); - } - #[test] fn delta_scan_errors_without_delta_feature() { let op = delta_spark_envelope(spark_operator::DeltaSparkScan { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 7547ba7222a..d8dc23565fc 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -377,8 +377,7 @@ enum IcebergPredicateOperator { NotIn = 14; } -// Common data shared by all partitions of a scan. Each task deserializes its own copy, so it -// holds no per-file pool entries (see IcebergScan.own_pools). +// Common data shared by all partitions in split mode (sent once, captured in closure) message IcebergScanCommon { // Catalog-specific configuration for FileIO (credentials, S3/GCS config, etc.) map catalog_properties = 1; @@ -389,10 +388,10 @@ message IcebergScanCommon { // Schema to read repeated SparkStructField required_schema = 3; - // Deduplication pools. In the common block the driver ships, the five per-file pools + // 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 ships the entries its tasks use in - // IcebergScan.own_pools, and the JVM merges them in before the plan reaches native. + // delete_file_path_pool) are empty: each partition ships the entries its tasks use in its own + // IcebergScan.common, and the JVM merges the two before the plan reaches native. 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 @@ -436,28 +435,12 @@ message IcebergScanCommon { } message IcebergScan { - // Common data shared across partitions (pools, metadata, catalog props). Unset in the - // per-partition message the driver ships. In the plan that reaches native it holds everything - // the partition's tasks index into. + // 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 repeated IcebergFileScanTask file_scan_tasks = 2; - - // Only in the per-partition message the driver ships: the per-file pool entries this - // partition's tasks use. The JVM merges them into common before the plan reaches native, and - // native rejects a plan that still has them (see IcebergScanPools). - IcebergFilePools own_pools = 3; -} - -// The five per-file pools of IcebergScanCommon, holding one partition's entries. Field meanings -// match the IcebergScanCommon fields of the same name. -message IcebergFilePools { - repeated PartitionData partition_data_pool = 1; - repeated DeleteFileList delete_files_pool = 2; - repeated IcebergPredicate residual_pool = 3; - repeated IcebergDeleteFile delete_file_pool = 4; - repeated string delete_file_path_pool = 5; } // Common data shared by all partitions for native Lance scans. @@ -517,8 +500,7 @@ message IcebergFileScanTask { // decodes this directly; no KMS unwrap is needed on the native side. optional bytes key_metadata = 6; - // Indices into the deduplication pools. The four non-file pools are in the common block; the - // per-file pools are in IcebergScan.own_pools until the JVM merges them (see IcebergScanPools). + // Indices into IcebergScan deduplication pools uint32 schema_idx = 15; optional uint32 partition_spec_idx = 17; optional uint32 name_mapping_idx = 18; 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 9608f3b5124..3d9ee41060d 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, PartitionPlanData} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeExec} 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._ @@ -427,6 +427,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 +587,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) @@ -951,16 +952,16 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * one pass from the DPP-filtered partitions. * * 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 (schemas, specs, - * partition data, residuals, delete files) 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. + * 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 has to stay small. The per-file pool entries - * (delete files, their paths and sets, partition values, residuals) ship with each partition - * whose tasks use them, and the common block holds none of them (see IcebergScanPools). Each - * partition's block is built when its task is serialized. Keeping every delete file in the - * common block made each partition's plan as large as the whole scan's delete metadata. + * 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, into that + * partition's IcebergScan.common, 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 @@ -974,40 +975,58 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * @param metadata * Pre-extracted Iceberg metadata from CometScanRule * @return - * The common block's bytes, and each partition's bytes built when its task is serialized + * Tuple of (commonBytes, perPartitionBytes) for native execution */ def serializePartitions( scanExec: BatchScanExec, output: Seq[Attribute], - metadata: CometIcebergNativeScanMetadata): (Array[Byte], PartitionPlanData) = { + metadata: CometIcebergNativeScanMetadata): (Array[Byte], Array[Array[Byte]]) = { 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]() + + // The per-file pools of one partition. Each partition ships the entries its own tasks use in + // its IcebergScan.common, and the injector merges them into the common block, 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]() + // Individual delete files are interned into a flat pool; deleteFiles 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 deleteFiles = mutable.HashMap[OperatorOuterClass.IcebergDeleteFile, Int]() + val deleteFileSets = 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 deleteFilePaths = mutable.HashMap[String, Int]() + val residuals = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Int]() + + def internDeleteFilePath(path: String): Int = + deleteFilePaths.getOrElseUpdate( + path, { + val idx = deleteFilePaths.size + builder.addDeleteFilePathPool(path) + idx + }) + + def isEmpty: Boolean = + partitionData.isEmpty && deleteFiles.isEmpty && deleteFilePaths.isEmpty && + residuals.isEmpty + } + // Entries in the per-file pools, summed over partitions, for the log. + var partitionDataEntries = 0 + var deleteFileEntries = 0 + var deleteFileSetEntries = 0 + var deleteFilePathEntries = 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. @@ -1094,6 +1113,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]) @@ -1155,24 +1175,24 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit contentFileClass, fileScanTaskClass, deleteFileClass, - internDeleteFilePath) + pools.internDeleteFilePath) 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( + pools.deleteFiles.getOrElseUpdate( df, { - val idx = deleteFileToPoolIndex.size - commonBuilder.addDeleteFilePool(df) + val idx = pools.deleteFiles.size + pools.builder.addDeleteFilePool(df) idx }) } - val deleteFilesIdx = deleteFilesToPoolIndex.getOrElseUpdate( + 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) @@ -1185,10 +1205,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) @@ -1201,7 +1221,8 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit taskBuilder, commonBuilder, partitionSpecToPoolIndex, - partitionDataToPoolIndex) + pools.builder, + pools.partitionData) metadata.nameMapping.foreach { nm => val nmIdx = nameMappingToPoolIndex.getOrElseUpdate( @@ -1218,6 +1239,12 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } + if (!pools.isEmpty) partitionBuilder.setCommon(pools.builder) + partitionDataEntries += pools.partitionData.size + deleteFileEntries += pools.deleteFiles.size + deleteFileSetEntries += pools.deleteFileSets.size + deleteFilePathEntries += pools.deleteFilePaths.size + residualEntries += pools.residuals.size perPartitionBuilders += partitionBuilder.build() } case other if other.getClass.getName == "org.apache.spark.rdd.ParallelCollectionRDD" => @@ -1251,11 +1278,11 @@ 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, + deleteFilePathEntries, + residualEntries) val avgDedup = if (totalTasks == 0) { "0.0" @@ -1269,81 +1296,60 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } - // The per-file pool entries ship with each partition that uses them, so the common block - // every task deserializes doesn't grow with the number of files. - val original = commonBuilder.build() - // Entries at least half of the partitions use stay in the common block; the rest ship with - // each partition that uses them. - val layout = - IcebergScanPools.Layout(new IcebergScanPools.MessagePools(original), perPartitionBuilders) - val common = IcebergScanPools.commonBlock(original, layout) - - // Per-pool byte sizes of the whole scan, to diagnose large plan data. Sizes sum as Long because - // a single pool at or past protobuf's 2 GiB message limit overflows the int getSerializedSize, - // and they are logged before any toByteArray so the breakdown survives a failed allocation. - // String pools carry JSON or file paths, whose serialized size is the UTF-8 length. + // Per-pool byte sizes to diagnose an oversized common message. Sizes sum as Long because a + // single pool at or past protobuf's 2 GiB message limit overflows the int getSerializedSize, + // and the logging runs before toByteArray so the breakdown survives even if that allocation + // fails. String pools carry JSON or file paths, whose serialized size is the UTF-8 length. def sumSizes(sizes: Iterator[Int]): Long = sizes.map(_.toLong).sum def sumStrBytes(strings: mutable.Buffer[String]): Long = strings.iterator.map(_.getBytes(UTF_8).length.toLong).sum - val scanPoolBytes: Seq[(String, Int, Long)] = Seq( - ("schema", original.getSchemaPoolCount, sumStrBytes(original.getSchemaPoolList.asScala)), + val commonPoolBytes: Seq[(String, Int, Long)] = Seq( + ( + "schema", + commonBuilder.getSchemaPoolCount, + sumStrBytes(commonBuilder.getSchemaPoolList.asScala)), ( "partition_type", - original.getPartitionTypePoolCount, - sumStrBytes(original.getPartitionTypePoolList.asScala)), + commonBuilder.getPartitionTypePoolCount, + sumStrBytes(commonBuilder.getPartitionTypePoolList.asScala)), ( "partition_spec", - original.getPartitionSpecPoolCount, - sumStrBytes(original.getPartitionSpecPoolList.asScala)), + commonBuilder.getPartitionSpecPoolCount, + sumStrBytes(commonBuilder.getPartitionSpecPoolList.asScala)), ( "name_mapping", - original.getNameMappingPoolCount, - sumStrBytes(original.getNameMappingPoolList.asScala)), + commonBuilder.getNameMappingPoolCount, + sumStrBytes(commonBuilder.getNameMappingPoolList.asScala)), ( "project_field_ids", - original.getProjectFieldIdsPoolCount, - sumSizes(original.getProjectFieldIdsPoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "partition_data", - original.getPartitionDataPoolCount, - sumSizes(original.getPartitionDataPoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "delete_file", - original.getDeleteFilePoolCount, - sumSizes(original.getDeleteFilePoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "delete_file_path", - original.getDeleteFilePathPoolCount, - sumStrBytes(original.getDeleteFilePathPoolList.asScala)), - ( - "delete_files_set", - original.getDeleteFilesPoolCount, - sumSizes(original.getDeleteFilesPoolList.asScala.iterator.map(_.getSerializedSize))), - ( - "residual", - original.getResidualPoolCount, - sumSizes(original.getResidualPoolList.asScala.iterator.map(_.getSerializedSize)))) + commonBuilder.getProjectFieldIdsPoolCount, + sumSizes( + commonBuilder.getProjectFieldIdsPoolList.asScala.iterator.map(_.getSerializedSize)))) + + val perPartitionSizes = perPartitionBuilders.map(_.getSerializedSize.toLong) + val perPartitionTotal = perPartitionSizes.sum + val perPartitionMax = if (perPartitionSizes.isEmpty) 0L else perPartitionSizes.max logInfo(s"IcebergScan: $totalTasks tasks, ${allPoolSizes.size} pools ($avgDedup% avg dedup)") logInfo( - " Scan pools (unique/bytes): " + - scanPoolBytes + " Common pools (unique/bytes): " + + commonPoolBytes .map { case (name, count, bytes) => s"$name=$count/$bytes" } .mkString(", ") + - s"; total=${scanPoolBytes.map(_._3).sum} bytes") - - val commonBytes = common.toByteArray - // Tasks only, still indexing the full pools. Each partition's message is built from these - // when its task is serialized, so the driver holds every partition's tasks once and never - // every partition's message at the same time. - val taskBytes = perPartitionBuilders.map(_.toByteArray).toArray - logInfo(s" Common block: ${commonBytes.length} bytes. Per-partition messages: " + - s"${taskBytes.length}, task bytes total=${taskBytes.iterator.map(_.length.toLong).sum}, " + - s"max=${if (taskBytes.isEmpty) 0 else taskBytes.iterator.map(_.length).max}") - - // Keep the full pools as bytes, not as the parsed message, which takes several times the heap. - val pools = new IcebergScanPools.SerializedPools(original.toByteArray) - (commonBytes, new IcebergPartitionPlanData(pools, layout, taskBytes)) + s"; total=${commonPoolBytes.map(_._3).sum} bytes") + logInfo( + " Per-partition pool entries, summed over partitions: " + + s"partition_data=$partitionDataEntries, delete_file=$deleteFileEntries, " + + s"delete_file_path=$deleteFilePathEntries, delete_files_set=$deleteFileSetEntries, " + + s"residual=$residualEntries") + logInfo( + s" Per-partition messages: ${perPartitionBuilders.size}, " + + s"total=$perPartitionTotal bytes, max=$perPartitionMax bytes") + + val commonBytes = commonBuilder.build().toByteArray + val perPartitionBytes = perPartitionBuilders.map(_.toByteArray).toArray + + (commonBytes, perPartitionBytes) } override def createExec(nativeOp: Operator, op: CometBatchScanExec): CometNativeExec = { @@ -1364,36 +1370,3 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit CometIcebergNativeScanExec(nativeOp, op.wrapped, op.session, metadataLocation, metadata) } } - -/** - * A scan's per-partition messages, each built when its task is serialized: the partition's tasks - * plus the per-file pool entries they use (see IcebergScanPools.partitionPayload). - */ -private[operator] final class IcebergPartitionPlanData( - private val pools: IcebergScanPools.SerializedPools, - layout: IcebergScanPools.Layout, - private val taskBytes: Array[Array[Byte]]) - extends PartitionPlanData { - - override def length: Int = taskBytes.length - - // Same table state read into the same partitions, compared without building any partition. - // The layout follows from the pools and the tasks. - override def sameAs(other: PartitionPlanData): Boolean = other match { - case that: IcebergPartitionPlanData => - java.util.Arrays.equals(that.pools.bytes, pools.bytes) && - that.taskBytes.length == taskBytes.length && - taskBytes.indices.forall(i => java.util.Arrays.equals(taskBytes(i), that.taskBytes(i))) - case _ => false - } - - override def apply(i: Int): Array[Byte] = - if (pools.isEmpty) { - // With no per-file entries in the scan, the stored task bytes already are each message. - taskBytes(i) - } else { - IcebergScanPools - .partitionPayload(pools, layout, OperatorOuterClass.IcebergScan.parseFrom(taskBytes(i))) - .toByteArray - } -} diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala b/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala deleted file mode 100644 index 5110c74cab6..00000000000 --- a/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala +++ /dev/null @@ -1,408 +0,0 @@ -/* - * 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.serde.operator - -import java.nio.charset.StandardCharsets.UTF_8 - -import scala.collection.mutable -import scala.jdk.CollectionConverters._ - -import com.google.protobuf.CodedInputStream - -import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteFile, IcebergFilePools, IcebergPredicate, IcebergScan, IcebergScanCommon, PartitionData} - -/** - * Decides where an Iceberg scan's pool entries ship. - * - * Five pools grow with the number of files in the scan: delete file paths, delete files, delete - * file sets, partition values and residuals. An entry that at least half of the partitions use, - * such as an unpartitioned equality delete, ships once in the common block. Every other entry - * ships with each partition whose tasks use it, as Iceberg's own Spark reader ships each task's - * delete files with the task. Since each entry in the common block is needed by at least half of - * the tasks, the common block is at most twice what an average partition's tasks need. - * - * A partition's tasks index the common entries first, then the partition's own. On the driver, - * [[partitionPayload]] builds one partition's message. On the executor, [[forPartition]] appends - * its own entries to the common block. - * - * Known limit: a coalesce without a shuffle puts several partitions' messages in one task, and - * each carries its own copy of the entries they share. Entries that fewer than half of the - * partitions use, such as the partition-scoped deletes of a large partition, are then repeated - * once per coalesced partition. - */ -private[apache] object IcebergScanPools { - - /** Read access to a scan's five per-file pools. */ - trait FilePools { - def pathCount: Int - def path(i: Int): String - def deleteCount: Int - def delete(i: Int): IcebergDeleteFile - def setCount: Int - def set(i: Int): DeleteFileList - def partitionValueCount: Int - def partitionValue(i: Int): PartitionData - def residualCount: Int - def residual(i: Int): IcebergPredicate - - def isEmpty: Boolean = - pathCount == 0 && deleteCount == 0 && setCount == 0 && partitionValueCount == 0 && - residualCount == 0 - } - - /** The pools of a built message. */ - final class MessagePools(common: IcebergScanCommon) extends FilePools { - override def pathCount: Int = common.getDeleteFilePathPoolCount - override def path(i: Int): String = common.getDeleteFilePathPool(i) - override def deleteCount: Int = common.getDeleteFilePoolCount - override def delete(i: Int): IcebergDeleteFile = common.getDeleteFilePool(i) - override def setCount: Int = common.getDeleteFilesPoolCount - override def set(i: Int): DeleteFileList = common.getDeleteFilesPool(i) - override def partitionValueCount: Int = common.getPartitionDataPoolCount - override def partitionValue(i: Int): PartitionData = common.getPartitionDataPool(i) - override def residualCount: Int = common.getResidualPoolCount - override def residual(i: Int): IcebergPredicate = common.getResidualPool(i) - } - - /** - * The pools of a serialized `IcebergScanCommon`, parsing only the entries asked for. The bytes - * plus one offset per entry take a fraction of the heap of the parsed message, which is several - * times its serialized size. - */ - final class SerializedPools(val bytes: Array[Byte]) extends FilePools { - // Each entry's offset (high 32 bits) and length (low 32 bits), per pool, in wire order. - private val located: Array[Array[Long]] = { - val entries = Array.fill(5)(mutable.ArrayBuilder.make[Long]) - val in = CodedInputStream.newInstance(bytes) - var tag = in.readTag() - while (tag != 0) { - val pool = SerializedPools.poolOf(tag >>> 3) - if (pool >= 0) { - val length = in.readRawVarint32() - entries(pool) += (in.getTotalBytesRead.toLong << 32) | length.toLong - in.skipRawBytes(length) - } else { - val _ = in.skipField(tag) - } - tag = in.readTag() - } - entries.map(_.result()) - } - - private def offset(pool: Int, i: Int): Int = (located(pool)(i) >>> 32).toInt - private def length(pool: Int, i: Int): Int = (located(pool)(i) & 0xffffffffL).toInt - - override def pathCount: Int = located(SerializedPools.Paths).length - override def path(i: Int): String = - new String(bytes, offset(SerializedPools.Paths, i), length(SerializedPools.Paths, i), UTF_8) - override def deleteCount: Int = located(SerializedPools.Deletes).length - override def delete(i: Int): IcebergDeleteFile = - IcebergDeleteFile - .parser() - .parseFrom(bytes, offset(SerializedPools.Deletes, i), length(SerializedPools.Deletes, i)) - override def setCount: Int = located(SerializedPools.Sets).length - override def set(i: Int): DeleteFileList = - DeleteFileList - .parser() - .parseFrom(bytes, offset(SerializedPools.Sets, i), length(SerializedPools.Sets, i)) - override def partitionValueCount: Int = located(SerializedPools.PartitionValues).length - override def partitionValue(i: Int): PartitionData = - PartitionData - .parser() - .parseFrom( - bytes, - offset(SerializedPools.PartitionValues, i), - length(SerializedPools.PartitionValues, i)) - override def residualCount: Int = located(SerializedPools.Residuals).length - override def residual(i: Int): IcebergPredicate = - IcebergPredicate - .parser() - .parseFrom( - bytes, - offset(SerializedPools.Residuals, i), - length(SerializedPools.Residuals, i)) - } - - private object SerializedPools { - val Paths = 0 - val Deletes = 1 - val Sets = 2 - val PartitionValues = 3 - val Residuals = 4 - - def poolOf(field: Int): Int = - if (field == IcebergScanCommon.DELETE_FILE_PATH_POOL_FIELD_NUMBER) Paths - else if (field == IcebergScanCommon.DELETE_FILE_POOL_FIELD_NUMBER) Deletes - else if (field == IcebergScanCommon.DELETE_FILES_POOL_FIELD_NUMBER) Sets - else if (field == IcebergScanCommon.PARTITION_DATA_POOL_FIELD_NUMBER) PartitionValues - else if (field == IcebergScanCommon.RESIDUAL_POOL_FIELD_NUMBER) Residuals - else -1 - } - - private val Unused = -1 - - private def checkIndex(name: String, i: Int, size: Int): Unit = - if (i < 0 || i >= size) { - throw new IllegalStateException( - s"Iceberg $name index $i is out of range: the pool has $size entries") - } - - /** - * Which entries at least half of the non-empty partitions use, and each one's index in the - * common block. With fewer than two non-empty partitions there are none. - */ - final class Layout private[IcebergScanPools] ( - val path: Array[Int], - val delete: Array[Int], - val set: Array[Int], - val partitionValue: Array[Int], - val residual: Array[Int]) { - - val pathCount: Int = path.count(_ != Unused) - val deleteCount: Int = delete.count(_ != Unused) - val setCount: Int = set.count(_ != Unused) - val partitionValueCount: Int = partitionValue.count(_ != Unused) - val residualCount: Int = residual.count(_ != Unused) - } - - object Layout { - - /** Finds the entries of `pools` that at least half of the non-empty `partitions` use. */ - def apply(pools: FilePools, partitions: collection.IndexedSeq[IcebergScan]): Layout = { - val paths = new Users("delete_file_path", pools.pathCount) - val deletes = new Users("delete_file", pools.deleteCount) - val sets = new Users("delete_files_set", pools.setCount) - val partitionValues = new Users("partition_data", pools.partitionValueCount) - val residuals = new Users("residual", pools.residualCount) - var nonEmpty = 0 - partitions.indices.foreach { p => - val tasks = partitions(p).getFileScanTasksList - if (!tasks.isEmpty) nonEmpty += 1 - tasks.asScala.foreach { task => - if (task.hasResidualIdx) residuals.use(task.getResidualIdx, p) - if (task.hasPartitionDataIdx) partitionValues.use(task.getPartitionDataIdx, p) - if (task.hasDeleteFilesIdx && sets.use(task.getDeleteFilesIdx, p)) { - val set = pools.set(task.getDeleteFilesIdx) - (0 until set.getDeleteFileIndicesCount).foreach { j => - val d = set.getDeleteFileIndices(j) - if (deletes.use(d, p)) paths.use(pools.delete(d).getFilePathIdx, p) - } - } - } - } - // A set's delete files and a delete file's path have at least the set's users, so an entry - // in the common block only points at entries that are there too. - val half = if (nonEmpty >= 2) (nonEmpty + 1) / 2 else Int.MaxValue - new Layout( - paths.usedBy(half), - deletes.usedBy(half), - sets.usedBy(half), - partitionValues.usedBy(half), - residuals.usedBy(half)) - } - - /** How many partitions use each entry of one pool. */ - private final class Users(name: String, size: Int) { - private val lastUser = Array.fill(size)(Unused) - private val users = new Array[Int](size) - - // Records that partition p uses entry i. Returns true the first time p does, so the caller - // records the entries it points to as well. - def use(i: Int, p: Int): Boolean = { - checkIndex(name, i, size) - if (lastUser(i) == p) { - false - } else { - lastUser(i) = p - users(i) += 1 - true - } - } - - // Each entry's index among the entries at least `partitions` partitions use, or Unused. - def usedBy(partitions: Int): Array[Int] = { - var next = 0 - users.map { n => - if (n >= partitions) { - next += 1 - next - 1 - } else { - Unused - } - } - } - } - } - - /** - * The common block every task gets: the scan without its per-file pools, plus the entries at - * least half of the partitions use. - */ - def commonBlock(original: IcebergScanCommon, layout: Layout): IcebergScanCommon = { - val builder = original.toBuilder - .clearDeleteFilePathPool() - .clearDeleteFilePool() - .clearDeleteFilesPool() - .clearPartitionDataPool() - .clearResidualPool() - def shared(positions: Array[Int]): Iterator[Int] = - positions.indices.iterator.filter(positions(_) != Unused) - shared(layout.path).foreach(i => - builder.addDeleteFilePathPool(original.getDeleteFilePathPool(i))) - // Every partition that uses a delete file uses its path, and every one that uses a set uses - // its delete files, so those are in the common block too. - shared(layout.delete).foreach { i => - val deleteFile = original.getDeleteFilePool(i) - builder.addDeleteFilePool( - deleteFile.toBuilder.setFilePathIdx(layout.path(deleteFile.getFilePathIdx))) - } - shared(layout.set).foreach { i => - val set = original.getDeleteFilesPool(i) - val remapped = DeleteFileList.newBuilder() - (0 until set.getDeleteFileIndicesCount).foreach { j => - remapped.addDeleteFileIndices(layout.delete(set.getDeleteFileIndices(j))) - } - builder.addDeleteFilesPool(remapped) - } - shared(layout.partitionValue).foreach(i => - builder.addPartitionDataPool(original.getPartitionDataPool(i))) - shared(layout.residual).foreach(i => builder.addResidualPool(original.getResidualPool(i))) - builder.build() - } - - /** - * One partition's message, built on the driver. `tasks` indexes into `pools`. The result - * carries the tasks, renumbered to index the common block's entries first, and in `own_pools` - * the other entries they use. - */ - def partitionPayload(pools: FilePools, layout: Layout, tasks: IcebergScan): IcebergScan = { - val own = IcebergFilePools.newBuilder() - val ownPath = mutable.HashMap.empty[Int, Int] - val ownDelete = mutable.HashMap.empty[Int, Int] - val ownSet = mutable.HashMap.empty[Int, Int] - val ownPartitionValue = mutable.HashMap.empty[Int, Int] - val ownResidual = mutable.HashMap.empty[Int, Int] - - def path(i: Int): Int = { - checkIndex("delete_file_path", i, pools.pathCount) - if (layout.path(i) != Unused) { - layout.path(i) - } else { - ownPath.getOrElseUpdate( - i, { - own.addDeleteFilePathPool(pools.path(i)) - layout.pathCount + own.getDeleteFilePathPoolCount - 1 - }) - } - } - - def delete(i: Int): Int = { - checkIndex("delete_file", i, pools.deleteCount) - if (layout.delete(i) != Unused) { - layout.delete(i) - } else { - ownDelete.getOrElseUpdate( - i, { - val deleteFile = pools.delete(i) - val pathIdx = path(deleteFile.getFilePathIdx) - own.addDeleteFilePool(deleteFile.toBuilder.setFilePathIdx(pathIdx)) - layout.deleteCount + own.getDeleteFilePoolCount - 1 - }) - } - } - - def set(i: Int): Int = { - checkIndex("delete_files_set", i, pools.setCount) - if (layout.set(i) != Unused) { - layout.set(i) - } else { - ownSet.getOrElseUpdate( - i, { - val set = pools.set(i) - val remapped = DeleteFileList.newBuilder() - (0 until set.getDeleteFileIndicesCount).foreach { j => - remapped.addDeleteFileIndices(delete(set.getDeleteFileIndices(j))) - } - own.addDeleteFilesPool(remapped) - layout.setCount + own.getDeleteFilesPoolCount - 1 - }) - } - } - - def partitionValue(i: Int): Int = { - checkIndex("partition_data", i, pools.partitionValueCount) - if (layout.partitionValue(i) != Unused) { - layout.partitionValue(i) - } else { - ownPartitionValue.getOrElseUpdate( - i, { - own.addPartitionDataPool(pools.partitionValue(i)) - layout.partitionValueCount + own.getPartitionDataPoolCount - 1 - }) - } - } - - def residual(i: Int): Int = { - checkIndex("residual", i, pools.residualCount) - if (layout.residual(i) != Unused) { - layout.residual(i) - } else { - ownResidual.getOrElseUpdate( - i, { - own.addResidualPool(pools.residual(i)) - layout.residualCount + own.getResidualPoolCount - 1 - }) - } - } - - val renumbered = tasks.getFileScanTasksList.asScala.map { task => - val builder = task.toBuilder - if (task.hasResidualIdx) builder.setResidualIdx(residual(task.getResidualIdx)) - if (task.hasPartitionDataIdx) { - builder.setPartitionDataIdx(partitionValue(task.getPartitionDataIdx)) - } - if (task.hasDeleteFilesIdx) builder.setDeleteFilesIdx(set(task.getDeleteFilesIdx)) - builder.build() - } - val result = tasks.toBuilder.clearFileScanTasks().addAllFileScanTasks(renumbered.asJava) - if (own.getDeleteFilePathPoolCount > 0 || own.getDeleteFilePoolCount > 0 || - own.getDeleteFilesPoolCount > 0 || own.getPartitionDataPoolCount > 0 || - own.getResidualPoolCount > 0) { - result.setOwnPools(own) - } - result.build() - } - - /** - * The scan one partition's native plan runs, built on the executor from the common block and - * the partition's message. Its tasks already index the common entries first, then its own, so - * this only appends its own entries to the common block. - */ - def forPartition(common: IcebergScanCommon, partition: IcebergScan): IcebergScan = { - val own = partition.getOwnPools - val merged = common.toBuilder - .addAllDeleteFilePathPool(own.getDeleteFilePathPoolList) - .addAllDeleteFilePool(own.getDeleteFilePoolList) - .addAllDeleteFilesPool(own.getDeleteFilesPoolList) - .addAllPartitionDataPool(own.getPartitionDataPoolList) - .addAllResidualPool(own.getResidualPoolList) - partition.toBuilder.clearOwnPools().setCommon(merged).build() - } -} 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 3b6ef900360..c903b7f70b0 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 @@ -33,16 +33,13 @@ import org.apache.comet.{CometExecIterator, CometRuntimeException, CometShuffleB /** * Partition that carries per-partition planning data, avoiding closure capture of all partitions. - * The data is built while the task is serialized (see [[PlanDataSlice]]). */ private[spark] class CometExecPartition( override val index: Int, val inputPartitions: Array[Partition], - planData: PlanDataSlice, + val planDataByKey: Map[String, Array[Byte]], val filePaths: Seq[String] = Seq.empty) - extends Partition { - def planDataByKey: Map[String, Array[Byte]] = planData.get -} + extends Partition /** * Unified RDD for Comet native execution. Non-shuffle input slots are `RDD[ArrowArrayStream]` @@ -61,7 +58,7 @@ private[spark] class CometExecRDD( sc: SparkContext, var inputRDDs: Seq[RDD[_]], commonByKey: Map[String, Array[Byte]], - @transient perPartitionByKey: Map[String, PartitionPlanData], + @transient perPartitionByKey: Map[String, Array[Array[Byte]]], serializedPlan: Array[Byte], planFingerprint: Long, defaultNumPartitions: Int, @@ -101,7 +98,7 @@ private[spark] class CometExecRDD( override protected def getPartitions: Array[Partition] = { (0 until numPartitions).map { idx => val inputParts = inputRDDs.map(_.partitions(idx)).toArray - val planData = new PlanDataSlice(perPartitionByKey, idx) + val planData = perPartitionByKey.map { case (key, arr) => key -> arr(idx) } val fp = if (perPartitionFilePaths.length > idx) perPartitionFilePaths(idx) else Seq.empty[String] new CometExecPartition(idx, inputParts, planData, fp) @@ -239,7 +236,7 @@ object CometExecRDD { sc: SparkContext, inputRDDs: Seq[RDD[_]], commonByKey: Map[String, Array[Byte]], - perPartitionByKey: Map[String, PartitionPlanData], + perPartitionByKey: Map[String, Array[Array[Byte]]], 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 9e4afca90b6..31e518cfcf8 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 @@ -79,7 +79,7 @@ case class CometIcebergNativeScanExec( /** * Lazy partition serialization, deferred until execution time. Triggered from `commonData` / - * `partitionPlanData` (via `PlanDataInjector.findAllPlanData`) and from + * `perPartitionData` (via `PlanDataInjector.findAllPlanData`) and from * `LazyIcebergMetric.value` (via Iceberg planning metrics). Lazy val semantics ensure single * evaluation across entry points. * @@ -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], PartitionPlanData) = { + @transient private lazy val serializedPartitionData: (Array[Byte], Array[Array[Byte]]) = { // 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,19 +122,11 @@ case class CometIcebergNativeScanExec( def commonData: Array[Byte] = serializedPartitionData._1 - /** Each partition's message, built when its task is serialized. */ - def partitionPlanData: PartitionPlanData = serializedPartitionData._2 - - /** - * Every partition's message, all built at once. For tests only: execution must use - * [[partitionPlanData]], which builds each one when its task is serialized. - */ - private[apache] def buildAllPartitionDataForTesting(): Array[Array[Byte]] = - partitionPlanData.toArray + def perPartitionData: Array[Array[Byte]] = serializedPartitionData._2 // numPartitions for execution - derived from actual DPP-filtered partitions // Only accessed during execution, not planning - def numPartitions: Int = partitionPlanData.length + def numPartitions: Int = perPartitionData.length // Like BatchScanExec on Spark 3.5 and later without a key-grouped partitioning, the scan does // not know its partition count while planning, because dynamic partition pruning decides which @@ -142,7 +134,7 @@ case class CometIcebergNativeScanExec( // optimizes a stage (ValidateRequirements after CoalesceShufflePartitions), which is before // CometPlanAdaptiveDynamicPruningFilters converts the DPP placeholders, so serializing the // tasks here would run a placeholder subquery. At execution, CometNativeExec sizes the native - // RDD from partitionPlanData instead. + // RDD from perPartitionData instead. override lazy val outputPartitioning: Partitioning = UnknownPartitioning(0) override lazy val outputOrdering: Seq[SortOrder] = Nil @@ -302,10 +294,10 @@ case class CometIcebergNativeScanExec( sparkContext, inputRDDs = Seq.empty, commonByKey = Map(injectorKey -> commonData), - perPartitionByKey = Map(injectorKey -> partitionPlanData), + perPartitionByKey = Map(injectorKey -> perPartitionData), serializedPlan = serializedPlan, planFingerprint = PlanDataInjector.planFingerprint(serializedPlan), - defaultNumPartitions = partitionPlanData.length, + defaultNumPartitions = perPartitionData.length, numOutputCols = output.length, nativeMetrics = nativeMetrics, subqueries = Seq.empty, 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 484d7919df7..8c8456b074a 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 -> partitionPlanData), + Map(sourceKey -> 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 deleted file mode 100644 index f18e19f2648..00000000000 --- a/spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala +++ /dev/null @@ -1,89 +0,0 @@ -/* - * 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 java.io.{ObjectInputStream, ObjectOutputStream} - -/** - * One scan's planning data for each partition. A scan can build a partition's bytes only when its - * task is serialized, the way Spark serializes each task's input partition, so the driver never - * holds every partition's bytes at once. - */ -private[apache] trait PartitionPlanData { - def length: Int - - /** Partition i's bytes. Must be safe to call from any thread, and more than once. */ - def apply(i: Int): Array[Byte] - - /** - * Whether `other` holds the same data, decided without building any partition. The default only - * knows identity. - */ - def sameAs(other: PartitionPlanData): Boolean = this eq other - - /** Every partition's bytes. For tests and diagnostics only: it builds them all. */ - def toArray: Array[Array[Byte]] = Array.tabulate(length)(apply) -} - -private[apache] object PartitionPlanData { - - private final class Built(val bytes: Array[Array[Byte]]) extends PartitionPlanData { - override def length: Int = bytes.length - override def apply(i: Int): Array[Byte] = bytes(i) - override def sameAs(other: PartitionPlanData): Boolean = other match { - case that: Built => - that.bytes.length == bytes.length && - bytes.indices.forall(i => java.util.Arrays.equals(bytes(i), that.bytes(i))) - case _ => false - } - } - - /** Data already built for every partition. */ - def apply(bytes: Array[Array[Byte]]): PartitionPlanData = new Built(bytes) - - /** Whether two scans carry the same data. */ - def same(a: PartitionPlanData, b: PartitionPlanData): Boolean = (a eq b) || a.sameAs(b) -} - -/** - * The planning data one task needs, keyed by scan. On the driver it holds only a reference to - * every scan's data and builds this partition's bytes while the task is serialized. On the - * executor it holds the bytes it received. - */ -private[comet] final class PlanDataSlice( - @transient private val source: Map[String, PartitionPlanData], - index: Int) - extends Serializable { - - @transient private var received: Map[String, Array[Byte]] = _ - - def get: Map[String, Array[Byte]] = - if (received != null) received else source.map { case (key, data) => key -> data(index) } - - private def writeObject(out: ObjectOutputStream): Unit = { - out.defaultWriteObject() - out.writeObject(get) - } - - private def readObject(in: ObjectInputStream): Unit = { - in.defaultReadObject() - received = in.readObject().asInstanceOf[Map[String, Array[Byte]]] - } -} 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 93468a7ba44..8770496bff2 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, PartitionPlanData, PlanDataSlice} +import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} 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, PartitionPlanData] = Map.empty, + @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, positionalRoundRobin: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, @@ -98,11 +98,13 @@ private[shuffle] class CometNativeShuffleInputRDD( // `leafRdd.partitions` on the executor, which would otherwise trigger getPartitions and // hit the @transient-null trap (e.g. CometExecRDD.perPartitionByKey). val inputParts = inputRDDs.map(_.partitions(i)).toArray - // Carry only this partition's slice of the plan data on the Partition object, built when the - // task is serialized. Keeping the full O(numPartitions) map out of the broadcast task binary - // avoids the 2GB ByteArrayOutputStream limit on jobs with tens of millions of partitions. - // Mirrors CometExecRDD.getPartitions. - new CometNativeShuffleInputPartition(i, inputParts, new PlanDataSlice(perPartitionByKey, i)) + // Slice this partition's plan data off the @transient full map here on the driver. Carrying + // only the per-partition slice on the Partition object (serialized per task) keeps the full + // O(numPartitions) map out of the broadcast task binary, which otherwise blows the 2GB + // ByteArrayOutputStream limit on jobs with tens of millions of partitions. Mirrors + // CometExecRDD.getPartitions. + val planDataByKey = perPartitionByKey.map { case (key, arr) => key -> arr(i) } + new CometNativeShuffleInputPartition(i, inputParts, planDataByKey) }.toArray override def compute( @@ -142,10 +144,8 @@ private[shuffle] class CometNativeShuffleInputRDD( private[shuffle] class CometNativeShuffleInputPartition( override val index: Int, val inputPartitions: Array[Partition], - planData: PlanDataSlice) - extends Partition { - def planDataByKey: Map[String, Array[Byte]] = planData.get -} + val planDataByKey: Map[String, Array[Byte]]) + extends Partition /** * Iterator handed to [[CometNativeShuffleWriter.write]] via Spark's ShuffleMapTask. Reports no 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 0ff0dbe86f7..43e66a36a2b 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 @@ -66,7 +66,7 @@ import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatibl import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} -import org.apache.comet.serde.operator.{CometSink, IcebergScanPools} +import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim /** @@ -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, PartitionPlanData]) = { + plan: SparkPlan): (Map[String, Array[Byte]], Map[String, Array[Array[Byte]]]) = { plan match { case iceberg: CometIcebergNativeScanExec => // Trigger Spark's standard prepare -> waitForSubqueries lifecycle so DPP @@ -469,7 +469,7 @@ private[comet] object PlanDataInjector extends Logging { // it. sendDriverMetrics guards on the execution id, so the standalone-scan path (which // posts from doExecuteColumnar) and this fused path never double-post. iceberg.sendDriverMetrics() - if (iceberg.commonData.nonEmpty && iceberg.partitionPlanData.length > 0) { + if (iceberg.commonData.nonEmpty && iceberg.perPartitionData.nonEmpty) { // A self-join/self-merge can put two scans of the same table (same metadata_location) // in one native plan, and they can match on scan_hash_code too while reading // different files, so the key also carries the op's plan_id (see withPlanId). @@ -477,7 +477,7 @@ private[comet] object PlanDataInjector extends Logging { // uses to look it up, keeps the two sides from drifting apart. IcebergPlanDataInjector.getKey(iceberg.nativeOp) match { case Some(key) => - (Map(key -> iceberg.commonData), Map(key -> iceberg.partitionPlanData)) + (Map(key -> iceberg.commonData), Map(key -> iceberg.perPartitionData)) case None => (Map.empty, Map.empty) } @@ -495,8 +495,8 @@ private[comet] object PlanDataInjector extends Logging { case s: CometLeafExec with CometScanWithPlanData => s.ensureSubqueriesResolved() s.sendDriverMetrics() - if (s.commonData.nonEmpty && s.partitionPlanData.length > 0) { - (Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.partitionPlanData)) + if (s.commonData.nonEmpty && s.perPartitionData.nonEmpty) { + (Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.perPartitionData)) } else { (Map.empty, Map.empty) } @@ -518,10 +518,13 @@ private[comet] object PlanDataInjector extends Logging { val results = plan.children.map(findAllPlanData) ( mergeByKey(results.map(_._1))(java.util.Arrays.equals(_, _)), - mergeByKey(results.map(_._2))(PartitionPlanData.same)) + mergeByKey(results.map(_._2))(samePartitionData)) } } + private def samePartitionData(a: Array[Array[Byte]], b: Array[Array[Byte]]): Boolean = + a.length == b.length && a.indices.forall(i => java.util.Arrays.equals(a(i), b(i))) + /** * Merges the planning data found under sibling subtrees. The injector looks data up by key, so * two scans sharing a key would both read whichever entry survived. The same data reached twice @@ -572,11 +575,16 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector { val partition = OperatorOuterClass.IcebergScan.parseFrom(partitionBytes) - // The common block plus the pool entries this partition's tasks use (see IcebergScanPools). - val pruned = IcebergScanPools.forPartition(common, partition) val scanBuilder = scan.toBuilder - scanBuilder.setCommon(pruned.getCommon) - scanBuilder.addAllFileScanTasks(pruned.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() } @@ -886,7 +894,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, PartitionPlanData], + @transient perPartitionByKey: Map[String, Array[Array[Byte]]], shuffleScanIndices: Set[Int], hasScanInput: Boolean, // Like plan data, file paths are sliced onto each task's Partition on the driver. @@ -1128,11 +1136,11 @@ abstract class CometNativeExec extends CometExec { // broadcast plan. val (firstNonBroadcastPlanRDD, firstNonBroadcastPlanNumPartitions) = firstNonBroadcastPlan.get._1 match { - // Plan-data scans, Iceberg included, are sized from partitionPlanData, which + // Plan-data scans, Iceberg included, are sized from perPartitionData, which // findAllPlanData above has resolved. A bucketed scan has one entry per bucket, after any // bucket coalescing, so the count still matches its HashPartitioning. case scan: CometScanWithPlanData => - (null.asInstanceOf[RDD[Any]], scan.partitionPlanData.length) + (null.asInstanceOf[RDD[Any]], scan.perPartitionData.length) case scan: CometIcebergNativeScanExec => (null.asInstanceOf[RDD[Any]], scan.numPartitions) case plan: CometNativeExec => @@ -1406,9 +1414,6 @@ trait CometScanWithPlanData { self: CometLeafExec => def commonData: Array[Byte] def perPartitionData: Array[Array[Byte]] - /** [[perPartitionData]] for execution. Override to build each partition's bytes on demand. */ - def partitionPlanData: PartitionPlanData = PartitionPlanData(perPartitionData) - // DPP / partition filters that may carry AQE SubqueryAdaptiveBroadcast // subqueries needing rewrite by CometPlanAdaptiveDynamicPruningFilters. // Default empty: scans with dedicated handling (CometNativeScanExec, 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 index 46cd0615d68..f986a07a466 100644 --- 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 @@ -15,10 +15,10 @@ -- specific language governing permissions and limitations -- under the License. --- The native scan ships a partition value, delete file or residual once in the common block when --- at least half of the Spark partitions use it, and otherwise with each Spark partition whose --- tasks use it (see IcebergScanPools). 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 +-- 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 diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala index 087a946983b..517ea6f27a9 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala @@ -22,12 +22,13 @@ package org.apache.comet import scala.util.Random import org.apache.spark.sql.Column +import org.apache.spark.sql.comet.CometIcebergNativeScanExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec 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.operator.IcebergScanPoolsTestUtils.storedPools +import org.apache.comet.serde.OperatorOuterClass class CometFuzzIcebergSuite extends CometFuzzIcebergBase { @@ -176,6 +177,12 @@ class CometFuzzIcebergSuite extends CometFuzzIcebergBase { } } + /** The pool entries a scan ships: the common block's, then each partition's own. */ + private def storedPools( + scan: CometIcebergNativeScanExec): Seq[OperatorOuterClass.IcebergScanCommon] = + OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) +: + scan.perPartitionData.toSeq.map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon) + // Asserts that predicate pushdown actually happens for the types that support it. For each // 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 diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 9c993d0f8fb..e1f70150d7d 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -49,7 +49,6 @@ import org.apache.comet.cloud.s3.MinioLocationScopedCredentialProvider import org.apache.comet.iceberg.{IcebergReflection, RESTCatalogHelper} import org.apache.comet.rules.CometScanRule import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde} -import org.apache.comet.serde.operator.IcebergScanPoolsTestUtils.{partitionScans, storedPools} import org.apache.comet.testing.{FuzzDataGenerator, SchemaGenOptions} /** @@ -145,6 +144,23 @@ class CometIcebergNativeSuite s"Expected fallback to Spark but found a CometIcebergNativeScanExec. Plan:\n$cometPlan") } + /** The pool entries a scan ships: the common block's, then each partition's own. */ + private def storedPools( + scan: CometIcebergNativeScanExec): Seq[OperatorOuterClass.IcebergScanCommon] = + OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) +: + scan.perPartitionData.toSeq.map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon) + + /** The pools each partition's tasks index into once the injector merges them. */ + private def partitionPools( + scan: CometIcebergNativeScanExec): Seq[OperatorOuterClass.IcebergScanCommon] = { + val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) + scan.perPartitionData.toSeq.map { bytes => + common.toBuilder + .mergeFrom(OperatorOuterClass.IcebergScan.parseFrom(bytes).getCommon) + .build() + } + } + /** Counts non-overlapping occurrences of `needle` within `haystack`. */ private def countByteOccurrences(haystack: Array[Byte], needle: Array[Byte]): Int = { require(needle.nonEmpty, "needle must be non-empty") @@ -804,8 +820,7 @@ class CometIcebergNativeSuite // 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 = partitionScans(scan).flatMap { injected => - val pools = injected.getCommon + val resolvedPaths = partitionPools(scan).flatMap { pools => pools.getDeleteFilePoolList.asScala.map(v => pools.getDeleteFilePathPool(v.getFilePathIdx)) } @@ -814,15 +829,12 @@ class CometIcebergNativeSuite s"expected one interned Puffin path referenced by every vector, got pool $pathPool " + s"and resolved paths $resolvedPaths") val occurrences = - scan - .buildAllPartitionDataForTesting() - .toSeq + scan.perPartitionData.toSeq .map(countByteOccurrences(_, pathPool.head.getBytes(UTF_8))) - // The path ships once in the common block when every partition uses it, otherwise once - // in each partition that does. + // 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( - inCommon <= 1 && occurrences.forall(_ <= 1) && (inCommon == 0 || occurrences.sum == 0), + inCommon == 0 && occurrences.forall(_ <= 1), s"Puffin path ${pathPool.head} serialized more than once: common $inCommon, " + s"partitions $occurrences") @@ -886,17 +898,15 @@ class CometIcebergNativeSuite .head .getLong(0) assert( - scan.buildAllPartitionDataForTesting().length == dataFileCount, + 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 = - scan - .buildAllPartitionDataForTesting() - .toSeq - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getOwnPools) + scan.perPartitionData.toSeq + .map(OperatorOuterClass.IcebergScan.parseFrom(_).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) @@ -986,7 +996,7 @@ class CometIcebergNativeSuite // 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.buildAllPartitionDataForTesting().toSeq + val planData = scans.head.commonData +: scans.head.perPartitionData.toSeq val duplicated = for { p <- distinctPaths (bytes, i) <- planData.zipWithIndex @@ -1080,10 +1090,8 @@ class CometIcebergNativeSuite s"per-file entries in the common block: $common") } def ownDeletes(scan: CometIcebergNativeScanExec): Seq[Int] = - scan - .buildAllPartitionDataForTesting() - .toSeq - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getOwnPools.getDeleteFilePoolCount) + scan.perPartitionData.toSeq + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon.getDeleteFilePoolCount) // One bucket per partition: each delete file goes to exactly one partition. val grouped = scanOf(grouped = true) @@ -1095,9 +1103,7 @@ class CometIcebergNativeSuite val perFile = scanOf(grouped = false) assertNoFileEntries(perFile) assert( - perFile - .buildAllPartitionDataForTesting() - .length > grouped.buildAllPartitionDataForTesting().length, + 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), @@ -1370,10 +1376,10 @@ class CometIcebergNativeSuite table.newRowDelta().addDeletes(deleteFile).commit() } - // An unpartitioned equality delete applies to every older data file, so most partitions use it - // and it ships once in the common block instead of in each task. The rows must still match - // Spark, including a newer row the delete does not apply to. - test("a global equality delete ships once in the common block and reads the right rows") { + // 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), @@ -1407,17 +1413,15 @@ class CometIcebergNativeSuite "id", Seq(2, 105, 210, 315), warehouseDir) - // A newer append the delete does not apply to: its id 2 row must survive. Partitions holding - // only this file do not use the delete, which still ships once while most partitions do. + // 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 = scan - .buildAllPartitionDataForTesting() - .toSeq - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getOwnPools) + val own = scan.perPartitionData.toSeq + .map(OperatorOuterClass.IcebergScan.parseFrom(_).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") @@ -1427,14 +1431,16 @@ class CometIcebergNativeSuite 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( - own.forall( - !_.getDeleteFilePoolList.asScala.exists(_.getContentType == "EQUALITY_DELETES")), - "a partition ships the equality delete itself") + 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.getDeleteFilePoolList.asScala.count(_.getContentType == "EQUALITY_DELETES") == 1, - s"expected the equality delete once in the common block: $common") + common.getDeleteFilePoolCount == 0, + s"the common block holds delete files: $common") spark.sql("DROP TABLE test_cat.db.global_eq_delete") } @@ -5161,8 +5167,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 - .buildAllPartitionDataForTesting() + val plannedTasks = icebergScans.head.perPartitionData .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum assert( @@ -5298,7 +5303,7 @@ class CometIcebergNativeSuite // 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 tasksWithResidualPerPartition = - scans.head.buildAllPartitionDataForTesting().toSeq.map { bytes => + scans.head.perPartitionData.toSeq.map { bytes => OperatorOuterClass.IcebergScan .parseFrom(bytes) .getFileScanTasksList @@ -5313,11 +5318,11 @@ class CometIcebergNativeSuite tasksWithResidualPerPartition.exists(_ > 1), "expected a partition with several tasks carrying a residual: " + tasksWithResidualPerPartition) - // Every task uses it, so it ships once: in the common block when several partitions have - // tasks, otherwise with the one partition that does. + // 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( - residualsPerPartition.sum == 1, - "expected the shared non-partition residual to dedup to 1 pool entry " + + 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") @@ -5970,8 +5975,7 @@ class CometIcebergNativeSuite val icebergScans = collectIcebergNativeScans(cometPlan) assert(icebergScans.nonEmpty, "Expected CometIcebergNativeScanExec in plan") val scan = icebergScans.head - val plannedTasks = scan - .buildAllPartitionDataForTesting() + val plannedTasks = scan.perPartitionData .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum val numSplits = scan.metrics("num_splits").value @@ -6240,8 +6244,7 @@ class CometIcebergNativeSuite val unprunedScan = collectIcebergNativeScans(unprunedPlan).head assert(unprunedScan.numPartitions == 1, s"Expected one split:\n$unprunedPlan") assert( - unprunedScan - .buildAllPartitionDataForTesting() + unprunedScan.perPartitionData .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum == 1, s"Expected one file scan task before pruning:\n$unprunedPlan") @@ -6278,8 +6281,7 @@ class CometIcebergNativeSuite collectIcebergDPPSubqueries(plan).nonEmpty, s"Expected executable DPP on the native scan:\n$plan") assert( - scan - .buildAllPartitionDataForTesting() + scan.perPartitionData .forall(OperatorOuterClass.IcebergScan.parseFrom(_).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") @@ -6317,8 +6319,7 @@ class CometIcebergNativeSuite val factB = "aqe_cat.db.single_split_join_b" val tables = Seq(factA, factB) def fileTaskCount(scan: CometIcebergNativeScanExec): Int = - scan - .buildAllPartitionDataForTesting() + scan.perPartitionData .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum @@ -6784,8 +6785,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 - .buildAllPartitionDataForTesting() + val plannedTasks = scan.perPartitionData .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) .sum val numSplits = scan.metrics("num_splits").value diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala deleted file mode 100644 index 2e2cf57682b..00000000000 --- a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsSuite.scala +++ /dev/null @@ -1,465 +0,0 @@ -/* - * 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.serde.operator - -import scala.jdk.CollectionConverters._ -import scala.util.Random - -import org.scalatest.funsuite.AnyFunSuite - -import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteFile, IcebergFilePools, IcebergFileScanTask, IcebergPredicate, IcebergScan, IcebergScanCommon, PartitionData} -import org.apache.comet.serde.operator.IcebergScanPoolsTestUtils._ - -class IcebergScanPoolsSuite extends AnyFunSuite { - - // A task with every pool index replaced by the entry it points at, so tasks compare equal no - // matter how their pools are numbered. - private case class ResolvedTask( - task: IcebergFileScanTask, - deletes: Option[Seq[(IcebergDeleteFile, String)]], - partitionData: Option[PartitionData], - residual: Option[IcebergPredicate]) - - private def resolveDelete(pools: IcebergScanCommon, idx: Int): (IcebergDeleteFile, String) = { - val deleteFile = pools.getDeleteFilePool(idx) - ( - deleteFile.toBuilder.clearFilePathIdx().build(), - pools.getDeleteFilePathPool(deleteFile.getFilePathIdx)) - } - - private def resolveSet(pools: IcebergScanCommon, idx: Int): Seq[(IcebergDeleteFile, String)] = - pools - .getDeleteFilesPool(idx) - .getDeleteFileIndicesList - .asScala - .toSeq - .map(d => resolveDelete(pools, d)) - - private def resolve(pools: IcebergScanCommon, task: IcebergFileScanTask): ResolvedTask = - ResolvedTask( - task.toBuilder.clearDeleteFilesIdx().clearPartitionDataIdx().clearResidualIdx().build(), - if (task.hasDeleteFilesIdx) Some(resolveSet(pools, task.getDeleteFilesIdx)) else None, - if (task.hasPartitionDataIdx) Some(pools.getPartitionDataPool(task.getPartitionDataIdx)) - else None, - if (task.hasResidualIdx) Some(pools.getResidualPool(task.getResidualIdx)) else None) - - // Each partition's message as the driver ships it, and the scan the executor builds from it. - private def shipAndInject(common: IcebergScanCommon, partitions: IndexedSeq[IcebergScan]) - : (IcebergScanCommon, IndexedSeq[(IcebergScan, IcebergScan)]) = { - // Read the pools the way the driver does after planning: from the serialized scan. - val pools = new IcebergScanPools.SerializedPools(common.toByteArray) - val layout = IcebergScanPools.Layout(pools, partitions) - val block = IcebergScanPools.commonBlock(common, layout) - ( - block, - partitions.map { p => - val payload = IcebergScanPools.partitionPayload(pools, layout, p) - payload -> IcebergScanPools.forPartition(block, payload) - }) - } - - /** Fails unless every task the executor runs resolves to the same entries as the original. */ - private def assertSameTasks( - common: IcebergScanCommon, - partitions: IndexedSeq[IcebergScan]): Unit = { - val (_, shipped) = shipAndInject(common, partitions) - partitions.indices.foreach { p => - val injected = shipped(p)._2 - val expected = partitions(p).getFileScanTasksList.asScala.map(resolve(common, _)) - val actual = injected.getFileScanTasksList.asScala.map(resolve(injected.getCommon, _)) - assert(actual == expected, s"partition $p resolves differently once shipped") - } - } - - // A pool entry, named by its pool and its index in the original pools. - private type Entry = (String, Int) - - // Which partitions use each entry: directly from a task, or through a set or delete file that - // holds it. - private def users( - common: IcebergScanCommon, - partitions: IndexedSeq[IcebergScan]): Map[Entry, Set[Int]] = { - val uses = for { - p <- partitions.indices - task <- partitions(p).getFileScanTasksList.asScala.toSeq - entry <- { - val set = if (task.hasDeleteFilesIdx) { - val deletes = common - .getDeleteFilesPool(task.getDeleteFilesIdx) - .getDeleteFileIndicesList - .asScala - .toSeq - .map(_.intValue) - Seq("set" -> task.getDeleteFilesIdx) ++ deletes.map("delete" -> _) ++ - deletes.map(d => "path" -> common.getDeleteFilePool(d).getFilePathIdx) - } else { - Seq.empty - } - val partitionData = - if (task.hasPartitionDataIdx) Seq("partition" -> task.getPartitionDataIdx) - else Seq.empty - val residual = - if (task.hasResidualIdx) Seq("residual" -> task.getResidualIdx) else Seq.empty - set ++ partitionData ++ residual - } - } yield entry -> p - uses.groupBy(_._1).map { case (entry, ps) => entry -> ps.map(_._2).toSet } - } - - // An entry's content, so entries can be matched across differently numbered pools. - private def content(pools: IcebergScanCommon, entry: Entry): Any = entry match { - case ("path", i) => pools.getDeleteFilePathPool(i) - case ("delete", i) => resolveDelete(pools, i) - case ("set", i) => resolveSet(pools, i) - case ("partition", i) => pools.getPartitionDataPool(i) - case ("residual", i) => pools.getResidualPool(i) - case other => throw new IllegalArgumentException(s"unknown pool entry $other") - } - - // Every entry of the five split pools, from position `from` on. - private def entries(pools: IcebergScanCommon, from: IcebergScanCommon): Seq[Entry] = { - def range(name: String, count: IcebergScanCommon => Int): Seq[Entry] = - (count(from) until count(pools)).map(name -> _) - range("path", _.getDeleteFilePathPoolCount) ++ - range("delete", _.getDeleteFilePoolCount) ++ - range("set", _.getDeleteFilesPoolCount) ++ - range("partition", _.getPartitionDataPoolCount) ++ - range("residual", _.getResidualPoolCount) - } - - /** - * A random scan. Entries are distinct by content, as the driver's interning makes them. Some - * delete files share a path, like deletion vectors in one Puffin file, and some entries are - * used by no task. - */ - private def randomScan(rnd: Random): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { - val paths = 1 + rnd.nextInt(30) - val deletes = 1 + rnd.nextInt(40) - val partitionValues = 1 + rnd.nextInt(10) - val residuals = 1 + rnd.nextInt(5) - val common = withOtherFields(IcebergScanCommon.newBuilder()) - (0 until paths).foreach(i => common.addDeleteFilePathPool(deleteFilePath(i))) - (0 until deletes).foreach { i => - common.addDeleteFilePool(positionDelete(rnd.nextInt(paths), recordCount = i + 1L)) - } - val sets = Seq - .fill(1 + rnd.nextInt(20))(Seq.fill(1 + rnd.nextInt(4))(rnd.nextInt(deletes)).distinct) - .distinct - sets.foreach { set => - val list = DeleteFileList.newBuilder() - set.foreach(d => list.addDeleteFileIndices(d)) - common.addDeleteFilesPool(list) - } - (0 until partitionValues).foreach(i => common.addPartitionDataPool(partitionValue(i))) - (0 until residuals).foreach(i => common.addResidualPool(notNull(s"c$i"))) - - def sometimes(percent: Int, bound: Int): Option[Int] = - if (rnd.nextInt(100) < percent) Some(rnd.nextInt(bound)) else None - var files = 0 - val partitions = (0 until 1 + rnd.nextInt(8)).map { _ => - val scan = IcebergScan.newBuilder() - (0 until rnd.nextInt(6)).foreach { _ => - files += 1 - scan.addFileScanTasks( - fileTask( - s"data-$files.parquet", - sometimes(70, sets.size), - sometimes(80, partitionValues), - sometimes(50, residuals))) - } - scan.build() - } - (common.build(), partitions) - } - - test("every task resolves to the same entries once shipped and injected") { - (0 until 200).foreach { seed => - val (common, partitions) = randomScan(new Random(seed)) - withClue(s"seed $seed: ") { - assertSameTasks(common, partitions) - } - } - } - - test("a partition ships exactly the entries its tasks use, each once") { - (0 until 200).foreach { seed => - val (common, partitions) = randomScan(new Random(seed)) - val (block, shipped) = shipAndInject(common, partitions) - val usedBy = users(common, partitions) - val empty = IcebergScanCommon.getDefaultInstance - // Entries at least half of the non-empty partitions use ship once in the common block. - val nonEmpty = partitions.indices.filter(partitions(_).getFileScanTasksCount > 0).toSet - val half = (nonEmpty.size + 1) / 2 - // toSeq first: Map.collect into (pool, index) tuples would build a Map and merge them. - val everywhere = usedBy.toSeq.collect { - case (e, ps) if nonEmpty.size >= 2 && ps.size >= half => e - }.toSet - withClue(s"seed $seed: ") { - val inCommon = entries(block, empty).map(content(block, _)) - assert(inCommon.size == everywhere.size, "the common block holds an entry twice or extra") - assert(inCommon.toSet == everywhere.map(content(common, _))) - nonEmpty.foreach { p => - val used = usedBy.toSeq.collect { case (e, ps) if ps.contains(p) => e }.toSet - val expected = (used -- everywhere).map(content(common, _)) - // Own entries follow the common ones and may point at them, so resolve them merged. - val injected = shipped(p)._2.getCommon - val sent = entries(injected, block).map(content(injected, _)) - assert(sent.size == expected.size, s"partition $p ships an entry twice or extra") - assert(sent.toSet == expected, s"partition $p ships the wrong entries") - // The plan holds the common entries, which some partitions skip, and its own. - val kept = entries(injected, empty).map(content(injected, _)) - val visible = used ++ everywhere - assert(kept.size == visible.size && kept.toSet == visible.map(content(common, _))) - } - } - } - } - - test("the common block and each partition's plan do not grow with the scan") { - // Partition 0's data is the same at every scan size, so its message and plan must be too. - // Shipping the whole scan's pools in the common block made both grow with the scan. - // From four partitions up, a bucket's own entries are used by fewer than half of them. - val sizes = Seq(4, 100, 10000).map { buckets => - val (common, partitions) = bucketedScan(buckets) - val (block, shipped) = shipAndInject(common, partitions) - ( - common.getSerializedSize, - block.getSerializedSize, - shipped.head._1.getSerializedSize, - shipped.head._2.getSerializedSize) - } - val unsplit = sizes.map(_._1) - assert(unsplit.last > 100 * unsplit.head, s"the test scan does not grow: $unsplit") - assert(sizes.map(_._2).distinct.size == 1, s"common block grows with the scan: $sizes") - assert(sizes.map(_._3).distinct.size == 1, s"partition message grows with the scan: $sizes") - assert(sizes.map(_._4).distinct.size == 1, s"injected plan grows with the scan: $sizes") - } - - test("a partition's message does not grow when partitions share more delete files") { - // Each delete file applies to two neighbouring partitions. Partition 0 must still get only - // the two it uses, however many other partitions there are. - def shippedSize(partitionCount: Int): (Int, Int) = { - val common = IcebergScanCommon.newBuilder().setMetadataLocation(metadataLocation) - (0 until partitionCount).foreach { d => - common.addDeleteFilePathPool(deleteFilePath(d)) - common.addDeleteFilePool(positionDelete(d, recordCount = 1)) - } - val partitions = (0 until partitionCount).map { p => - val set = common.getDeleteFilesPoolCount - common.addDeleteFilesPool( - DeleteFileList - .newBuilder() - .addDeleteFileIndices(p) - .addDeleteFileIndices((p + 1) % partitionCount)) - IcebergScan - .newBuilder() - .addFileScanTasks(fileTask(f"data-$p%08d.parquet", Some(set), None, None)) - .build() - } - val (block, shipped) = shipAndInject(common.build(), partitions) - (block.getSerializedSize, shipped.head._1.getSerializedSize) - } - // From five partitions up, a delete file two neighbours share is used by fewer than half. - val sizes = Seq(5, 100, 10000).map(shippedSize) - assert(sizes.distinct.size == 1, s"shipped sizes grow with shared delete files: $sizes") - } - - test("a delete file every partition uses ships once in the common block") { - // An unpartitioned equality delete applies to every data file. - val global = IcebergDeleteFile - .newBuilder() - .setFilePathIdx(0) - .setContentType("EQUALITY_DELETES") - .setFileFormat("PARQUET") - .addEqualityIds(1) - .setRecordCount(5) - .build() - val common = IcebergScanCommon - .newBuilder() - .setMetadataLocation(metadataLocation) - .addDeleteFilePathPool("s3://bucket/table/data/global-eq-deletes.parquet") - .addDeleteFilePool(global) - .addDeleteFilesPool(DeleteFileList.newBuilder().addDeleteFileIndices(0)) - .build() - val partitions = (0 until 6).map { b => - IcebergScan - .newBuilder() - .addFileScanTasks(fileTask(s"a-$b.parquet", Some(0), None, None)) - .build() - } - val (block, shipped) = shipAndInject(common, partitions) - assert(block.getDeleteFilePoolList.asScala.toSeq == Seq(global)) - assert(block.getDeleteFilePathPoolCount == 1 && block.getDeleteFilesPoolCount == 1) - shipped.foreach { case (payload, _) => assert(!payload.hasOwnPools) } - assertSameTasks(common, partitions) - } - - test("a partition with no tasks ships an empty message") { - val (common, partitions) = bucketedScan(3) - val withEmpty = partitions :+ IcebergScan.getDefaultInstance - val pools = new IcebergScanPools.MessagePools(common) - val payload = IcebergScanPools.partitionPayload( - pools, - IcebergScanPools.Layout(pools, withEmpty), - IcebergScan.getDefaultInstance) - assert(!payload.hasOwnPools) - assert(payload.getSerializedSize == 0) - assertSameTasks(common, withEmpty) - } - - test("an index past the end of its pool fails with the pool's name") { - val (common, partitions) = bucketedScan(2) - val bad = partitions.head.toBuilder - .setFileScanTasks(0, partitions.head.getFileScanTasks(0).toBuilder.setDeleteFilesIdx(99)) - .build() - val pools = new IcebergScanPools.MessagePools(common) - val layout = IcebergScanPools.Layout(pools, partitions) - val error = intercept[IllegalStateException] { - IcebergScanPools.partitionPayload(pools, layout, bad) - } - assert(error.getMessage.contains("delete_files_set index 99 is out of range")) - } - - // The common block with the five per-file pools emptied: everything shipping must leave alone. - private def otherFields(common: IcebergScanCommon): IcebergScanCommon = - common.toBuilder - .clearDeleteFilePathPool() - .clearDeleteFilePool() - .clearDeleteFilesPool() - .clearPartitionDataPool() - .clearResidualPool() - .build() - - test("shipping keeps every field outside the five pools") { - (0 until 50).foreach { seed => - val (common, partitions) = randomScan(new Random(seed)) - val (block, shipped) = shipAndInject(common, partitions) - withClue(s"seed $seed: ") { - assert(otherFields(block) == otherFields(common)) - shipped.foreach { case (_, injected) => - assert(otherFields(injected.getCommon) == otherFields(common)) - } - } - } - } - - test("every pool index field is one the pruning knows about") { - // A new index into a pool must be handled by IcebergScanPools, or tasks keep indices into - // pools that no longer match. This fails when one is added without updating it. - def indexFields(descriptor: com.google.protobuf.Descriptors.Descriptor): Set[String] = - descriptor.getFields.asScala - .map(_.getName) - .filter(n => n.endsWith("_idx") || n.endsWith("_indices")) - .toSet - assert( - indexFields(IcebergFileScanTask.getDescriptor) == Set( - "schema_idx", - "partition_spec_idx", - "name_mapping_idx", - "project_field_ids_idx", - "partition_data_idx", - "delete_files_idx", - "residual_idx")) - assert(indexFields(IcebergDeleteFile.getDescriptor) == Set("file_path_idx")) - assert(indexFields(DeleteFileList.getDescriptor) == Set("delete_file_indices")) - // A per-partition field added to IcebergScan must be carried by partitionPayload. - assert( - IcebergScan.getDescriptor.getFields.asScala.map(_.getName).toSet == - Set("common", "file_scan_tasks", "own_pools")) - // A partition's own pools must be exactly the common block's per-file pools. - assert( - IcebergFilePools.getDescriptor.getFields.asScala.map(_.getName).toSet == Set( - "partition_data_pool", - "delete_files_pool", - "residual_pool", - "delete_file_pool", - "delete_file_path_pool")) - assert( - IcebergScanCommon.getDescriptor.getFields.asScala - .map(_.getName) - .filter(_.endsWith("_pool")) - .toSet == - Set( - "schema_pool", - "partition_type_pool", - "partition_spec_pool", - "name_mapping_pool", - "project_field_ids_pool", - "partition_data_pool", - "delete_files_pool", - "residual_pool", - "delete_file_pool", - "delete_file_path_pool")) - } - - test("lazily built scan data compares by content without building partitions") { - val (common, partitions) = bucketedScan(3) - def data(c: IcebergScanCommon) = { - val pools = new IcebergScanPools.SerializedPools(c.toByteArray) - new IcebergPartitionPlanData( - pools, - IcebergScanPools.Layout(pools, partitions), - partitions.map(_.toByteArray).toArray) - } - assert(data(common).sameAs(data(common))) - assert(!data(common).sameAs(data(common.toBuilder.setScanHashCode(7).build()))) - } - - test("reading the serialized pools gives the same entries as the parsed message") { - (0 until 50).foreach { seed => - val (common, _) = randomScan(new Random(seed)) - val parsed = new IcebergScanPools.MessagePools(common) - val serialized = new IcebergScanPools.SerializedPools(common.toByteArray) - withClue(s"seed $seed: ") { - assert(serialized.pathCount == parsed.pathCount) - assert((0 until parsed.pathCount).forall(i => serialized.path(i) == parsed.path(i))) - assert(serialized.deleteCount == parsed.deleteCount) - assert((0 until parsed.deleteCount).forall(i => serialized.delete(i) == parsed.delete(i))) - assert(serialized.setCount == parsed.setCount) - assert((0 until parsed.setCount).forall(i => serialized.set(i) == parsed.set(i))) - assert(serialized.partitionValueCount == parsed.partitionValueCount) - assert((0 until parsed.partitionValueCount).forall(i => - serialized.partitionValue(i) == parsed.partitionValue(i))) - assert(serialized.residualCount == parsed.residualCount) - assert((0 until parsed.residualCount).forall(i => - serialized.residual(i) == parsed.residual(i))) - } - } - } - - test("a delete file most partitions use ships once, one fewer than half use ships per task") { - // Newer appends that a delete does not apply to, or a bucket with no deletes, leave some - // partitions without it. It still ships once while at least half of the partitions use it. - def layout(users: Int): (IcebergScanCommon, IndexedSeq[(IcebergScan, IcebergScan)]) = { - val (common, partitions) = sharedDeleteScan(8, users) - assertSameTasks(common, partitions) - shipAndInject(common, partitions) - } - val (most, mostShipped) = layout(users = 7) - assert(most.getDeleteFilePoolCount == 1) - assert(mostShipped.forall(!_._1.hasOwnPools)) - val (half, _) = layout(users = 4) - assert(half.getDeleteFilePoolCount == 1) - val (few, fewShipped) = layout(users = 3) - assert(few.getDeleteFilePoolCount == 0) - assert(fewShipped.take(3).forall(_._1.getOwnPools.getDeleteFilePoolCount == 1)) - assert(fewShipped.drop(3).forall(!_._1.hasOwnPools)) - } -} diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala b/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala deleted file mode 100644 index 308cfaa0157..00000000000 --- a/spark/src/test/scala/org/apache/comet/serde/operator/IcebergScanPoolsTestUtils.scala +++ /dev/null @@ -1,175 +0,0 @@ -/* - * 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.serde.operator - -import org.apache.spark.sql.comet.CometIcebergNativeScanExec - -import org.apache.comet.serde.OperatorOuterClass.{DeleteFileList, IcebergDeleteFile, IcebergFilePools, IcebergFileScanTask, IcebergLiteral, IcebergPredicate, IcebergPredicateOperator, IcebergScan, IcebergScanCommon, IcebergUnaryPredicate, PartitionData, PartitionValue, ProjectFieldIdList, SparkStructField} - -/** Builders and views for tests of how a native Iceberg scan splits its pools. */ -object IcebergScanPoolsTestUtils { - - val metadataLocation = "s3://bucket/table/metadata/v1.metadata.json" - - def deleteFilePath(i: Int): String = f"s3://bucket/table/data/part-$i%08d-deletes.parquet" - - def positionDelete(pathIdx: Int, recordCount: Long): IcebergDeleteFile = - IcebergDeleteFile - .newBuilder() - .setFilePathIdx(pathIdx) - .setContentType("POSITION_DELETES") - .setFileFormat("PARQUET") - .setRecordCount(recordCount) - .build() - - def partitionValue(value: Int): PartitionData = - PartitionData - .newBuilder() - .addValues( - PartitionValue - .newBuilder() - .setFieldId(1000) - .setLiteral(IcebergLiteral.newBuilder().setIntVal(value))) - .build() - - def notNull(column: String): IcebergPredicate = - IcebergPredicate - .newBuilder() - .setUnary( - IcebergUnaryPredicate - .newBuilder() - .setColumn(column) - .setOp(IcebergPredicateOperator.NotNull)) - .build() - - def fileTask( - dataFile: String, - deleteSet: Option[Int], - partitionData: Option[Int], - residual: Option[Int]): IcebergFileScanTask = { - val task = IcebergFileScanTask - .newBuilder() - .setDataFilePath(dataFile) - .setLength(1024) - .setFileSizeInBytes(1024) - deleteSet.foreach(task.setDeleteFilesIdx) - partitionData.foreach(task.setPartitionDataIdx) - residual.foreach(task.setResidualIdx) - task.build() - } - - /** Sets the common-block fields the split leaves alone, so tests can check they survive. */ - def withOtherFields(common: IcebergScanCommon.Builder): IcebergScanCommon.Builder = - common - .setMetadataLocation(metadataLocation) - .setScanHashCode(42) - .setCatalogName("test_cat") - .setDataFileConcurrencyLimit(4) - .putCatalogProperties("s3.region", "us-west-2") - .addRequiredSchema(SparkStructField.newBuilder().setName("id").setNullable(true)) - .addSchemaPool("{\"type\":\"struct\",\"schema-id\":0,\"fields\":[]}") - .addPartitionSpecPool("{\"spec-id\":0,\"fields\":[]}") - .addPartitionTypePool("{\"type\":\"struct\",\"fields\":[]}") - .addNameMappingPool("[]") - .addProjectFieldIdsPool(ProjectFieldIdList.newBuilder().addFieldIds(1)) - - /** - * A scan like a bucketed merge-on-read table read one bucket per partition. Partition b holds - * four data files, three delete files only it uses, and its own partition value. Every task - * carries the same residual. - */ - def bucketedScan(buckets: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { - val common = IcebergScanCommon - .newBuilder() - .setMetadataLocation(metadataLocation) - .addResidualPool(notNull("payload")) - val partitions = (0 until buckets).map { b => - val set = DeleteFileList.newBuilder() - (0 until 3).foreach { d => - val idx = common.getDeleteFilePoolCount - common.addDeleteFilePathPool(deleteFilePath(b * 3 + d)) - common.addDeleteFilePool(positionDelete(idx, recordCount = 10)) - set.addDeleteFileIndices(idx) - } - val setIdx = common.getDeleteFilesPoolCount - common.addDeleteFilesPool(set) - val partitionIdx = common.getPartitionDataPoolCount - common.addPartitionDataPool(partitionValue(b)) - val scan = IcebergScan.newBuilder() - (0 until 4).foreach { f => - scan.addFileScanTasks( - fileTask(f"data-$b%08d-$f.parquet", Some(setIdx), Some(partitionIdx), Some(0))) - } - scan.build() - } - (common.build(), partitions) - } - - /** - * A scan where each partition reads one data file, and the first `users` of them share one - * delete file. Whether it ships in the common block depends on how many partitions use it. - */ - def sharedDeleteScan( - partitionCount: Int, - users: Int): (IcebergScanCommon, IndexedSeq[IcebergScan]) = { - val common = IcebergScanCommon - .newBuilder() - .setMetadataLocation(metadataLocation) - .addDeleteFilePathPool(deleteFilePath(0)) - .addDeleteFilePool(positionDelete(0, recordCount = 100)) - .addDeleteFilesPool(DeleteFileList.newBuilder().addDeleteFileIndices(0)) - .build() - val partitions = (0 until partitionCount).map { p => - val deletes = if (p < users) Some(0) else None - IcebergScan - .newBuilder() - .addFileScanTasks(fileTask(f"data-$p%08d.parquet", deletes, None, None)) - .build() - } - (common, partitions) - } - - /** A partition's own pools in a common block, so tests can read both the same way. */ - def asCommon(pools: IcebergFilePools): IcebergScanCommon = - IcebergScanCommon - .newBuilder() - .addAllDeleteFilePathPool(pools.getDeleteFilePathPoolList) - .addAllDeleteFilePool(pools.getDeleteFilePoolList) - .addAllDeleteFilesPool(pools.getDeleteFilesPoolList) - .addAllPartitionDataPool(pools.getPartitionDataPoolList) - .addAllResidualPool(pools.getResidualPoolList) - .build() - - /** 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) +: - scan - .buildAllPartitionDataForTesting() - .toSeq - .map(bytes => asCommon(IcebergScan.parseFrom(bytes).getOwnPools)) - - /** The scan each partition's native plan runs, as the injector builds it. */ - def partitionScans(scan: CometIcebergNativeScanExec): Seq[IcebergScan] = { - val common = IcebergScanCommon.parseFrom(scan.commonData) - scan.buildAllPartitionDataForTesting().toSeq.map { bytes => - IcebergScanPools.forPartition(common, IcebergScan.parseFrom(bytes)) - } - } -} 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 af78973665a..9bf49125d7d 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,8 +120,7 @@ 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").toArray.map(_.toSeq).toSeq == perPart.map(_.toSeq).toSeq) + assert(perPartitionByKey("contrib-key").map(_.toSeq).toSeq == perPart.map(_.toSeq).toSeq) assert( scan.subqueriesResolved, "findAllPlanData must drive the scan's subquery-resolution lifecycle before reading data") @@ -146,9 +145,7 @@ class CometScanWithPlanDataSuite extends AnyFunSuite { assert(commonByKey.keySet == Set("shared-key")) assert(commonByKey("shared-key") sameElements Array[Byte](1, 2)) assert( - perPartitionByKey("shared-key").toArray.map(_.toSeq).toSeq == Seq( - Seq[Byte](10), - Seq[Byte](20))) + perPartitionByKey("shared-key").map(_.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 4ab3ee594c3..c805fdf9374 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 @@ -28,7 +28,6 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager import org.apache.comet.serde.OperatorOuterClass import org.apache.comet.serde.OperatorOuterClass.Operator -import org.apache.comet.serde.operator.{IcebergScanPools, IcebergScanPoolsTestUtils} class PlanDataInjectorSuite extends AnyFunSuite { @@ -129,6 +128,84 @@ 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") + } + 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 @@ -168,102 +245,6 @@ class PlanDataInjectorSuite extends AnyFunSuite { assert(IcebergPlanDataInjector.getKey(opA) != IcebergPlanDataInjector.getKey(opB)) } - private def layoutOf( - common: OperatorOuterClass.IcebergScanCommon, - partitions: IndexedSeq[OperatorOuterClass.IcebergScan]): IcebergScanPools.Layout = - IcebergScanPools.Layout(new IcebergScanPools.MessagePools(common), partitions) - - test("an injected Iceberg plan does not grow with the rest of the scan") { - // Partition 0's data is the same at every scan size, while every other partition adds its - // own delete files. Tasks allocate the injected plan in serializeOperator, so its size must - // not follow the whole scan. - val scanOp = icebergScanOp(IcebergScanPoolsTestUtils.metadataLocation, scanHashCode = 7) - val key = IcebergPlanDataInjector.getKey(scanOp).get - // From four partitions up, a bucket's own entries are used by fewer than half of them. - val injected = Seq(4, 100, 10000).map { buckets => - val (common, partitions) = IcebergScanPoolsTestUtils.bucketedScan(buckets) - PlanDataInjector.injectPlanData( - parseBasePlan(scanOp.toByteArray), - Map( - key -> IcebergScanPools.commonBlock(common, layoutOf(common, partitions)).toByteArray), - Map( - key -> IcebergScanPools - .partitionPayload( - new IcebergScanPools.MessagePools(common), - layoutOf(common, partitions), - partitions.head) - .toByteArray)) - } - val sizes = injected.map(PlanDataInjector.serializeOperator(_).length) - assert(sizes.distinct.size == 1, s"injected plan grows with the scan: $sizes") - - // The partition's own delete files are reachable from its tasks after injection. - val scan = injected.last.getIcebergScan - val pools = scan.getCommon - val deletes = pools.getDeleteFilesPool(scan.getFileScanTasks(0).getDeleteFilesIdx) - val paths = deletes.getDeleteFileIndicesList.asScala.map { d => - pools.getDeleteFilePathPool(pools.getDeleteFilePool(d).getFilePathIdx) - } - assert(paths == (0 until 3).map(IcebergScanPoolsTestUtils.deleteFilePath)) - } - - test("a coalesced task carries a copy of a shared delete file per input partition") { - // Pins a documented limit: without a shuffle, a coalesced task serializes all its input - // partitions, and each brings its own copy of the delete files it uses. Spark's Iceberg - // reader writes a shared delete file once per task. Update this test if that changes. - // Three of eight partitions use the delete file: too few for the common block. - val (common, partitions) = IcebergScanPoolsTestUtils.sharedDeleteScan(8, users = 3) - val layout = layoutOf(common, partitions) - val data = new PartitionPlanData { - override def length: Int = partitions.length - override def apply(i: Int): Array[Byte] = - IcebergScanPools - .partitionPayload(new IcebergScanPools.MessagePools(common), layout, partitions(i)) - .toByteArray - } - def taskSize(inputs: Int): Int = { - val parents = (0 until inputs).map { i => - new CometExecPartition(i, Array.empty, new PlanDataSlice(Map("k" -> data), i)) - }.toArray - val bytes = new java.io.ByteArrayOutputStream() - val out = new java.io.ObjectOutputStream(bytes) - out.writeObject(parents) - out.close() - bytes.size() - } - val perPartition = taskSize(2) - taskSize(1) - assert(taskSize(3) - taskSize(1) == 2 * perPartition, "each input partition adds the same") - assert( - perPartition > data(0).length, - "each input partition carries its own copy of the shared delete file") - } - - test("a task's plan data is built only when its partition is serialized") { - // The driver holds every partition object for the job, so a partition must not keep its - // bytes. They are built while Spark serializes the task, and the executor gets them intact. - val built = new java.util.concurrent.atomic.AtomicInteger() - val data = new PartitionPlanData { - override def length: Int = 3 - override def apply(i: Int): Array[Byte] = { - built.incrementAndGet() - Array.fill(4)(i.toByte) - } - } - val partition = new CometExecPartition(1, Array.empty, new PlanDataSlice(Map("k" -> data), 1)) - assert(built.get == 0, "plan data built before the task was serialized") - - val bytes = new java.io.ByteArrayOutputStream() - val out = new java.io.ObjectOutputStream(bytes) - out.writeObject(partition) - out.close() - assert(built.get == 1) - - val in = new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(bytes.toByteArray)) - val received = in.readObject().asInstanceOf[CometExecPartition] - assert(received.planDataByKey("k").toSeq == Seq.fill(4)(1.toByte)) - assert(built.get == 1, "the executor side rebuilt the plan data") - } - /** * Builds an un-injected NativeScan operator the way the driver ships it: hasCommon, no * file_partition, source_key_hash embedded (see CometNativeScanExec.apply). 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 c0ca8290715..6ca87258b92 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, CometExecRDD, CometMetricNode, NativeExecContext, PartitionPlanData, PlanDataSlice} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometMetricNode, NativeExecContext} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -119,7 +119,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 2, Set.empty, CometMetricNode(Map.empty), - planData.map { case (key, bytes) => key -> PartitionPlanData(bytes) }) + planData) val local = remote.copyForLocalShuffle() assert(local.id != remote.id) @@ -207,7 +207,7 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { CometNativeShuffleInputRDD, CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]) = { val perPartitionByKey = - Map("scan-0" -> PartitionPlanData(Array.fill(numPartitions)(new Array[Byte](1024)))) + Map("scan-0" -> Array.fill(numPartitions)(new Array[Byte](1024))) val childMetricNode = CometMetricNode(Map.empty) val writerMetrics = Map( "spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "disk spilled bytes"), @@ -263,55 +263,4 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { assert(part.planDataByKey.keySet == Set("scan-0")) assert(part.planDataByKey("scan-0").length == 1024) } - - test("a shuffle map task's plan data is built only when its partition is serialized") { - val built = new java.util.concurrent.atomic.AtomicInteger() - val data = new PartitionPlanData { - override def length: Int = 2 - override def apply(i: Int): Array[Byte] = { - built.incrementAndGet() - Array.fill(3)(i.toByte) - } - } - val partition = - new CometNativeShuffleInputPartition(1, Array.empty, new PlanDataSlice(Map("k" -> data), 1)) - assert(built.get == 0, "plan data built before the task was serialized") - - val bytes = new java.io.ByteArrayOutputStream() - val out = new java.io.ObjectOutputStream(bytes) - out.writeObject(partition) - out.close() - assert(built.get == 1) - - val in = new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(bytes.toByteArray)) - val received = in.readObject().asInstanceOf[CometNativeShuffleInputPartition] - assert(received.planDataByKey("k").toSeq == Seq.fill(3)(1.toByte)) - assert(built.get == 1, "the executor side rebuilt the plan data") - } - - test("an RDD's partitions don't build any partition's plan data") { - // Spark keeps every partition object for the job's lifetime, so building the partitions must - // not build the plan data. Only serializing a task does. - val built = new java.util.concurrent.atomic.AtomicInteger() - val data = new PartitionPlanData { - override def length: Int = 1000 - override def apply(i: Int): Array[Byte] = { - built.incrementAndGet() - Array(i.toByte) - } - } - val rdd = CometExecRDD( - spark.sparkContext, - inputRDDs = Seq.empty, - commonByKey = Map("k" -> Array[Byte](1)), - perPartitionByKey = Map("k" -> data), - serializedPlan = Array.empty, - planFingerprint = 0L, - numPartitions = 1000, - numOutputCols = 1, - nativeMetrics = CometMetricNode(Map.empty), - subqueries = Seq.empty) - assert(rdd.partitions.length == 1000) - assert(built.get == 0, s"building the partitions built ${built.get} partitions' plan data") - } } From 44cf0855ce5e03a207675f152f8fabfa60b4f319 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Thu, 5 Nov 2026 13:46:17 -0800 Subject: [PATCH 6/8] docs: note that the driver keeps a copy of a shared delete file per partition --- docs/source/user-guide/latest/iceberg.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 8e241fe0a3f..8d98b02ab11 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -204,6 +204,13 @@ uses it. Scans where many partitions share delete files, such as tables with unp equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning instead of coalescing, avoids this. +The driver builds every task's plan data before the scan starts, and keeps it until the query +ends. It keeps one copy of a shared delete file for each partition that uses it. A table with +many unpartitioned equality deletes, read as many partitions, can therefore use a lot of driver +memory. Compact the delete files before reading such a table natively. Building plan data only +when each task is sent is tracked in +[#7523](https://github.com/apache/datafusion-comet/issues/7523). + ### Current limitations The following scenarios will fall back to the JVM Iceberg reader: From 19c4f141d8276ced0ac17676aa4ba09cd87aa436 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Thu, 5 Nov 2026 15:57:53 -0800 Subject: [PATCH 7/8] fix: share Iceberg delete files between partitions so a coalesced task carries each once Review feedback: a coalesce without a shuffle writes several input partitions into one task, and each partition held its own copy of a shared delete file, so the task carried one copy per input partition where the latest release carried one. Per-partition plan data is now a PartitionPlanData object instead of bytes. An Iceberg partition holds its own message plus references to delete files serialized once for the whole scan, so Java serialization writes a shared delete file once per task stream and the driver keeps one copy. The executor sets each delete file's path index when it builds the partition's message, since the path pool order differs per partition. Native Parquet and contrib scans wrap their bytes, so the CometScanWithPlanData trait is unchanged. Move the Iceberg pool test helpers into IcebergScanPlanData, with distinctResiduals for tests that read the residual pool. --- .../sql/comet/CometDeltaNativeScanExec.scala | 2 +- docs/source/user-guide/latest/iceberg.md | 18 +-- native/proto/src/proto/operator.proto | 6 +- .../operator/CometIcebergNativeScan.scala | 145 +++++++++--------- .../apache/spark/sql/comet/CometExecRDD.scala | 10 +- .../comet/CometIcebergNativeScanExec.scala | 4 +- .../spark/sql/comet/CometNativeScanExec.scala | 2 +- .../spark/sql/comet/PartitionPlanData.scala | 104 +++++++++++++ .../shuffle/CometNativeShuffleInputRDD.scala | 8 +- .../shuffle/CometNativeShuffleWriter.scala | 2 +- .../apache/spark/sql/comet/operators.scala | 13 +- .../apache/comet/CometFuzzIcebergSuite.scala | 15 +- .../comet/CometIcebergNativeSuite.scala | 80 ++++------ .../apache/comet/IcebergScanPlanData.scala | 55 +++++++ .../CometIcebergNativeScanSuite.scala | 12 +- .../comet/CometScanWithPlanDataSuite.scala | 7 +- .../sql/comet/PlanDataInjectorSuite.scala | 138 +++++++++++++++++ .../CometNativeShuffleInputRDDSuite.scala | 13 +- 18 files changed, 455 insertions(+), 179 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala create mode 100644 spark/src/test/scala/org/apache/comet/IcebergScanPlanData.scala 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 8d98b02ab11..d64492c1658 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -197,19 +197,11 @@ These `s3.*` storage properties are not specific to the Hive catalog shown here. ### 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 coalesce without a shuffle -(`coalesce(n)`, or a write that coalesces its input) combines several input partitions into one -task, and that task then carries a copy of a shared delete file for every input partition that -uses it. Scans where many partitions share delete files, such as tables with unpartitioned -equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning -instead of coalescing, avoids this. - -The driver builds every task's plan data before the scan starts, and keeps it until the query -ends. It keeps one copy of a shared delete file for each partition that uses it. A table with -many unpartitioned equality deletes, read as many partitions, can therefore use a lot of driver -memory. Compact the delete files before reading such a table natively. Building plan data only -when each task is sent is tracked in -[#7523](https://github.com/apache/datafusion-comet/issues/7523). +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 diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index d8dc23565fc..7761ff838bf 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -390,8 +390,10 @@ message IcebergScanCommon { // 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 ships the entries its tasks use in its own - // IcebergScan.common, and the JVM merges the two before the plan reaches native. + // 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 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 3d9ee41060d..e9a0779330a 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()) } /** @@ -959,14 +953,16 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit * * 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, into that - * partition's IcebergScan.common, 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. + * 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 each partition stores it once and references it many times, as Iceberg's own - * DeleteFileIndex does. + * 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 @@ -975,12 +971,12 @@ 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() @@ -988,44 +984,61 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val partitionSpecToPoolIndex = mutable.HashMap[String, Int]() val nameMappingToPoolIndex = mutable.HashMap[String, Int]() - // The per-file pools of one partition. Each partition ships the entries its own tasks use in - // its IcebergScan.common, and the injector merges them into the common block, so a task never - // holds entries that only other partitions use. Within a partition each entry is stored once. + // 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]() - // Individual delete files are interned into a flat pool; deleteFiles 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 deleteFiles = mutable.HashMap[OperatorOuterClass.IcebergDeleteFile, 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]() - // 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 deleteFilePaths = mutable.HashMap[String, Int]() val residuals = mutable.HashMap[OperatorOuterClass.IcebergPredicate, Int]() - def internDeleteFilePath(path: String): Int = - deleteFilePaths.getOrElseUpdate( - path, { - val idx = deleteFilePaths.size - builder.addDeleteFilePathPool(path) - idx + 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 isEmpty: Boolean = - partitionData.isEmpty && deleteFiles.isEmpty && deleteFilePaths.isEmpty && - residuals.isEmpty + 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 deleteFilePathEntries = 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 @@ -1060,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 @@ -1174,19 +1187,11 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit task, contentFileClass, fileScanTaskClass, - deleteFileClass, - pools.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 => - pools.deleteFiles.getOrElseUpdate( - df, { - val idx = pools.deleteFiles.size - pools.builder.addDeleteFilePool(df) - idx - }) - } + // 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 = pools.deleteFileSets.size @@ -1239,13 +1244,11 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } - if (!pools.isEmpty) partitionBuilder.setCommon(pools.builder) partitionDataEntries += pools.partitionData.size deleteFileEntries += pools.deleteFiles.size deleteFileSetEntries += pools.deleteFileSets.size - deleteFilePathEntries += pools.deleteFilePaths.size residualEntries += pools.residuals.size - perPartitionBuilders += partitionBuilder.build() + 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 @@ -1264,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( @@ -1281,7 +1285,6 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit partitionDataEntries, deleteFileEntries, deleteFileSetEntries, - deleteFilePathEntries, residualEntries) val avgDedup = if (totalTasks == 0) { @@ -1326,7 +1329,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit sumSizes( 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 @@ -1340,16 +1343,16 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit logInfo( " Per-partition pool entries, summed over partitions: " + s"partition_data=$partitionDataEntries, delete_file=$deleteFileEntries, " + - s"delete_file_path=$deleteFilePathEntries, delete_files_set=$deleteFileSetEntries, " + - s"residual=$residualEntries") + s"delete_files_set=$deleteFileSetEntries, residual=$residualEntries") + logInfo( + s" 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: ${perPartitionBuilders.size}, " + + 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 43e66a36a2b..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 @@ -894,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/scala/org/apache/comet/CometFuzzIcebergSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala index 517ea6f27a9..ef1f0d40253 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzIcebergSuite.scala @@ -22,13 +22,12 @@ package org.apache.comet import scala.util.Random import org.apache.spark.sql.Column -import org.apache.spark.sql.comet.CometIcebergNativeScanExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec 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,12 +176,6 @@ class CometFuzzIcebergSuite extends CometFuzzIcebergBase { } } - /** The pool entries a scan ships: the common block's, then each partition's own. */ - private def storedPools( - scan: CometIcebergNativeScanExec): Seq[OperatorOuterClass.IcebergScanCommon] = - OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) +: - scan.perPartitionData.toSeq.map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon) - // Asserts that predicate pushdown actually happens for the types that support it. For each // 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 @@ -233,7 +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 pushed = storedPools(scans.head).exists(_.getResidualPoolCount > 0) + val pushed = distinctResiduals(scans.head).nonEmpty if (isPushable(field.dataType)) { assert( pushed, @@ -259,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 residuals = storedPools(scans.head).map(_.getResidualPoolCount).sum - assert(residuals == 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 e1f70150d7d..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 @@ -144,23 +145,6 @@ class CometIcebergNativeSuite s"Expected fallback to Spark but found a CometIcebergNativeScanExec. Plan:\n$cometPlan") } - /** The pool entries a scan ships: the common block's, then each partition's own. */ - private def storedPools( - scan: CometIcebergNativeScanExec): Seq[OperatorOuterClass.IcebergScanCommon] = - OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) +: - scan.perPartitionData.toSeq.map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon) - - /** The pools each partition's tasks index into once the injector merges them. */ - private def partitionPools( - scan: CometIcebergNativeScanExec): Seq[OperatorOuterClass.IcebergScanCommon] = { - val common = OperatorOuterClass.IcebergScanCommon.parseFrom(scan.commonData) - scan.perPartitionData.toSeq.map { bytes => - common.toBuilder - .mergeFrom(OperatorOuterClass.IcebergScan.parseFrom(bytes).getCommon) - .build() - } - } - /** Counts non-overlapping occurrences of `needle` within `haystack`. */ private def countByteOccurrences(haystack: Array[Byte], needle: Array[Byte]): Int = { require(needle.nonEmpty, "needle must be non-empty") @@ -830,7 +814,7 @@ class CometIcebergNativeSuite s"and resolved paths $resolvedPaths") val occurrences = scan.perPartitionData.toSeq - .map(countByteOccurrences(_, pathPool.head.getBytes(UTF_8))) + .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( @@ -904,9 +888,7 @@ class CometIcebergNativeSuite assert( common.getDeleteFilePathPoolCount == 0 && common.getDeleteFilePoolCount == 0, s"expected no delete files in the common block: $common") - val own = - scan.perPartitionData.toSeq - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon) + 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) @@ -918,6 +900,14 @@ class CometIcebergNativeSuite .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") } @@ -996,7 +986,7 @@ class CometIcebergNativeSuite // 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 + val planData = scans.head.commonData +: scans.head.perPartitionData.toSeq.map(_.bytes) val duplicated = for { p <- distinctPaths (bytes, i) <- planData.zipWithIndex @@ -1090,8 +1080,7 @@ class CometIcebergNativeSuite s"per-file entries in the common block: $common") } def ownDeletes(scan: CometIcebergNativeScanExec): Seq[Int] = - scan.perPartitionData.toSeq - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon.getDeleteFilePoolCount) + partitionScans(scan).map(_.getCommon.getDeleteFilePoolCount) // One bucket per partition: each delete file goes to exactly one partition. val grouped = scanOf(grouped = true) @@ -1420,8 +1409,7 @@ class CometIcebergNativeSuite val (_, cometPlan) = checkSparkAnswerAndOperator("SELECT * FROM test_cat.db.global_eq_delete ORDER BY id") val scan = collectIcebergNativeScans(cometPlan).head - val own = scan.perPartitionData.toSeq - .map(OperatorOuterClass.IcebergScan.parseFrom(_).getCommon) + 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") @@ -1441,6 +1429,16 @@ class CometIcebergNativeSuite 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") } @@ -5167,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") @@ -5303,11 +5299,8 @@ class CometIcebergNativeSuite // 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 tasksWithResidualPerPartition = - scans.head.perPartitionData.toSeq.map { bytes => - OperatorOuterClass.IcebergScan - .parseFrom(bytes) - .getFileScanTasksList - .asScala + partitionScans(scans.head).map { partition => + partition.getFileScanTasksList.asScala .count(_.hasResidualIdx) } val tasksWithResidual = tasksWithResidualPerPartition.sum @@ -5975,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, @@ -6244,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 @@ -6281,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") } @@ -6319,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 => @@ -6785,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 c805fdf9374..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 @@ -206,6 +206,144 @@ class PlanDataInjectorSuite extends AnyFunSuite { 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) } } From d7aafc343751aca873053bb93e74b08686009aac Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Thu, 5 Nov 2026 20:03:05 -0800 Subject: [PATCH 8/8] style: drop an interpolator with no placeholders --- .../apache/comet/serde/operator/CometIcebergNativeScan.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e9a0779330a..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 @@ -1345,7 +1345,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit s"partition_data=$partitionDataEntries, delete_file=$deleteFileEntries, " + s"delete_files_set=$deleteFileSetEntries, residual=$residualEntries") logInfo( - s" Delete files, each held once and shared by the partitions that use it: " + + " Delete files, each held once and shared by the partitions that use it: " + s"${sharedDeleteFiles.size} files, ${sharedDeleteFilePaths.size} paths, " + s"$sharedDeleteFileBytes bytes") logInfo(