Skip to content

[VL] Unify gpu hash shuffle writer output#12499

Open
marin-ma wants to merge 3 commits into
apache:mainfrom
marin-ma:unify-gpu-shuffle-writer-output
Open

[VL] Unify gpu hash shuffle writer output#12499
marin-ma wants to merge 3 commits into
apache:mainfrom
marin-ma:unify-gpu-shuffle-writer-output

Conversation

@marin-ma

Copy link
Copy Markdown
Contributor

Currently, gpu shuffle writer splits the boolean and timestamp vector to buffers that is same as arrow/cudf format because it assumes the shuffle reader stage will execute on GPU and use gpu shuffle reader. However, the stage execution is decided at runtime when AQE is on. If the next stage is planed to execute on CPU, the shuffle read will fail.

In order to support a more general hybrid execution mode, it's necessary to unify the cpu/gpu shuffle writer output for a specific shuffle writer type (currently hash or sort).

Copilot AI review requested due to automatic review settings July 10, 2026 13:38
@github-actions github-actions Bot added the VELOX label Jul 10, 2026
@marin-ma marin-ma marked this pull request as draft July 10, 2026 13:44
@marin-ma marin-ma removed the request for review from Copilot July 10, 2026 13:44
@marin-ma marin-ma marked this pull request as ready for review July 10, 2026 14:42
Copilot AI review requested due to automatic review settings July 10, 2026 14:42

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR aligns GPU-hash shuffle writer output with the CPU/hash shuffle writer so that shuffle data remains readable when AQE changes the downstream stage to CPU execution (hybrid CPU/GPU mode).

Changes:

  • Route kGpuHashShuffle writer creation through VeloxHashShuffleWriter to unify on a single hash-shuffle output format.
  • Remove the dedicated VeloxGpuHashShuffleWriter implementation and shift boolean/timestamp value normalization into GpuBufferColumnarBatch::compose.
  • Update JNI/test option wiring to use HashShuffleWriterOptions.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
cpp/velox/tests/VeloxGpuShuffleWriterTest.cc Updates GPU-hash test configuration to use HashShuffleWriterOptions.
cpp/velox/shuffle/VeloxShuffleWriter.cc Treats kGpuHashShuffle as an alias for hash shuffle writer creation.
cpp/velox/shuffle/VeloxGpuShuffleWriter.h Removes the GPU hash shuffle writer class definition.
cpp/velox/shuffle/VeloxGpuShuffleWriter.cc Removes the GPU hash shuffle writer implementation.
cpp/velox/memory/GpuBufferColumnarBatch.cc Adds boolean/timestamp value conversions when composing GPU buffer batches.
cpp/velox/CMakeLists.txt Drops the removed GPU shuffle writer source from the GPU build.
cpp/core/jni/JniWrapper.cc Switches GPU-hash JNI writer creation to HashShuffleWriterOptions.
Comments suppressed due to low confidence (1)

cpp/core/jni/JniWrapper.cc:1105

  • createGpuHashShuffleWriter now builds a HashShuffleWriterOptions, which defaults shuffleWriterType to kHashShuffle. Since VeloxRuntime::createShuffleWriter selects the writer using options->shuffleWriterType, this JNI API will no longer create a kGpuHashShuffle writer (behavior change vs the previous GpuHashShuffleWriterOptions). If the caller still expects a GPU-hash writer type, set shuffleWriterType explicitly after construction.
  auto shuffleWriterOptions = std::make_shared<HashShuffleWriterOptions>(
      toPartitioning(jStringToCString(env, partitioningNameJstr)),
      startPartitionId,
      splitBufferSize,
      splitBufferReallocThreshold);

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

@@ -21,7 +21,6 @@
#include "config/GlutenConfig.h"
#include "memory/GpuBufferColumnarBatch.h"
#include "shuffle/VeloxGpuShuffleWriter.h"
Comment on lines +166 to +172
VELOX_CHECK_LE(valueBufferOffset + bufferSize, returnBuffers[bufferIdx + 1]->size());
const auto* src = reinterpret_cast<const int64_t*>(batch->bufferAt(bufferIdx + 1)->data());
auto* dst = reinterpret_cast<int64_t*>(returnBuffers[bufferIdx + 1]->mutable_data());
for (auto j = 0; j < batch->numRows(); ++j) {
dst[valueBufferOffset + j] = src[j << 1] * 1'000'000'000L + src[j << 1 | 1];
}
valueBufferOffset += bufferSize;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants