This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/branch-0.7 by this push:
new 6b0c0b0214 [CELEBORN-2481] Fix NPE when persisting MapFileMeta without
segment granularity
6b0c0b0214 is described below
commit 6b0c0b0214751a06dc2340736bd6f7ac6242f33e
Author: Rudy Attias <[email protected]>
AuthorDate: Mon Oct 19 12:05:21 2026 +0800
[CELEBORN-2481] Fix NPE when persisting MapFileMeta without segment
granularity
### What changes were proposed in this pull request?
`PbSerDeUtils.toPbFileInfo` now adds the MAP partition's writing-segment
map and segment indexes only when they are non-null.
`MapFileMeta` creates both lazily, on the first `addPartitionSegmentId` /
`addSegmentIdAndFirstBufferIndex` call. The worker builds MAP file metas with
`new MapFileMeta()`, so for a MAP partition written without segment granularity
both fields stay `null`, and `putAllPartitionWritingSegment(null)` threw:
```
java.lang.NullPointerException: Cannot invoke "java.util.Map.keySet()"
because "m" is null
at
com.google.protobuf.MapField$MutabilityAwareMap.putAll(MapField.java:312)
at
org.apache.celeborn.common.protocol.PbFileInfo$Builder.putAllPartitionWritingSegment(PbFileInfo.java:1674)
at
org.apache.celeborn.common.util.PbSerDeUtils$.toPbFileInfo(PbSerDeUtils.scala:164)
```
On restore, `fromPbFileInfo` restores both fields as `null` when they are
empty, matching a live meta that never recorded a segment, so
`MapFileMeta#hasPartitionSegmentIds()` gives the same answer before and after a
restart. Restored metas are already marked writer-closed, so nothing writes to
them. Partitions with segment granularity serialize and restore exactly as
before.
### Why are the changes needed?
When graceful shutdown is enabled, the worker persists committed file infos
through `toPbFileInfoMap` → `toPbFileInfo` in two places, and the NPE breaks
both:
- **The periodic saver**
(`worker-storage-manager-committed-fileinfo-saver`). It runs with
`scheduleWithFixedDelay` and does not catch the NPE, so the first failure
cancels it and nothing is saved afterwards.
- **`saveAllCommittedFileInfosToDB` at shutdown.** It logs the error and
skips any shuffle that contains such a MAP partition.
As a result, after a graceful restart the worker cannot recover MAP
partitions written without segment granularity, such as Flink shuffle
partitions. Readers fail to find them, and the producing tasks have to be
re-run.
### Does this PR resolve a correctness bug?
- [ ] Yes
Shuffle data is never incorrect. Partitions are lost on graceful restart
instead of being recovered.
### Does this PR introduce _any_ user-facing change?
- [x] Yes
With graceful shutdown enabled, a worker restart now recovers committed MAP
partitions without segment granularity. No configuration or protocol change.
### How was this patch tested?
Added two tests to `PbSerDeUtilsTest`:
- `fromAndToPbFileInfo with MapFileMeta without segment granularity` builds
the meta the way the worker does (no-arg constructor plus buffer size and
subpartition count). It round-trips it through `toPbFileInfo`/`fromPbFileInfo`
and through `toPbFileInfoMap`/`fromPbFileInfoMap`, checking the buffer size,
subpartition count, segment-granularity flag, writer-closed flag, that
`hasPartitionSegmentIds()` is false and the segment indexes are null, a segment
lookup, path and length. On `mai [...]
- `fromAndToPbFileInfo with MapFileMeta with segment granularity` is a
regression check. It confirms that `hasPartitionSegmentIds()` stays true and
that recorded segment IDs and segment indexes still round-trip.
```
build/mvn -pl common -am -Dtest=none
-DwildcardSuites=org.apache.celeborn.common.util.PbSerDeUtilsTest test
Tests: succeeded 36, failed 0
```
Closes #3874 from attiasr/fix-map-fileinfo-serde.
Authored-by: Rudy Attias <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/19
AI-Contributed/UT: 0/60
(cherry picked from commit bc06a1830148672ef50edb3c946d9a9b349ef962)
Signed-off-by: zhengtao <[email protected]>
---
.../apache/celeborn/common/util/PbSerDeUtils.scala | 19 +++++--
.../celeborn/common/util/PbSerDeUtilsTest.scala | 60 ++++++++++++++++++++++
2 files changed, 74 insertions(+), 5 deletions(-)
diff --git
a/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
b/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
index e9c407ce80..cac7550871 100644
--- a/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
@@ -104,12 +104,16 @@ object PbSerDeUtils {
case PartitionType.REDUCE =>
new ReduceFileMeta(pbFileInfo.getChunkOffsetsList)
case PartitionType.MAP =>
+ // Protobuf cannot tell an empty field from an unset one; restore both
as null, which
+ // is what MapFileMeta#hasPartitionSegmentIds reads as "no segments
recorded".
val fileMeta = new MapFileMeta(
pbFileInfo.getBufferSize,
pbFileInfo.getNumSubpartitions,
pbFileInfo.getIsSegmentGranularityVisible,
- pbFileInfo.getPartitionWritingSegmentMap,
- fromPbSegmentIndexList(pbFileInfo.getSegmentIndexList))
+ if (pbFileInfo.getPartitionWritingSegmentCount == 0) null
+ else pbFileInfo.getPartitionWritingSegmentMap,
+ if (pbFileInfo.getSegmentIndexCount == 0) null
+ else fromPbSegmentIndexList(pbFileInfo.getSegmentIndexList))
// writer always closed as this is committed file info.
fileMeta.setIsWriterClosed(true)
fileMeta
@@ -161,9 +165,14 @@ object PbSerDeUtils {
builder.setBufferSize(mapFileMeta.getBufferSize)
builder.setNumSubpartitions(mapFileMeta.getNumSubpartitions)
builder.setIsSegmentGranularityVisible(mapFileMeta.isSegmentGranularityVisible)
-
builder.putAllPartitionWritingSegment(mapFileMeta.getSubPartitionWritingSegmentId)
- builder.addAllSegmentIndex(
-
mapFileMeta.getSubPartitionSegmentIndexes.asScala.map(toPbSegmentIndex).asJava)
+ // Both are created lazily on the first segment: null without segment
granularity.
+ if (mapFileMeta.getSubPartitionWritingSegmentId != null) {
+
builder.putAllPartitionWritingSegment(mapFileMeta.getSubPartitionWritingSegmentId)
+ }
+ if (mapFileMeta.getSubPartitionSegmentIndexes != null) {
+ builder.addAllSegmentIndex(
+
mapFileMeta.getSubPartitionSegmentIndexes.asScala.map(toPbSegmentIndex).asJava)
+ }
} else {
val reduceFileMeta = fileInfo.getFileMeta.asInstanceOf[ReduceFileMeta]
builder.setPartitionType(PartitionType.REDUCE.getValue)
diff --git
a/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
b/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
index 15c4df0fc7..3010f7e049 100644
---
a/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
+++
b/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
@@ -410,6 +410,66 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
assert(restoredFileInfo6.isDFS && restoredFileInfo6.isS3)
}
+ test("fromAndToPbFileInfo with MapFileMeta without segment granularity") {
+ // A worker creates MAP file metas with the no-arg constructor and only
sets the
+ // buffer size and subpartition count on handshake.
+ val mapFileMeta = new MapFileMeta()
+ mapFileMeta.setBufferSize(32 * 1024)
+ mapFileMeta.setNumSubPartitions(8)
+ val mapFileInfo = new DiskFileInfo(
+ userIdentifier1,
+ true,
+ mapFileMeta,
+ file1.getAbsolutePath,
+ StorageInfo.Type.HDD,
+ 6000L)
+
+ val restoredMeta =
PbSerDeUtils.fromPbFileInfo(PbSerDeUtils.toPbFileInfo(mapFileInfo))
+ .getFileMeta.asInstanceOf[MapFileMeta]
+ assert(restoredMeta.getBufferSize == 32 * 1024)
+ assert(restoredMeta.getNumSubpartitions == 8)
+ assert(!restoredMeta.isSegmentGranularityVisible)
+ assert(restoredMeta.isWriterClosed)
+ assert(!restoredMeta.hasPartitionSegmentIds)
+ assert(restoredMeta.getSubPartitionSegmentIndexes == null)
+ assert(restoredMeta.getSegmentIdByFirstBufferIndex(0, 0) == null)
+
+ val files = JavaUtils.newConcurrentHashMap[String, DiskFileInfo]()
+ files.put("mapFile", mapFileInfo)
+ val restoredFiles = PbSerDeUtils.fromPbFileInfoMap(
+ PbSerDeUtils.toPbFileInfoMap(files),
+ JavaUtils.newConcurrentHashMap[String, UserIdentifier](),
+ mountPoints)
+ val restoredFile = restoredFiles.get("mapFile")
+ assert(restoredFile.getFilePath == mapFileInfo.getFilePath)
+ assert(restoredFile.getFileLength == mapFileInfo.getFileLength)
+
assert(restoredFile.getFileMeta.asInstanceOf[MapFileMeta].getNumSubpartitions
== 8)
+ }
+
+ test("fromAndToPbFileInfo with MapFileMeta with segment granularity") {
+ val mapFileMeta = new MapFileMeta()
+ mapFileMeta.setBufferSize(32 * 1024)
+ mapFileMeta.setNumSubPartitions(2)
+ mapFileMeta.setSegmentGranularityVisible(true)
+ mapFileMeta.addPartitionSegmentId(1, 3)
+ mapFileMeta.addSegmentIdAndFirstBufferIndex(1, 5, 3)
+ val mapFileInfo = new DiskFileInfo(
+ userIdentifier1,
+ true,
+ mapFileMeta,
+ file1.getAbsolutePath,
+ StorageInfo.Type.HDD,
+ 6000L)
+
+ val restoredMeta =
PbSerDeUtils.fromPbFileInfo(PbSerDeUtils.toPbFileInfo(mapFileInfo))
+ .getFileMeta.asInstanceOf[MapFileMeta]
+ assert(restoredMeta.isSegmentGranularityVisible)
+ assert(restoredMeta.hasPartitionSegmentIds)
+ assert(restoredMeta.getPartitionWritingSegmentId(1) == 3)
+ assert(restoredMeta.getSegmentIdByFirstBufferIndex(1, 5) == 3)
+ assert(restoredMeta.getSegmentIdByFirstBufferIndex(0, 5) == null)
+ }
+
test("fromAndToPbUserIdentifier") {
val pbUserIdentifier = PbSerDeUtils.toPbUserIdentifier(userIdentifier1)
val restoredUserIdentifier =
PbSerDeUtils.fromPbUserIdentifier(pbUserIdentifier)
