parthchandra commented on code in PR #7491:
URL: https://github.com/apache/datafusion-comet/pull/7491#discussion_r4484543536
##########
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:
Done. The five interning maps are created per partition inside the
partitions.foreach loop, and each partition's message comes out of the single
pass with indices starting at 0. IcebergScanPools and its suite are removed.
Coverage is in PlanDataInjectorSuite (a new merge test),
CometIcebergNativeSuite and the SQL file
##########
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:
Done. A partition's pools now ship in its existing IcebergScan.common, and
inject merges them with mergeFrom. IcebergFilePools, own_pools, the pool mirror
types and the guard test are gone
##########
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:
Agreed. Since commonByKey rides in the task binary, the threshold only saved
bytes on the wire, not task heap. It's removed, and the common block holds none
of the five pools.
--
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]
