This is an automated email from the ASF dual-hosted git repository.
zhouyuan 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 f5764eb1bb [GLUTEN-12809][VL] Fix short decimal / timestamp width in
hash shuffle partition buffer sizing (#12810)
f5764eb1bb is described below
commit f5764eb1bb59b791ad022f25ed3b71acf63d26da
Author: Yuan <[email protected]>
AuthorDate: Fri Aug 21 17:42:11 2026 +0100
[GLUTEN-12809][VL] Fix short decimal / timestamp width in hash shuffle
partition buffer sizing (#12810)
---
cpp/velox/shuffle/VeloxHashShuffleWriter.cc | 9 +-
cpp/velox/shuffle/VeloxHashShuffleWriter.h | 7 +
cpp/velox/tests/CMakeLists.txt | 3 +
.../VeloxHashShuffleWriterBufferSizingTest.cc | 151 +++++++++++++++++++++
4 files changed, 167 insertions(+), 3 deletions(-)
diff --git a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
index 3cec6e4d41..14c55861bc 100644
--- a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
+++ b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
@@ -893,9 +893,12 @@ inline bool
VeloxHashShuffleWriter::beyondThreshold(uint32_t partitionId, uint32
void VeloxHashShuffleWriter::calculateSimpleColumnBytes() {
fixedWidthBufferBytes_ = 0;
for (size_t col = 0; col < fixedWidthColumnCount_; ++col) {
- auto colIdx = simpleColumnIndices_[col];
- // `bool(1) >> 3` gets 0, so +7
- fixedWidthBufferBytes_ +=
((arrow::bit_width(arrowColumnTypes_[colIdx]->id()) + 7) >> 3);
+ // Reuse the same per-column sizing as the actual buffer allocation,
otherwise this estimate can
+ // drift from it: `arrow::bit_width` mis-counts the types whose Arrow bit
width differs from the
+ // width the partition buffer allocates, i.e. short decimal (allocated as
int64, 8 bytes not 16)
+ // and timestamp (allocated as int128, 16 bytes not 8). Note bool is still
rounded up to one byte
+ // per row.
+ fixedWidthBufferBytes_ += valueBufferSizeForFixedWidthArray(col, 1);
}
fixedWidthBufferBytes_ += kSizeOfStringLength * binaryColumnIndices_.size();
}
diff --git a/cpp/velox/shuffle/VeloxHashShuffleWriter.h
b/cpp/velox/shuffle/VeloxHashShuffleWriter.h
index ea7a659f2e..a0cf4a1544 100644
--- a/cpp/velox/shuffle/VeloxHashShuffleWriter.h
+++ b/cpp/velox/shuffle/VeloxHashShuffleWriter.h
@@ -142,6 +142,13 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter {
// For test only.
void setPartitionBufferSize(uint32_t newSize) override;
+ // For test only. Per-row byte estimate of the fixed-width part of the
partition buffers,
+ // computed by calculateSimpleColumnBytes(). Must stay consistent with the
sizes actually
+ // allocated by valueBufferSizeForFixedWidthArray().
+ uint32_t fixedWidthBufferBytes() const {
+ return fixedWidthBufferBytes_;
+ }
+
// Read-only counters of incoming column-vector encodings observed at the
// entry of every `write(...)` call, BEFORE flattening. Mirrors the layout
// of `cpuWallTimingList_` (always-on counter, logged via `stat()` when
diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt
index 87331de818..136884503a 100644
--- a/cpp/velox/tests/CMakeLists.txt
+++ b/cpp/velox/tests/CMakeLists.txt
@@ -102,6 +102,9 @@ add_velox_test(velox_shuffle_writer_test SOURCES
VeloxShuffleWriterTest.cc)
add_velox_test(velox_hash_shuffle_writer_input_encoding_test SOURCES
VeloxHashShuffleWriterInputEncodingTest.cc)
+add_velox_test(velox_hash_shuffle_writer_buffer_sizing_test SOURCES
+ VeloxHashShuffleWriterBufferSizingTest.cc)
+
add_velox_test(velox_shuffle_writer_spill_test SOURCES
VeloxShuffleWriterSpillTest.cc)
diff --git a/cpp/velox/tests/VeloxHashShuffleWriterBufferSizingTest.cc
b/cpp/velox/tests/VeloxHashShuffleWriterBufferSizingTest.cc
new file mode 100644
index 0000000000..46e3107d03
--- /dev/null
+++ b/cpp/velox/tests/VeloxHashShuffleWriterBufferSizingTest.cc
@@ -0,0 +1,151 @@
+/*
+ * 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 <gtest/gtest.h>
+
+#include "shuffle/VeloxHashShuffleWriter.h"
+
+#include "VeloxShuffleWriterTestBase.h"
+#include "utils/Macros.h"
+#include "utils/TestUtils.h"
+#include "velox/vector/tests/utils/VectorTestBase.h"
+
+namespace gluten {
+
+namespace {
+
+std::shared_ptr<PartitionWriter> makeLocalPartitionWriter(
+ uint32_t numPartitions,
+ const std::string& dataFile,
+ const std::vector<std::string>& localDirs) {
+ GLUTEN_ASSIGN_OR_THROW(auto codec,
arrow::util::Codec::Create(arrow::Compression::LZ4_FRAME));
+ auto options = std::make_shared<LocalPartitionWriterOptions>();
+ return std::make_shared<LocalPartitionWriter>(
+ numPartitions, std::move(codec), getDefaultMemoryManager(), options,
dataFile, localDirs);
+}
+
+} // namespace
+
+// Verifies that calculateSimpleColumnBytes() estimates each fixed-width
column with the same
+// per-row width the partition buffers are actually allocated with (see
+// valueBufferSizeForFixedWidthArray). Short decimal is stored as int64 (8
bytes, not the 16 of
+// arrow::bit_width(Decimal128)), and timestamp is stored as int128 (16 bytes,
not the 8 of
+// arrow::bit_width(Timestamp)).
+class HashShuffleWriterBufferSizingTest : public ::testing::Test, public
VeloxShuffleWriterTestBase {
+ protected:
+ static void SetUpTestSuite() {
+ setUpVeloxBackend();
+ }
+
+ static void TearDownTestSuite() {
+ tearDownVeloxBackend();
+ }
+
+ void SetUp() override {
+ GLUTEN_THROW_NOT_OK(setLocalDirsAndDataFile());
+ }
+
+ std::shared_ptr<VeloxHashShuffleWriter> createWriter(uint32_t numPartitions)
{
+ auto options = std::make_shared<HashShuffleWriterOptions>();
+ options->partitioning = Partitioning::kHash;
+ options->splitBufferSize = 4096;
+
+ auto partitionWriter = makeLocalPartitionWriter(numPartitions, dataFile_,
localDirs_);
+
+ GLUTEN_ASSIGN_OR_THROW(
+ auto base,
+ VeloxShuffleWriter::create(
+ ShuffleWriterType::kHashShuffle, numPartitions, partitionWriter,
options, getDefaultMemoryManager()));
+ return std::dynamic_pointer_cast<VeloxHashShuffleWriter>(base);
+ }
+
+ // Writes a single batch whose first child is the hash partition key
(stripped before the
+ // writer initializes column types), then returns the writer's per-row
fixed-width estimate
+ // for the remaining data columns.
+ uint32_t bytesPerRowFor(std::vector<facebook::velox::VectorPtr>
dataChildren) {
+ auto writer = createWriter(2);
+ EXPECT_NE(writer, nullptr);
+
+ std::vector<facebook::velox::VectorPtr> children;
+ children.push_back(makeFlatVector<int32_t>({0, 1}));
+ children.insert(children.end(), dataChildren.begin(), dataChildren.end());
+ auto rv = makeRowVector(children);
+
+ std::shared_ptr<ColumnarBatch> cb =
std::make_shared<VeloxColumnarBatch>(rv);
+ EXPECT_TRUE(writer->write(cb, ShuffleWriter::kMinMemLimit).ok());
+ return writer->fixedWidthBufferBytes();
+ }
+};
+
+// Short decimal is split and allocated as int64: 8 bytes per row, not the 16
bytes implied by
+// arrow::bit_width(Decimal128Type).
+TEST_F(HashShuffleWriterBufferSizingTest, shortDecimalCountedAsInt64) {
+ auto bytesPerRow = bytesPerRowFor({
+ makeFlatVector<int64_t>({232, 34567235}, facebook::velox::DECIMAL(12,
4)),
+ });
+ EXPECT_EQ(bytesPerRow, 8);
+}
+
+// Long decimal really is 16 bytes per row.
+TEST_F(HashShuffleWriterBufferSizingTest, longDecimalCountedAsInt128) {
+ auto bytesPerRow = bytesPerRowFor({
+ makeFlatVector<facebook::velox::int128_t>({232, 34567235},
facebook::velox::DECIMAL(20, 4)),
+ });
+ EXPECT_EQ(bytesPerRow, 16);
+}
+
+// Timestamp is split and allocated as int128 (velox Timestamp is 16 bytes),
not the 8 bytes
+// implied by arrow::bit_width(TimestampType).
+TEST_F(HashShuffleWriterBufferSizingTest, timestampCountedAsInt128) {
+ auto bytesPerRow = bytesPerRowFor({
+
makeFlatVector<facebook::velox::Timestamp>({facebook::velox::Timestamp(1, 0),
facebook::velox::Timestamp(2, 0)}),
+ });
+ EXPECT_EQ(bytesPerRow,
static_cast<uint32_t>(sizeof(facebook::velox::Timestamp)));
+}
+
+// Mixed schema: the estimate is the sum of the widths the buffers are
allocated with, plus the
+// length-buffer width for each binary column.
+TEST_F(HashShuffleWriterBufferSizingTest, mixedSchema) {
+ auto bytesPerRow = bytesPerRowFor({
+ makeFlatVector<bool>({true, false}),
+ makeFlatVector<int8_t>({1, 2}),
+ makeFlatVector<int32_t>({1, 2}),
+ makeFlatVector<int64_t>({1, 2}),
+ makeFlatVector<double>({1.0, 2.0}),
+ makeFlatVector<int64_t>({232, 34567235}, facebook::velox::DECIMAL(12,
4)),
+ makeFlatVector<facebook::velox::int128_t>({232, 34567235},
facebook::velox::DECIMAL(20, 4)),
+
makeFlatVector<facebook::velox::Timestamp>({facebook::velox::Timestamp(1, 0),
facebook::velox::Timestamp(2, 0)}),
+ makeFlatVector<facebook::velox::StringView>({"a", "bb"}),
+ });
+ uint32_t expected = 1 // bool, rounded up to one byte per row
+ + 1 // int8
+ + 4 // int32
+ + 8 // int64
+ + 8 // double
+ + 8 // short decimal, stored as int64
+ + 16 // long decimal
+ + 16 // timestamp, stored as int128
+ + kSizeOfStringLength; // varchar length buffer
+ EXPECT_EQ(bytesPerRow, expected);
+}
+
+} // namespace gluten
+
+int main(int argc, char** argv) {
+ testing::InitGoogleTest(&argc, argv);
+ return RUN_ALL_TESTS();
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
