This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 204b8d7128 [VL] Unify gpu hash shuffle writer output (#12499)
204b8d7128 is described below
commit 204b8d71288a7f709e22331d28f704fc4cb78593
Author: Rong Ma <[email protected]>
AuthorDate: Wed Jul 15 14:16:00 2026 +0100
[VL] Unify gpu hash shuffle writer output (#12499)
---
cpp/core/jni/JniWrapper.cc | 2 +-
cpp/velox/CMakeLists.txt | 1 -
cpp/velox/memory/GpuBufferColumnarBatch.cc | 73 ++++++++++++++++++++--------
cpp/velox/shuffle/VeloxGpuShuffleWriter.cc | 72 ---------------------------
cpp/velox/shuffle/VeloxGpuShuffleWriter.h | 57 ----------------------
cpp/velox/shuffle/VeloxShuffleWriter.cc | 9 +---
cpp/velox/tests/VeloxGpuShuffleWriterTest.cc | 5 +-
7 files changed, 57 insertions(+), 162 deletions(-)
diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
index 29036132a4..a01b6a8d0a 100644
--- a/cpp/core/jni/JniWrapper.cc
+++ b/cpp/core/jni/JniWrapper.cc
@@ -1048,7 +1048,7 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe
}
ObjectStore::release(partitionWriterHandle);
- auto shuffleWriterOptions = std::make_shared<GpuHashShuffleWriterOptions>(
+ auto shuffleWriterOptions = std::make_shared<HashShuffleWriterOptions>(
toPartitioning(jStringToCString(env, partitioningNameJstr)),
startPartitionId,
splitBufferSize,
diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt
index b9cffc92fa..636b073045 100644
--- a/cpp/velox/CMakeLists.txt
+++ b/cpp/velox/CMakeLists.txt
@@ -227,7 +227,6 @@ if(ENABLE_GPU)
operators/plannodes/CudfVectorStream.cc
shuffle/VeloxGpuAsyncShuffleReader.cc
shuffle/VeloxGpuShuffleReader.cc
- shuffle/VeloxGpuShuffleWriter.cc
operators/serializer/VeloxGpuColumnarBatchSerializer.cc
utils/GpuBufferBatchResizer.cc
memory/GpuBufferColumnarBatch.cc)
diff --git a/cpp/velox/memory/GpuBufferColumnarBatch.cc
b/cpp/velox/memory/GpuBufferColumnarBatch.cc
index b76a849c82..eab90eb2f8 100644
--- a/cpp/velox/memory/GpuBufferColumnarBatch.cc
+++ b/cpp/velox/memory/GpuBufferColumnarBatch.cc
@@ -28,7 +28,7 @@
namespace gluten {
namespace {
-enum class BufferType { kNull, kLength, kValue };
+enum class BufferType { kNull, kLength, kValue, kBooleanValue, kTimestampValue
};
}
using namespace facebook;
@@ -67,16 +67,20 @@ std::shared_ptr<GpuBufferColumnarBatch>
GpuBufferColumnarBatch::compose(
GLUTEN_CHECK(!batches.empty(), "No batches to compose");
// Compute the returned GpuBufferColumnarBatch buffers.
auto& type = batches[0]->getRowType();
- const auto bufferSize = batches[0]->buffers().size();
+ const auto numBuffers = batches[0]->buffers().size();
std::vector<size_t> bufferSizes;
- bufferSizes.resize(bufferSize);
+ bufferSizes.resize(numBuffers);
std::vector<BufferType> bufferTypes;
- bufferTypes.reserve(bufferSize);
+ bufferTypes.reserve(numBuffers);
for (const auto& colType : type->children()) {
bufferSizes[bufferTypes.size()] = arrow::bit_util::BytesForBits(numRows);
bufferTypes.push_back(BufferType::kNull);
- if (colType->isFixedWidth()) {
+ if (colType->isTimestamp()) {
+ bufferTypes.push_back(BufferType::kTimestampValue);
+ } else if (colType->isBoolean()) {
+ bufferTypes.push_back(BufferType::kBooleanValue);
+ } else if (colType->isFixedWidth()) {
bufferTypes.push_back(BufferType::kValue);
} else {
// Add the first offset 0.
@@ -86,27 +90,37 @@ std::shared_ptr<GpuBufferColumnarBatch>
GpuBufferColumnarBatch::compose(
bufferTypes.push_back(BufferType::kValue);
}
}
- VELOX_CHECK_EQ(bufferTypes.size(), bufferSize);
+ VELOX_CHECK_EQ(bufferTypes.size(), numBuffers);
// This buffer may be more than the actual reauired buffer for null buffer.
for (const auto& batch : batches) {
if (batch->numRows() == 0) {
continue;
}
- for (auto i = 0; i < bufferSize; ++i) {
+ for (auto i = 0; i < numBuffers; ++i) {
// The null buffer may be null or length = 0.
// Maybe optimize later, detect if the null buffer is all true. And set
the return null buffer to 0.
if (bufferTypes[i] == BufferType::kNull || bufferTypes[i] ==
BufferType::kLength) {
continue;
}
+
auto& buffer = batch->bufferAt(i);
VELOX_CHECK_NOT_NULL(buffer);
- bufferSizes[i] += buffer->size();
+
+ if (bufferTypes[i] == BufferType::kTimestampValue) {
+ // Velox Timestamp value is 16 bytes. CUDF Timestamp value is 8 bytes.
+ bufferSizes[i] += static_cast<size_t>(batch->numRows()) *
sizeof(int64_t);
+ } else if (bufferTypes[i] == BufferType::kBooleanValue) {
+ // Velox Boolean value is 1 bit. CUDF Boolean value is 1 byte.
+ bufferSizes[i] += batch->numRows();
+ } else {
+ bufferSizes[i] += buffer->size();
+ }
}
}
std::vector<std::shared_ptr<arrow::Buffer>> returnBuffers;
- returnBuffers.reserve(bufferSize);
- for (auto i = 0; i < bufferSize; ++i) {
+ returnBuffers.reserve(numBuffers);
+ for (auto i = 0; i < numBuffers; ++i) {
// Defer the null buffer to really contains null.
if (bufferTypes[i] == BufferType::kNull) {
returnBuffers.emplace_back(nullptr);
@@ -121,8 +135,8 @@ std::shared_ptr<GpuBufferColumnarBatch>
GpuBufferColumnarBatch::compose(
int32_t bufferIdx = 0;
for (const auto& colType : type->children()) {
size_t rowNumber = 0;
- // Also records the value buffer offset.
- size_t stringOffset = 0;
+ // Also records the length buffer offset for strings.
+ size_t valueBufferOffset = 0;
for (auto i = 0; i < batches.size(); ++i) {
const auto& batch = batches[i];
if (batch->numRows() == 0) {
@@ -147,28 +161,47 @@ std::shared_ptr<GpuBufferColumnarBatch>
GpuBufferColumnarBatch::compose(
arrow::internal::CopyBitmap(batch->bufferAt(bufferIdx)->data(), 0,
batch->numRows(), dst, rowNumber);
}
- if (colType->isFixedWidth()) {
+ if (colType->isTimestamp()) {
+ const auto bufferSize = batch->numRows() * sizeof(int64_t);
+ 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() + valueBufferOffset);
+ for (auto j = 0; j < batch->numRows(); ++j) {
+ // src[0] is seconds, src[1] is nanoseconds in Timestamp.
+ dst[j] = src[j << 1] * 1'000'000'000L + src[j << 1 | 1];
+ }
+ valueBufferOffset += bufferSize;
+ } else if (colType->isBoolean()) {
+ const auto bufferSize = batch->numRows();
+ VELOX_CHECK_LE(valueBufferOffset + bufferSize, returnBuffers[bufferIdx
+ 1]->size());
+ const auto* src = batch->bufferAt(bufferIdx + 1)->data();
+ auto* dst = returnBuffers[bufferIdx + 1]->mutable_data() +
valueBufferOffset;
+ for (auto j = 0; j < batch->numRows(); ++j) {
+ dst[j] = (src[j >> 3] >> (j & 7)) & 1;
+ }
+ valueBufferOffset += bufferSize;
+ } else if (colType->isFixedWidth()) {
// The buffer is values.
const auto bufferSize = batch->bufferAt(bufferIdx + 1)->size();
- VELOX_CHECK_LE(stringOffset + bufferSize, returnBuffers[bufferIdx +
1]->size());
+ VELOX_CHECK_LE(valueBufferOffset + bufferSize, returnBuffers[bufferIdx
+ 1]->size());
memcpy(
- returnBuffers[bufferIdx + 1]->mutable_data() + stringOffset,
+ returnBuffers[bufferIdx + 1]->mutable_data() + valueBufferOffset,
batch->bufferAt(bufferIdx + 1)->data(),
bufferSize);
- stringOffset += bufferSize;
+ valueBufferOffset += bufferSize;
} else {
// String, lengths, values
memcpy(
- returnBuffers[bufferIdx + 2]->mutable_data() + stringOffset,
+ returnBuffers[bufferIdx + 2]->mutable_data() + valueBufferOffset,
batch->bufferAt(bufferIdx + 2)->data(),
batch->bufferAt(bufferIdx + 2)->size());
const auto* lengths = reinterpret_cast<const
int32_t*>(batch->bufferAt(bufferIdx + 1)->data());
auto* offsetBuffer =
reinterpret_cast<int32_t*>(returnBuffers[bufferIdx + 1]->mutable_data());
for (auto j = 0; j < batch->numRows(); ++j) {
- offsetBuffer[rowNumber + j] = stringOffset;
- stringOffset += lengths[j];
+ offsetBuffer[rowNumber + j] = valueBufferOffset;
+ valueBufferOffset += lengths[j];
}
- offsetBuffer[rowNumber + batch->numRows()] = stringOffset;
+ offsetBuffer[rowNumber + batch->numRows()] = valueBufferOffset;
}
rowNumber += batch->numRows();
}
diff --git a/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc
b/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc
deleted file mode 100644
index 1df6cb91ad..0000000000
--- a/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc
+++ /dev/null
@@ -1,72 +0,0 @@
-/*
- * 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 "VeloxGpuShuffleWriter.h"
-
-namespace gluten {
-
-// Split bool bit to bytes.
-void VeloxGpuHashShuffleWriter::splitBoolValueType(const uint8_t* srcAddr,
const std::vector<uint8_t*>& dstAddrs) {
- for (auto& pid : partitionUsed_) {
- auto dstaddr = dstAddrs[pid];
- if (dstaddr == nullptr) {
- continue;
- }
- auto dstPidBase = reinterpret_cast<uint8_t*>(dstaddr +
partitionBufferBase_[pid] * sizeof(uint8_t));
- auto pos = partition2RowOffsetBase_[pid];
- auto end = partition2RowOffsetBase_[pid + 1];
- for (; pos < end; ++pos) {
- auto rowId = rowOffset2RowId_[pos];
- // === Extract one bool value from the bit-packed source ===
- uint8_t byte = srcAddr[rowId >> 3]; // find the source byte
- uint8_t bit = (byte >> (rowId & 7)) & 0x01; // extract 0 or 1
- *dstPidBase++ = bit; // copy
- }
- }
-}
-
-// Split timestamp from int128_t to int64_t, both of them represents the
timestamp nanoseconds.
-arrow::Status VeloxGpuHashShuffleWriter::splitTimestamp(const uint8_t*
srcAddr, const std::vector<uint8_t*>& dstAddrs) {
- for (auto& pid : partitionUsed_) {
- auto dstPidBase = reinterpret_cast<int64_t*>(dstAddrs[pid] +
partitionBufferBase_[pid] * sizeof(int64_t));
- auto pos = partition2RowOffsetBase_[pid];
- auto end = partition2RowOffsetBase_[pid + 1];
- for (; pos < end; ++pos) {
- auto rowId = rowOffset2RowId_[pos];
- auto* src = reinterpret_cast<const int64_t*>(srcAddr) + rowId * 2;
- // src[0] is seconds, src[1] is nanoseconds in Timestamp.
- *dstPidBase++ = src[0] * 1'000'000'000LL + src[1];
- }
- }
- return arrow::Status::OK();
-}
-
-arrow::Result<std::shared_ptr<VeloxShuffleWriter>>
VeloxGpuHashShuffleWriter::create(
- uint32_t numPartitions,
- const std::shared_ptr<PartitionWriter>& partitionWriter,
- const std::shared_ptr<ShuffleWriterOptions>& options,
- MemoryManager* memoryManager) {
- if (auto hashOptions =
std::dynamic_pointer_cast<GpuHashShuffleWriterOptions>(options)) {
- std::shared_ptr<VeloxGpuHashShuffleWriter> res =
- std::make_shared<VeloxGpuHashShuffleWriter>(numPartitions,
partitionWriter, hashOptions, memoryManager);
- RETURN_NOT_OK(res->init());
- return res;
- }
- return arrow::Status::Invalid("Error casting ShuffleWriterOptions to
GpuHashShuffleWriterOptions. ");
-}
-
-} // namespace gluten
diff --git a/cpp/velox/shuffle/VeloxGpuShuffleWriter.h
b/cpp/velox/shuffle/VeloxGpuShuffleWriter.h
deleted file mode 100644
index e6f0e1a358..0000000000
--- a/cpp/velox/shuffle/VeloxGpuShuffleWriter.h
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * 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.
- */
-
-#pragma once
-
-#include "VeloxHashShuffleWriter.h"
-
-namespace gluten {
-
-class VeloxGpuHashShuffleWriter : public VeloxHashShuffleWriter {
- public:
- static arrow::Result<std::shared_ptr<VeloxShuffleWriter>> create(
- uint32_t numPartitions,
- const std::shared_ptr<PartitionWriter>& partitionWriter,
- const std::shared_ptr<ShuffleWriterOptions>& options,
- MemoryManager* memoryManager);
-
- VeloxGpuHashShuffleWriter(
- uint32_t numPartitions,
- const std::shared_ptr<PartitionWriter>& partitionWriter,
- const std::shared_ptr<GpuHashShuffleWriterOptions>& options,
- MemoryManager* memoryManager)
- : VeloxHashShuffleWriter(numPartitions, partitionWriter, options,
memoryManager) {}
-
- private:
- // Split the bool to byte.
- void splitBoolValueType(const uint8_t* srcAddr, const std::vector<uint8_t*>&
dstAddrs) override;
-
- uint64_t valueBufferSizeForBool(uint32_t newSize) override {
- return newSize;
- }
-
- bool boolIsBit() override {
- return false;
- }
-
- arrow::Status splitTimestamp(const uint8_t* srcAddr, const
std::vector<uint8_t*>& dstAddrs) override;
-
- uint64_t valueBufferSizeForTimestamp(uint32_t newSize) override {
- return sizeof(int64_t) * newSize;
- }
-};
-} // namespace gluten
diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.cc
b/cpp/velox/shuffle/VeloxShuffleWriter.cc
index 25ed099014..a83e3a1947 100644
--- a/cpp/velox/shuffle/VeloxShuffleWriter.cc
+++ b/cpp/velox/shuffle/VeloxShuffleWriter.cc
@@ -20,10 +20,6 @@
#include "shuffle/VeloxRssSortShuffleWriter.h"
#include "shuffle/VeloxSortShuffleWriter.h"
-#ifdef GLUTEN_ENABLE_GPU
-#include "VeloxGpuShuffleWriter.h"
-#endif
-
namespace gluten {
arrow::Result<std::shared_ptr<VeloxShuffleWriter>> VeloxShuffleWriter::create(
@@ -35,15 +31,12 @@ arrow::Result<std::shared_ptr<VeloxShuffleWriter>>
VeloxShuffleWriter::create(
std::shared_ptr<VeloxShuffleWriter> shuffleWriter;
switch (type) {
case ShuffleWriterType::kHashShuffle:
+ case ShuffleWriterType::kGpuHashShuffle:
return VeloxHashShuffleWriter::create(numPartitions,
std::move(partitionWriter), options, memoryManager);
case ShuffleWriterType::kSortShuffle:
return VeloxSortShuffleWriter::create(numPartitions,
std::move(partitionWriter), options, memoryManager);
case ShuffleWriterType::kRssSortShuffle:
return VeloxRssSortShuffleWriter::create(numPartitions,
std::move(partitionWriter), options, memoryManager);
-#ifdef GLUTEN_ENABLE_GPU
- case ShuffleWriterType::kGpuHashShuffle:
- return VeloxGpuHashShuffleWriter::create(numPartitions,
std::move(partitionWriter), options, memoryManager);
-#endif
default:
return arrow::Status::Invalid("Unsupported shuffle writer type: ",
typeToString(type));
}
diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
index cc846eb4eb..896fa9e626 100644
--- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
+++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
@@ -20,8 +20,7 @@
#include "config/GlutenConfig.h"
#include "memory/GpuBufferColumnarBatch.h"
-#include "shuffle/VeloxGpuShuffleWriter.h"
-#include "shuffle/VeloxHashShuffleWriter.h"
+#include "shuffle/VeloxShuffleWriter.h"
#include "tests/VeloxShuffleWriterTestBase.h"
#include "tests/utils/TestAllocationListener.h"
#include "tests/utils/TestStreamReader.h"
@@ -227,7 +226,7 @@ class GpuVeloxShuffleWriterTest : public
::testing::TestWithParam<GpuShuffleTest
const auto& params = GetParam();
switch (params.shuffleWriterType) {
case ShuffleWriterType::kGpuHashShuffle: {
- auto hashOptions = std::make_shared<GpuHashShuffleWriterOptions>();
+ auto hashOptions = std::make_shared<HashShuffleWriterOptions>();
hashOptions->splitBufferSize = splitBufferSize;
options = hashOptions;
} break;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
