mbutrovich commented on code in PR #7491:
URL: https://github.com/apache/datafusion-comet/pull/7491#discussion_r4483182458
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala:
##########
@@ -1265,76 +1269,81 @@ 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()
+ // 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)
Review Comment:
Could we build each partition's pools in the existing pass instead of
building global pools and splitting them afterwards? The serde already interns
partition values, delete files, delete-file sets, paths, and residuals with the
maps at [lines
989-1010](https://github.com/apache/datafusion-comet/blob/b1fc2b80b6241fe589a321d47e0b77f67c7aeadf/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala#L989-L1010).
`IcebergScanPools.partitionPayload` then interns the same set, delete file,
and path chain a second time with five more maps ([`IcebergScanPools.scala`
lines
295-391](https://github.com/apache/datafusion-comet/blob/b1fc2b80b6241fe589a321d47e0b77f67c7aeadf/spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala#L295-L391)),
and `SerializedPools` exists so the driver can read entries back out of the
global pools. If those five maps are created inside the [`partitions.foreach`
loop](https://github.com/apache/datafusion-comet/blob/b1fc
2b80b6241fe589a321d47e0b77f67c7aeadf/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala#L1095)
and write into that partition's builder, as I read it each partition's message
falls out of the single pass with indices that start at 0. `Layout`,
`commonBlock`, `partitionPayload`, and `SerializedPools` would then have
nothing left to do.
Most of `IcebergScanPoolsSuite` and `IcebergScanPoolsTestUtils` test those
structures. The regression the issue reports is covered by `an injected Iceberg
plan does not grow with the rest of the scan` in `PlanDataInjectorSuite`, and
per-block dedup is covered by the `countByteOccurrences` checks in
`CometIcebergNativeSuite` that this PR already adapts. Could the coverage live
in those suites?
##########
native/proto/src/proto/operator.proto:
##########
@@ -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;
}
Review Comment:
Could we reuse `IcebergScanCommon` for a partition's own entries instead of
adding `IcebergFilePools` and `own_pools`? The per-partition `IcebergScan`
already has a `common` field, and the driver leaves it unset in the message it
ships. If the driver puts the partition's five pools there, the executor side
of `IcebergPlanDataInjector.inject` becomes
`common.toBuilder.mergeFrom(partition.getCommon)`, because protobuf-java's
[`mergeFrom`](https://github.com/protocolbuffers/protobuf/blob/9d0ec0f92b5b5fdeeda11f9dcecc1872ff378014/java/core/src/main/java/com/google/protobuf/Message.java#L94-L95)
concatenates repeated fields. That would replace `forPartition`, the
`FilePools`/`MessagePools` types that mirror the five fields, and the new
message type. Together with the next comment, it would also remove the need for
the `every pool index field is one the pruning knows about` guard test, since
nothing would renumber indices after the fact.
##########
spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala:
##########
@@ -0,0 +1,408 @@
+/*
+ * 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
Review Comment:
What does keeping entries used by at least half of the partitions in the
common block buy at the executor? Every task deserializes its own copy of the
common block, because `commonByKey` is a field of `CometExecRDD` ([line
63](https://github.com/apache/datafusion-comet/blob/b1fc2b80b6241fe589a321d47e0b77f67c7aeadf/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala#L63),
read in `compute` at [line
141](https://github.com/apache/datafusion-comet/blob/b1fc2b80b6241fe589a321d47e0b77f67c7aeadf/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala#L141))
and so rides in the task binary. As I read it, an entry in the common block
costs every task, including the up to half that don't use it, while the same
entry in a partition's own pools costs only the tasks that use it. So the
threshold doesn't lower any task's heap. It saves network bytes and driver
serialization for widely shared entries, such as a global equality delete.
Spark's reader pays that cost too,
since it ships each task's delete files with the task. Was there a workload
where those bytes mattered? If not, could we drop the threshold so that the
common block holds none of the five pools?
##########
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]]]
+ }
+}
Review Comment:
Should the lazy driver-side build go in its own PR? #7490 reports executor
OOM. `PartitionPlanData`, `PlanDataSlice`, `SerializedPools`, and
`IcebergPartitionPlanData` address driver memory, and they change
`CometExecRDD`, `CometNativeShuffleInputRDD`, `NativeExecContext`,
`CometNativeScanExec`, and the Delta contrib scan. They also move each
partition's build inside Spark's scheduler lock, as the description notes.
On the table from #7490 (partition delete granularity, grouped read), as I
read it each delete file is used by one partition. So eager per-partition
messages from the previous comment would total about what the common block
holds on `main` today. The cost of eager messages is driver duplication of
entries many partitions share, such as a global equality delete. For context,
the driver already holds every `FileScanTask` and `DeleteFile` for the life of
the plan, because Iceberg's `SparkPartitioningAwareScan` caches its task groups
([line
69](https://github.com/apache/iceberg/blob/9f04b05090ade75d50eef5f244914181393adbb7/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkPartitioningAwareScan.java#L69),
filled at [lines
204-213](https://github.com/apache/iceberg/blob/9f04b05090ade75d50eef5f244914181393adbb7/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkPartitioningAwareScan.java#L204-L213)).
How were the 1.1x and 3.9x heap figures in the descri
ption measured? Driver heap numbers for eager messages on a shared-delete
layout would tell us whether the lazy build is needed, and a separate issue and
PR could carry them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
