This is an automated email from the ASF dual-hosted git repository.
RexXiong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 38f823141c [CELEBORN-2355] Return an empty stream from C++
ShuffleClientImpl::readPartition for an empty partition
38f823141c is described below
commit 38f823141cbb794d754185d9d24623535fd63302
Author: Nicholas Jiang <[email protected]>
AuthorDate: Tue Jun 30 10:15:03 2026 +0800
[CELEBORN-2355] Return an empty stream from C++
ShuffleClientImpl::readPartition for an empty partition
### What changes were proposed in this pull request?
The C++ `ShuffleClientImpl::readPartition` now returns an empty stream when
a shuffle has no locations for the requested partition, instead of constructing
a `CelebornInputStream` over zero locations. Specifically:
- Add a `CelebornInputStream::empty()` factory that builds a no-op stream
whose `read()` immediately yields `-1`.
- In `readPartition`, when `locations` is empty, log a warning and return
`CelebornInputStream::empty()`; the shuffle key is now built only on the
non-empty path.
Tests:
- A unit test for `CelebornInputStream::empty()` (immediate EOF, repeated
reads, zero-length read).
- Two `readPartition` tests (empty file groups, missing partition) backed
by a `VisibleForTesting` seam that seeds the cached file-group response so no
RPC to the `LifecycleManager` is needed.
### Why are the changes needed?
For a partition with no shuffle data, the previous code constructed a full
reader over an empty location list (also building the decompressor when
compression is enabled). Returning an empty stream short-circuits this and
makes the empty-shuffle case explicit, consistent with the read path's intent.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
Added unit tests under `cpp/celeborn/client/tests/`:
- `CelebornInputStreamTest.emptyStreamReadsReturnEof`
- `ShuffleClientImplTest.readPartitionEmptyWhenFileGroupsEmpty`
- `ShuffleClientImplTest.readPartitionEmptyWhenPartitionMissing`
Run with:
```
cd cpp && mkdir -p build && cd build && cmake .. && make
celeborn_client_test
./celeborn/client/tests/celeborn_client_test \
--gtest_filter='CelebornInputStreamTest.*:ShuffleClientImplTest.readPartition*'
```
Closes #3748 from SteNicholas/CELEBORN-2355.
Authored-by: Nicholas Jiang <[email protected]>
Signed-off-by: Shuang <[email protected]>
---
cpp/celeborn/client/ShuffleClient.cpp | 8 +++-
cpp/celeborn/client/ShuffleClient.h | 8 ++++
cpp/celeborn/client/reader/CelebornInputStream.cpp | 20 ++++++++
cpp/celeborn/client/reader/CelebornInputStream.h | 7 +++
cpp/celeborn/client/tests/CMakeLists.txt | 1 +
.../client/tests/CelebornInputStreamTest.cpp | 40 ++++++++++++++++
.../client/tests/ShuffleClientImplTest.cpp | 54 ++++++++++++++++++++++
7 files changed, 137 insertions(+), 1 deletion(-)
diff --git a/cpp/celeborn/client/ShuffleClient.cpp
b/cpp/celeborn/client/ShuffleClient.cpp
index 2549e5f992..45f6798a9d 100644
--- a/cpp/celeborn/client/ShuffleClient.cpp
+++ b/cpp/celeborn/client/ShuffleClient.cpp
@@ -818,13 +818,19 @@ std::unique_ptr<CelebornInputStream>
ShuffleClientImpl::readPartition(
bool needCompression) {
const auto reducerFileGroupInfo = getReducerFileGroupInfo(shuffleId);
CELEBORN_CHECK_NOT_NULL(reducerFileGroupInfo);
- std::string shuffleKey = utils::makeShuffleKey(appUniqueId_, shuffleId);
std::vector<std::shared_ptr<const protocol::PartitionLocation>> locations;
if (!reducerFileGroupInfo->fileGroups.empty() &&
reducerFileGroupInfo->fileGroups.count(partitionId)) {
locations = std::move(utils::toVector(
reducerFileGroupInfo->fileGroups.find(partitionId)->second));
}
+ // No locations: return an empty stream instead of building a reader.
+ if (locations.empty()) {
+ LOG(WARNING) << "Shuffle data is empty for shuffle " << shuffleId
+ << " partition " << partitionId << ".";
+ return CelebornInputStream::empty();
+ }
+ std::string shuffleKey = utils::makeShuffleKey(appUniqueId_, shuffleId);
return std::make_unique<CelebornInputStream>(
shuffleKey,
conf_,
diff --git a/cpp/celeborn/client/ShuffleClient.h
b/cpp/celeborn/client/ShuffleClient.h
index e2db120dc8..ee06dec0ba 100644
--- a/cpp/celeborn/client/ShuffleClient.h
+++ b/cpp/celeborn/client/ShuffleClient.h
@@ -289,6 +289,14 @@ class ShuffleClientImpl
std::optional<protocol::StatusCode> getPushTargetWorkerExcludeCause(
const protocol::PartitionLocation& location);
+ // @VisibleForTesting. Seeds the cached file-group response so readPartition
+ // runs without an RPC: getReducerFileGroupInfo then hits the cache.
+ void setReducerFileGroupInfoForTest(
+ int shuffleId,
+ std::shared_ptr<protocol::GetReducerFileGroupResponse> info) {
+ reducerFileGroupInfos_.set(shuffleId, std::move(info));
+ }
+
private:
std::shared_ptr<PushState> getPushState(const std::string& mapKey);
diff --git a/cpp/celeborn/client/reader/CelebornInputStream.cpp
b/cpp/celeborn/client/reader/CelebornInputStream.cpp
index f81ca0d1a9..6788cbcb08 100644
--- a/cpp/celeborn/client/reader/CelebornInputStream.cpp
+++ b/cpp/celeborn/client/reader/CelebornInputStream.cpp
@@ -77,6 +77,26 @@ CelebornInputStream::CelebornInputStream(
moveToNextReader();
}
+CelebornInputStream::CelebornInputStream()
+ : attemptNumber_(0),
+ startMapIndex_(0),
+ endMapIndex_(0),
+ shouldDecompress_(false),
+ currLocationIndex_(0),
+ currBatchPos_(0),
+ currBatchSize_(0),
+ fetchChunkRetryCnt_(0),
+ fetchChunkMaxRetry_(0),
+ retryWait_(utils::MS::zero()),
+ fetchExcludedWorkerExpireTimeoutMs_(0),
+ readSkewPartitionWithoutMapRange_(false),
+ shuffleClient_(nullptr) {}
+
+std::unique_ptr<CelebornInputStream> CelebornInputStream::empty() {
+ // Private constructor: not reachable via make_unique.
+ return std::unique_ptr<CelebornInputStream>(new CelebornInputStream());
+}
+
int CelebornInputStream::read(uint8_t* buffer, size_t offset, size_t len) {
CELEBORN_CHECK_NOT_NULL(buffer);
uint8_t* buf = buffer + offset;
diff --git a/cpp/celeborn/client/reader/CelebornInputStream.h
b/cpp/celeborn/client/reader/CelebornInputStream.h
index c333c9a45b..3af0ee47fc 100644
--- a/cpp/celeborn/client/reader/CelebornInputStream.h
+++ b/cpp/celeborn/client/reader/CelebornInputStream.h
@@ -54,9 +54,16 @@ class CelebornInputStream {
std::shared_ptr<const protocol::PartitionPushFailedBatches>
pushFailedBatches = nullptr);
+ // Returns an empty, no-op stream whose read() immediately yields -1, for a
+ // partition with no data.
+ static std::unique_ptr<CelebornInputStream> empty();
+
int read(uint8_t* buffer, size_t offset, size_t len);
private:
+ // The empty, no-op stream for empty(); members are zero/null-initialized.
+ CelebornInputStream();
+
bool fillBuffer();
bool moveToNextChunk();
diff --git a/cpp/celeborn/client/tests/CMakeLists.txt
b/cpp/celeborn/client/tests/CMakeLists.txt
index 796ecbc469..08623486b2 100644
--- a/cpp/celeborn/client/tests/CMakeLists.txt
+++ b/cpp/celeborn/client/tests/CMakeLists.txt
@@ -19,6 +19,7 @@ add_executable(
PushDataCallbackTest.cpp
PushMergedDataCallbackTest.cpp
ShuffleClientImplTest.cpp
+ CelebornInputStreamTest.cpp
CelebornInputStreamRetryTest.cpp
PushStateTest.cpp
ReviveManagerTest.cpp
diff --git a/cpp/celeborn/client/tests/CelebornInputStreamTest.cpp
b/cpp/celeborn/client/tests/CelebornInputStreamTest.cpp
new file mode 100644
index 0000000000..81299e3c26
--- /dev/null
+++ b/cpp/celeborn/client/tests/CelebornInputStreamTest.cpp
@@ -0,0 +1,40 @@
+/*
+ * 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.
+ */
+
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include "celeborn/client/reader/CelebornInputStream.h"
+
+using namespace celeborn;
+using namespace celeborn::client;
+
+// empty() returns a no-op stream: read() reports EOF (-1) immediately and
stays
+// at EOF across repeated and larger reads.
+TEST(CelebornInputStreamTest, emptyStreamReadsReturnEof) {
+ auto stream = CelebornInputStream::empty();
+ ASSERT_NE(stream, nullptr);
+
+ std::vector<uint8_t> buffer(16, 0);
+ EXPECT_EQ(stream->read(buffer.data(), 0, buffer.size()), -1);
+ // No state to advance: a second read still reports EOF.
+ EXPECT_EQ(stream->read(buffer.data(), 0, buffer.size()), -1);
+ // A zero-length read is a no-op that returns 0, not EOF.
+ EXPECT_EQ(stream->read(buffer.data(), 0, 0), 0);
+}
diff --git a/cpp/celeborn/client/tests/ShuffleClientImplTest.cpp
b/cpp/celeborn/client/tests/ShuffleClientImplTest.cpp
index 35be191725..22e85fbee2 100644
--- a/cpp/celeborn/client/tests/ShuffleClientImplTest.cpp
+++ b/cpp/celeborn/client/tests/ShuffleClientImplTest.cpp
@@ -15,10 +15,15 @@
* limitations under the License.
*/
+#include <climits>
+#include <cstdint>
+#include <vector>
+
#include <gtest/gtest.h>
#include "celeborn/client/ShuffleClient.h"
#include "celeborn/client/tests/ShuffleClientTestUtils.h"
+#include "celeborn/protocol/ControlMessages.h"
using namespace celeborn;
using namespace celeborn::client;
@@ -50,6 +55,12 @@ class TestableShuffleClient : public TestShuffleClientBase {
return getPushTargetWorkerExcludeCause(location);
}
+ void seedReducerFileGroupInfo(
+ int shuffleId,
+ std::shared_ptr<protocol::GetReducerFileGroupResponse> info) {
+ setReducerFileGroupInfoForTest(shuffleId, std::move(info));
+ }
+
private:
explicit TestableShuffleClient(
const std::shared_ptr<const conf::CelebornConf>& conf)
@@ -268,3 +279,46 @@ TEST(ShuffleClientImplTest, shutdownClearsExcludedWorkers)
{
EXPECT_FALSE(client->callGetPushTargetWorkerExcludeCause(*loc).has_value());
}
+
+// Asserts the stream reports EOF immediately, i.e. it carries no data.
+static void expectEmptyStream(
+ const std::unique_ptr<CelebornInputStream>& stream) {
+ ASSERT_NE(stream, nullptr);
+ std::vector<uint8_t> buffer(8, 0);
+ EXPECT_EQ(stream->read(buffer.data(), 0, buffer.size()), -1);
+}
+
+// readPartition returns an empty (EOF) stream when the shuffle's file groups
+// hold no locations for the partition. The seeded cache makes
+// getReducerFileGroupInfo hit, so no RPC is attempted.
+TEST(ShuffleClientImplTest, readPartitionEmptyWhenFileGroupsEmpty) {
+ auto client = TestableShuffleClient::create(makeConf(/*enabled=*/false));
+ auto info = std::make_shared<protocol::GetReducerFileGroupResponse>();
+ info->status = protocol::StatusCode::SUCCESS;
+ client->seedReducerFileGroupInfo(/*shuffleId=*/1, info);
+
+ expectEmptyStream(client->readPartition(
+ /*shuffleId=*/1,
+ /*partitionId=*/0,
+ /*attemptNumber=*/0,
+ /*startMapIndex=*/0,
+ /*endMapIndex=*/INT_MAX));
+}
+
+// Same empty-stream result when the file-group map is populated but has no
+// entry for the requested partition.
+TEST(ShuffleClientImplTest, readPartitionEmptyWhenPartitionMissing) {
+ auto client = TestableShuffleClient::create(makeConf(/*enabled=*/false));
+ auto info = std::make_shared<protocol::GetReducerFileGroupResponse>();
+ info->status = protocol::StatusCode::SUCCESS;
+ info->fileGroups[7].insert(makeLocation(7, 0, "host-7", 9007));
+ client->seedReducerFileGroupInfo(/*shuffleId=*/2, info);
+
+ // Partition 0 is absent (only partition 7 has locations).
+ expectEmptyStream(client->readPartition(
+ /*shuffleId=*/2,
+ /*partitionId=*/0,
+ /*attemptNumber=*/0,
+ /*startMapIndex=*/0,
+ /*endMapIndex=*/INT_MAX));
+}
