[CELEBORN-2219][CIP-14] Support PushMergedData in CppClient by afterincomparableyum · Pull Request #3611 · apache/celeborn

afterincomparableyum · 2026-03-02T04:43:08Z

What changes were proposed in this pull request?

Implement PushMergedData functionality in the C++ client, enabling batch merging and pushing of shuffle data grouped by worker address.

Key changes:

Add mergeData() and pushMergedData() to ShuffleClient, which accumulateper partition data batches and push them as merged payloads when the buffer threshold is exceeded or at mapper end.
Introduce DataBatches class to manage batch accumulation thread-safe add/take operations and size-bounded requireBatches().
Add PushMergedDataCallback to handle success responses (split handling, congestion control, MAP_ENDED) and failure responses with revive-based retry via submitRetryPushMergedData().
Add PushMergedData network message type with encoding for partitionUniqueIds and batchOffsets arrays.
Extend Encoders with encode/decode support for vector and vector<int32_t>.
Add pushMergedDataAsync() to TransportClient.
Add unit tests for DataBatches, PushMergedData message encoding, and array encoders.

Why are the changes needed?

This is needed to extend the functionality of C++ Client and there is a Bolt dependency on it with bytedance/bolt#370

Does this PR resolve a correctness bug?

No

Does this PR introduce any user-facing change?

Yes because it is a new functionality in the c++ client

How was this patch tested?

Tested through running unit tests and compiling locally.

SteNicholas · 2026-03-02T06:18:00Z

@afterincomparableyum, could you please update the description of the pull request which should refer to the template of pull request?

Implement PushMergedData functionality in the C++ client, enabling batch merging and pushing of shuffle data grouped by worker address. Key changes: - Add mergeData() and pushMergedData() to ShuffleClient, which accumulate per partition data batches and push them as merged payloads when the buffer threshold is exceeded or at mapper end. - Introduce DataBatches class to manage batch accumulation with thread-safe add/take operations and size-bounded requireBatches(). - Add PushMergedDataCallback to handle success responses (split handling, congestion control, MAP_ENDED) and failure responses with revive-based retry via submitRetryPushMergedData(). - Add PushMergedData network message type with encoding for partitionUniqueIds and batchOffsets arrays. - Extend Encoders with encode/decode support for vector<string> and vector<int32_t>. - Add pushMergedDataAsync() to TransportClient. - Add unit tests for DataBatches, PushMergedData message encoding, and array encoders. Tested through running unit tests and compiling locally.

afterincomparableyum · 2026-03-09T03:32:51Z

ping @SteNicholas @FMX @RexXiong @HolyLow for review. Thank you for your review too.

SteNicholas · 2026-03-10T02:46:51Z

@afterincomparableyum, the description of this pull request does not follow the template of the pull request yet. Please update the description.

Copilot

Pull request overview

This PR adds PushMergedData support to the C++ Celeborn client by buffering per-partition batches grouped by worker address, merging buffered payloads, and pushing them as a single request (with split/congestion/MAP_ENDED handling and revive-based retries).

Changes:

Add array encode/decode helpers in protocol::Encoders and tests for vector/vector<int32_t>.
Introduce network::PushMergedData message + TransportClient::pushMergedDataAsync() and associated message tests.
Add DataBatches + PushMergedDataCallback and integrate merged pushing into ShuffleClientImpl (flush at mapper end) with new unit tests.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
cpp/celeborn/protocol/Encoders.h	Adds declarations for vector/vector<int32_t> encoding helpers
cpp/celeborn/protocol/Encoders.cpp	Implements array encodedLength/encode/decode helpers
cpp/celeborn/protocol/tests/EncodersTest.cpp	Adds unit tests for the new array encoders
cpp/celeborn/network/Message.h	Introduces `PushMergedData` message class
cpp/celeborn/network/Message.cpp	Implements `PushMergedData` encoding and copy ctor
cpp/celeborn/network/MessageDispatcher.cpp	Ensures dispatcher extracts requestId for PUSH_MERGED_DATA
cpp/celeborn/network/TransportClient.h	Adds `pushMergedDataAsync()` API
cpp/celeborn/network/TransportClient.cpp	Implements async send path for `PushMergedData`
cpp/celeborn/network/tests/MessageTest.cpp	Adds encoding/copy-constructor tests for PushMergedData
cpp/celeborn/client/writer/DataBatches.h	New thread-safe batch accumulator abstraction
cpp/celeborn/client/writer/DataBatches.cpp	Implements batch accumulation and size-bounded draining
cpp/celeborn/client/writer/PushMergedDataCallback.h	New callback for PushMergedData response/failure handling
cpp/celeborn/client/writer/PushMergedDataCallback.cpp	Implements split parsing, congestion control, revive-based retry
cpp/celeborn/client/writer/PushState.h	Adds batch-buffering map and merge-related APIs
cpp/celeborn/client/writer/PushState.cpp	Implements buffering/take operations for merged pushing
cpp/celeborn/client/ShuffleClient.h	Adds `mergeData()` + `pushMergedData()` interface and impl hooks
cpp/celeborn/client/ShuffleClient.cpp	Implements merge/buffer/flush/push logic and retry regrouping
cpp/celeborn/client/tests/DataBatchesTest.cpp	Adds unit tests for DataBatches behavior
cpp/celeborn/client/tests/CelebornInputStreamRetryTest.cpp	Updates test stub to satisfy new ShuffleClient interface
cpp/celeborn/client/tests/CMakeLists.txt	Adds DataBatchesTest to client test target
cpp/celeborn/client/CMakeLists.txt	Adds new writer sources to the client library

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

cpp/celeborn/protocol/Encoders.cpp

cpp/celeborn/client/writer/DataBatches.cpp

cpp/celeborn/client/ShuffleClient.cpp

cpp/celeborn/client/writer/PushMergedDataCallback.cpp

cpp/celeborn/client/ShuffleClient.cpp

cpp/celeborn/client/writer/PushState.cpp

afterincomparableyum · 2026-03-11T04:05:18Z

Resolved comments @SteNicholas, and left a couple of TODOs in this PR as well based off of some of the comments from CoPilot.

afterincomparableyum · 2026-03-16T22:49:03Z

@SteNicholas wanted to follow up and see if you had a chance to look over the PR again. Thanks!

FMX · 2026-03-17T02:54:10Z

It's an interesting PR, let me watch it.

afterincomparableyum · 2026-03-17T09:07:54Z

@SteNicholas you might notice some small gaps between java and c++ implementation of this.

that’s mainly because SlowStartPushStrategy has not been implemented yet and PushExcludedWorker has not been implemented yet too.

I’ll add some detailed TODO comments to document what I need to do for this as well after I double check what needs to be done.

Copilot

Pull request overview

Copilot reviewed 22 out of 22 changed files in this pull request and generated 8 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

cpp/celeborn/protocol/Encoders.cpp

+      static_cast<size_t>(count) * sizeof(int32_t),
+      buffer.remainingSize(),


cpp/celeborn/protocol/Encoders.cpp

+    const std::vector<std::string>& arr) {
+  buffer.write<int>(static_cast<int>(arr.size()));
+  for (const auto& s : arr) {
+    encode(buffer, s);
+  }


cpp/celeborn/client/writer/PushMergedDataCallback.cpp

+    return;
+  }
+
+  if (response->remainingSize() <= 0) {


cpp/celeborn/client/writer/PushMergedDataCallback.cpp

+    case protocol::StatusCode::MAP_ENDED: {
+      auto mapperEndSet = sharedClient->mapperEndSets().computeIfAbsent(
+          shuffleId_,
+          []() { return std::make_shared<utils::ConcurrentHashSet<int>>(); });
+      mapperEndSet->insert(mapId_);
+      pushState_->removeBatch(groupedBatchId_, hostAndPushPort_);
+      break;


cpp/celeborn/client/writer/PushMergedDataCallback.cpp

+    default: {
+      LOG(WARNING) << "unhandled PushMergedData success StatusCode: " << reason;
+      pushState_->removeBatch(groupedBatchId_, hostAndPushPort_);
+    }


cpp/celeborn/client/ShuffleClient.cpp

+  auto host = hostAndPushPort.substr(0, hostAndPushPort.find(':'));
+  auto portStr = hostAndPushPort.substr(hostAndPushPort.find(':') + 1);
+  auto port = static_cast<uint16_t>(std::stoi(portStr));


cpp/celeborn/protocol/Encoders.cpp

+  int count = buffer.read<int>();
+  CELEBORN_CHECK_GE(count, 0, "Invalid string array count: {}", count);
+  CELEBORN_CHECK_LE(
+      static_cast<size_t>(count) * sizeof(int),
+      buffer.remainingSize(),
+      "String array count {} exceeds remaining buffer size {}",
+      count,
+      buffer.remainingSize());


cpp/celeborn/client/writer/PushState.cpp

+  auto result = batchesMap_.get(addressPairKey);
+  if (result.has_value()) {
+    batchesMap_.erase(addressPairKey);
+    return result.value();


afterincomparableyum · 2026-03-26T00:58:40Z

I will take a look at comments soon

github-actions bot added the module:cpp label Mar 2, 2026

afterincomparableyum force-pushed the cpp-client/celeborn-2219 branch from 8149bb9 to a234845 Compare March 9, 2026 03:25

afterincomparableyum changed the title ~~[WIP][CELEBORN-2219][CIP-14] Support PushMergedData in CppClient~~ [CELEBORN-2219][CIP-14] Support PushMergedData in CppClient Mar 9, 2026

afterincomparableyum marked this pull request as ready for review March 9, 2026 03:32

afterincomparableyum mentioned this pull request Mar 9, 2026

[CELEBORN-2275][CIP-14] Add C++ merge-write and Java-read hybrid integration test #3619

Draft

SteNicholas requested a review from Copilot March 10, 2026 02:44

Copilot started reviewing on behalf of SteNicholas March 10, 2026 02:44 View session

Copilot AI reviewed Mar 10, 2026

View reviewed changes

This was referenced Mar 10, 2026

[Feature] use mergeData for bolt NativeCelebornClient bytedance/bolt#370

Open

Support integration with Celeborn C++ client bytedance/bolt#13

Open

address the comments of copilot

85022fa

afterincomparableyum force-pushed the cpp-client/celeborn-2219 branch from 722dae1 to 85022fa Compare March 11, 2026 04:04

github-actions bot added the module:common label Mar 11, 2026

add TODO in java client

5d5c4b5

afterincomparableyum force-pushed the cpp-client/celeborn-2219 branch from 4eef732 to 5d5c4b5 Compare March 11, 2026 04:07

SteNicholas force-pushed the main branch from 740993b to 8c2c952 Compare March 11, 2026 11:43

SteNicholas requested a review from Copilot March 19, 2026 09:49

Copilot AI reviewed Mar 19, 2026

View reviewed changes

		static_cast<size_t>(count) * sizeof(int32_t),
		buffer.remainingSize(),

Conversation

afterincomparableyum commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR resolve a correctness bug?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

SteNicholas commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

afterincomparableyum commented Mar 9, 2026

Uh oh!

SteNicholas commented Mar 10, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

afterincomparableyum commented Mar 11, 2026

Uh oh!

afterincomparableyum commented Mar 16, 2026

Uh oh!

FMX commented Mar 17, 2026

Uh oh!

afterincomparableyum commented Mar 17, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

afterincomparableyum commented Mar 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

afterincomparableyum commented Mar 2, 2026 •

edited

Loading

SteNicholas commented Mar 2, 2026 •

edited

Loading