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 14683bb930 [VL] Add timestamp to shuffle test (#12503)
14683bb930 is described below
commit 14683bb9303ceb5f962e0615434413f70ea88ca2
Author: Rong Ma <[email protected]>
AuthorDate: Tue Jul 14 13:30:56 2026 +0100
[VL] Add timestamp to shuffle test (#12503)
---
cpp/velox/shuffle/VeloxHashShuffleWriter.cc | 5 +-
cpp/velox/shuffle/VeloxRssSortShuffleWriter.cc | 2 +
cpp/velox/shuffle/VeloxRssSortShuffleWriter.h | 2 -
cpp/velox/shuffle/VeloxShuffleWriter.h | 2 +
cpp/velox/shuffle/VeloxSortShuffleWriter.cc | 4 +-
cpp/velox/shuffle/VeloxSortShuffleWriter.h | 1 -
cpp/velox/tests/VeloxGpuShuffleWriterTest.cc | 135 -----------------------
cpp/velox/tests/VeloxRssSortShuffleWriterTest.cc | 120 +++++++++++---------
cpp/velox/tests/VeloxShuffleWriterTest.cc | 27 +----
cpp/velox/tests/VeloxShuffleWriterTestBase.h | 58 +++++++++-
10 files changed, 138 insertions(+), 218 deletions(-)
diff --git a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
index 157d421644..3cec6e4d41 100644
--- a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
+++ b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc
@@ -873,10 +873,13 @@ arrow::Status
VeloxHashShuffleWriter::initColumnTypes(const facebook::velox::Row
}
arrow::Status VeloxHashShuffleWriter::initFromRowVector(const
facebook::velox::RowVector& rv) {
- if (veloxColumnTypes_.empty()) {
+ if (!rowType_) {
RETURN_NOT_OK(initColumnTypes(rv));
RETURN_NOT_OK(initPartitions());
calculateSimpleColumnBytes();
+ rowType_ = rv.rowType();
+ } else {
+ VELOX_CHECK(rowType_->equivalent(*rv.rowType()));
}
return arrow::Status::OK();
}
diff --git a/cpp/velox/shuffle/VeloxRssSortShuffleWriter.cc
b/cpp/velox/shuffle/VeloxRssSortShuffleWriter.cc
index 0f9297ac5f..fd8af57aa7 100644
--- a/cpp/velox/shuffle/VeloxRssSortShuffleWriter.cc
+++ b/cpp/velox/shuffle/VeloxRssSortShuffleWriter.cc
@@ -219,6 +219,8 @@ arrow::Status
VeloxRssSortShuffleWriter::initFromRowVector(const facebook::velox
serdeOptions_ = {false, compressionKind_};
batch_ =
std::make_unique<facebook::velox::VectorStreamGroup>(veloxPool_.get(),
serde_.get());
batch_->createStreamTree(rowType_, splitBufferSize_, &serdeOptions_);
+ } else {
+ VELOX_CHECK(rowType_->equivalent(*rv.rowType()));
}
return arrow::Status::OK();
}
diff --git a/cpp/velox/shuffle/VeloxRssSortShuffleWriter.h
b/cpp/velox/shuffle/VeloxRssSortShuffleWriter.h
index dbb5051bbd..e2914f2322 100644
--- a/cpp/velox/shuffle/VeloxRssSortShuffleWriter.h
+++ b/cpp/velox/shuffle/VeloxRssSortShuffleWriter.h
@@ -103,8 +103,6 @@ class VeloxRssSortShuffleWriter final : public
VeloxShuffleWriter {
int64_t sortBufferMaxSize_;
facebook::velox::common::CompressionKind compressionKind_;
- facebook::velox::RowTypePtr rowType_;
-
std::unique_ptr<facebook::velox::VectorStreamGroup> batch_;
std::unique_ptr<BufferOutputStream> bufferOutputStream_;
diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.h
b/cpp/velox/shuffle/VeloxShuffleWriter.h
index 45276e3cd4..5a478a7755 100644
--- a/cpp/velox/shuffle/VeloxShuffleWriter.h
+++ b/cpp/velox/shuffle/VeloxShuffleWriter.h
@@ -145,6 +145,8 @@ class VeloxShuffleWriter : public ShuffleWriter {
int32_t maxBatchSize_{0};
+ facebook::velox::RowTypePtr rowType_{nullptr};
+
enum EvictState { kEvictable, kUnevictable };
// stat
diff --git a/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
b/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
index 7a4066153b..0a31dd3d7d 100644
--- a/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
+++ b/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
@@ -121,8 +121,10 @@ arrow::Status VeloxSortShuffleWriter::init() {
void VeloxSortShuffleWriter::initRowType(const facebook::velox::RowVectorPtr&
rv) {
if (UNLIKELY(!rowType_)) {
- rowType_ = facebook::velox::asRowType(rv->type());
+ rowType_ = rv->rowType();
fixedRowSize_ = facebook::velox::row::CompactRow::fixedRowSize(rowType_);
+ } else {
+ VELOX_CHECK(rowType_->equivalent(*rv->rowType()));
}
}
diff --git a/cpp/velox/shuffle/VeloxSortShuffleWriter.h
b/cpp/velox/shuffle/VeloxSortShuffleWriter.h
index a30f0edb83..59541d2496 100644
--- a/cpp/velox/shuffle/VeloxSortShuffleWriter.h
+++ b/cpp/velox/shuffle/VeloxSortShuffleWriter.h
@@ -122,7 +122,6 @@ class VeloxSortShuffleWriter final : public
VeloxShuffleWriter {
// Updated for each input RowVector.
std::vector<uint32_t> row2Partition_;
- std::shared_ptr<const facebook::velox::RowType> rowType_;
std::optional<int32_t> fixedRowSize_;
std::vector<RowSizeType> rowSize_;
std::vector<uint64_t> rowSizePrefixSum_;
diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
index 7853ad418a..cc846eb4eb 100644
--- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
+++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
@@ -582,141 +582,6 @@ TEST_P(GpuRoundRobinPartitioningShuffleWriterTest,
roundRobin) {
testShuffleRoundTrip(*shuffleWriter, {inputVector1_, inputVector2_,
inputVector1_}, 2, {blockPid1, blockPid2});
}
-TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, preAllocForceRealloc) {
- if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) {
- return;
- }
-
- auto shuffleWriterOptions = std::make_shared<HashShuffleWriterOptions>();
- shuffleWriterOptions->splitBufferSize = 4096;
- shuffleWriterOptions->splitBufferReallocThreshold = 0; // Force re-alloc on
buffer size changed.
- auto shuffleWriter = createShuffleWriter(2, shuffleWriterOptions);
-
- // First spilt no null.
- auto inputNoNull = inputVectorNoNull_;
-
- // Second split has null. Continue filling current partition buffers.
- std::vector<VectorPtr> intHasNull = {
- makeNullableFlatVector<int8_t>({std::nullopt, 1}),
- makeNullableFlatVector<int8_t>({std::nullopt, -1}),
- makeNullableFlatVector<int32_t>({std::nullopt, 100}),
- makeNullableFlatVector<int64_t>({0, 1}),
- makeNullableFlatVector<float>({0, 0.142857}),
- makeNullableFlatVector<bool>({false, true}),
- makeNullableFlatVector<StringView>({"", "alice"}),
- makeNullableFlatVector<StringView>({"alice", ""}),
- };
-
- auto inputHasNull = makeRowVector(intHasNull);
- // Split first input no null.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull));
- // Split second input, continue filling but update null.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputHasNull));
-
- // Split first input again.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull));
- // Check when buffer is full, evict current buffers and reuse.
- auto cachedPayloadSize = shuffleWriter->cachedPayloadSize();
- auto partitionBufferBeforeEvict = shuffleWriter->partitionBufferSize();
- int64_t evicted;
- ASSERT_NOT_OK(shuffleWriter->reclaimFixedSize(cachedPayloadSize, &evicted));
- // Check only cached data being spilled.
- ASSERT_EQ(evicted, cachedPayloadSize);
- VELOX_CHECK_EQ(shuffleWriter->partitionBufferSize(),
partitionBufferBeforeEvict);
-
- // Split more data with null. New buffer size is larger.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_));
-
- // Split more data with null. New buffer size is smaller.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector2_));
-
- // Split more data with null. New buffer size is larger and current data is
preserved.
- // Evict cached data first.
-
ASSERT_NOT_OK(shuffleWriter->reclaimFixedSize(shuffleWriter->cachedPayloadSize(),
&evicted));
- // Set a large buffer size.
- shuffleWriter->setPartitionBufferSize(100);
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_));
- // No data got evicted so the cached size is 0.
- ASSERT_EQ(shuffleWriter->cachedPayloadSize(), 0);
-
- ASSERT_NOT_OK(shuffleWriter->stop());
-}
-
-TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, preAllocForceReuse) {
- if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) {
- return;
- }
- auto shuffleWriterOptions = std::make_shared<HashShuffleWriterOptions>();
- shuffleWriterOptions->splitBufferSize = 4096;
- shuffleWriterOptions->splitBufferReallocThreshold = 1; // Force reuse on
buffer size changed.
- auto shuffleWriter = createShuffleWriter(2, shuffleWriterOptions);
-
- // First spilt no null.
- auto inputNoNull = inputVectorNoNull_;
- // Second split has null int, null string and non-null string,
- auto inputFixedWidthHasNull = inputVector1_;
- // Third split has null string.
- std::vector<VectorPtr> stringHasNull = {
- makeNullableFlatVector<int8_t>({0, 1}),
- makeNullableFlatVector<int8_t>({0, -1}),
- makeNullableFlatVector<int32_t>({0, 100}),
- makeNullableFlatVector<int64_t>({0, 1}),
- makeNullableFlatVector<float>({0, 0.142857}),
- makeNullableFlatVector<bool>({false, true}),
- makeNullableFlatVector<StringView>({std::nullopt, std::nullopt}),
- makeNullableFlatVector<StringView>({std::nullopt, std::nullopt}),
- };
- auto inputStringHasNull = makeRowVector(stringHasNull);
-
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull));
- // Split more data with null. Already filled + to be filled > buffer size,
Buffer is resized larger.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputFixedWidthHasNull));
- // Split more data with null. Already filled + to be filled > buffer size,
newSize is smaller so buffer is not
- // resized.
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputStringHasNull));
-
- ASSERT_NOT_OK(shuffleWriter->stop());
-}
-
-TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, spillVerifyResult) {
- if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) {
- return;
- }
-
- auto shuffleWriter = createShuffleWriter(2);
-
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_));
-
- // Clear buffers and evict payloads and cache.
- for (auto pid : {0, 1}) {
- ASSERT_NOT_OK(shuffleWriter->evictPartitionBuffers(pid, true));
- }
-
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_));
-
- // Evict all payloads and spill.
- int64_t evicted;
- auto cachedPayloadSize = shuffleWriter->cachedPayloadSize();
- auto partitionBufferSize = shuffleWriter->partitionBufferSize();
- ASSERT_NOT_OK(shuffleWriter->reclaimFixedSize(cachedPayloadSize +
partitionBufferSize, &evicted));
-
- ASSERT_EQ(evicted, cachedPayloadSize + partitionBufferSize);
-
- // No more cached payloads after spill.
- ASSERT_EQ(shuffleWriter->cachedPayloadSize(), 0);
- ASSERT_EQ(shuffleWriter->partitionBufferSize(), 0);
-
- ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_));
-
- auto blockPid1 =
- takeRows({inputVector1_, inputVector1_, inputVector1_}, {{0, 2, 4, 6,
8}, {0, 2, 4, 6, 8}, {0, 2, 4, 6, 8}});
- auto blockPid2 =
- takeRows({inputVector1_, inputVector1_, inputVector1_}, {{1, 3, 5, 7,
9}, {1, 3, 5, 7, 9}, {1, 3, 5, 7, 9}});
-
- // Stop and verify.
- shuffleWriteReadMultiBlocks(*shuffleWriter, 2, {blockPid1, blockPid2});
-}
-
INSTANTIATE_TEST_SUITE_P(
SinglePartitioningShuffleWriterGroup,
GpuSinglePartitioningShuffleWriterTest,
diff --git a/cpp/velox/tests/VeloxRssSortShuffleWriterTest.cc
b/cpp/velox/tests/VeloxRssSortShuffleWriterTest.cc
index 32fe57b241..129097106c 100644
--- a/cpp/velox/tests/VeloxRssSortShuffleWriterTest.cc
+++ b/cpp/velox/tests/VeloxRssSortShuffleWriterTest.cc
@@ -44,6 +44,7 @@ class VeloxRssSortShuffleWriterTest : public
VeloxShuffleWriterTestBase, public
void SetUp() override {
setUpTestData();
+ strBuffer_ = AlignedBuffer::allocate<char>(200, pool());
}
std::shared_ptr<VeloxShuffleWriter> createShuffleWriter(uint32_t
numPartitions) {
@@ -64,77 +65,96 @@ class VeloxRssSortShuffleWriterTest : public
VeloxShuffleWriterTestBase, public
numPartitions, std::move(partitionWriter),
std::move(writerOptions), getDefaultMemoryManager()));
return shuffleWriter;
}
+
+ std::shared_ptr<FlatVector<StringView>> makeNonSharedStringVector() const {
+ auto strBuffer = AlignedBuffer::allocate<char>(200, pool());
+ auto vector = BaseVector::create<FlatVector<StringView>>(VARCHAR(), 100,
pool());
+ vector->setStringBuffers({std::move(strBuffer)});
+ return vector;
+ }
+
+ std::shared_ptr<FlatVector<StringView>> makeSharedStringVector() {
+ auto vector = BaseVector::create<FlatVector<StringView>>(VARCHAR(), 100,
pool());
+ vector->setStringBuffers({strBuffer_});
+ return vector;
+ }
+
+ BufferPtr strBuffer_;
+
+ static constexpr int64_t kMemLimit = std::numeric_limits<int64_t>::max();
};
TEST_F(VeloxRssSortShuffleWriterTest, calculateBatchesSize) {
auto shuffleWriter =
std::dynamic_pointer_cast<VeloxRssSortShuffleWriter>(createShuffleWriter(10));
- // Do not trigger resetBatches by shuffle writer.
- const int64_t memLimit = INT64_MAX;
- // Shared string buffer in FlatVector<StringView>.
- BufferPtr strBuffer = AlignedBuffer::allocate<char>(200, pool());
- auto vector1 = BaseVector::create<FlatVector<StringView>>(VARCHAR(), 100,
pool());
- vector1->setStringBuffers({strBuffer});
- auto vector2 = BaseVector::create<FlatVector<StringView>>(VARCHAR(), 100,
pool());
- vector2->setStringBuffers({strBuffer});
- auto vector3 = BaseVector::create<FlatVector<StringView>>(VARCHAR(), 100,
pool());
- vector3->setStringBuffers({strBuffer});
- auto vector4 = BaseVector::create<FlatVector<int64_t>>(INTEGER(), 100,
pool());
-
- auto rowVector1 = makeRowVector({vector1, vector2});
- auto rowVector2 = makeRowVector({vector3, vector4});
- std::shared_ptr<ColumnarBatch> cb1 =
std::make_shared<VeloxColumnarBatch>(rowVector1);
- std::shared_ptr<ColumnarBatch> cb2 =
std::make_shared<VeloxColumnarBatch>(rowVector2);
-
- ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit));
- ASSERT_NOT_OK(shuffleWriter->write(cb2, memLimit));
- auto expectedSize = rowVector1->retainedSize() + rowVector2->retainedSize()
- strBuffer->capacity() * 2;
+ auto rowVector1 = makeRowVector({makeSharedStringVector(),
makeSharedStringVector()});
+ auto rowVector2 = makeRowVector({makeSharedStringVector(),
makeNonSharedStringVector()});
+ auto cb1 = std::make_shared<VeloxColumnarBatch>(rowVector1);
+ auto cb2 = std::make_shared<VeloxColumnarBatch>(rowVector2);
+
+ ASSERT_NOT_OK(shuffleWriter->write(cb1, kMemLimit));
+ ASSERT_NOT_OK(shuffleWriter->write(cb2, kMemLimit));
+ auto expectedSize = rowVector1->retainedSize() + rowVector2->retainedSize()
- strBuffer_->capacity() * 2;
EXPECT_EQ(expectedSize, shuffleWriter->getInputColumnBytes());
- shuffleWriter->resetBatches();
+}
+
+TEST_F(VeloxRssSortShuffleWriterTest, sharedStringInArray) {
+ auto shuffleWriter =
std::dynamic_pointer_cast<VeloxRssSortShuffleWriter>(createShuffleWriter(10));
+
+ // Shared string buffer in FlatVector<StringView>.
+ auto vector = makeSharedStringVector();
// Shared string buffer in ArrayVector.
- BufferPtr offsets = allocateOffsets(1, vector1->pool());
- BufferPtr sizes = allocateOffsets(1, vector1->pool());
- sizes->asMutable<vector_size_t>()[0] = vector1->size();
+ BufferPtr offsets = allocateOffsets(1, vector->pool());
+ BufferPtr sizes = allocateOffsets(1, vector->pool());
+ sizes->asMutable<vector_size_t>()[0] = vector->size();
auto arrayVector =
- std::make_shared<facebook::velox::ArrayVector>(pool(), ARRAY(VARCHAR()),
nullptr, 1, offsets, sizes, vector1);
- auto rowVector3 = makeRowVector({arrayVector, vector1});
- cb1 = std::make_shared<VeloxColumnarBatch>(rowVector3);
- ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit));
- expectedSize = rowVector3->retainedSize() - strBuffer->capacity();
+ std::make_shared<facebook::velox::ArrayVector>(pool(), ARRAY(VARCHAR()),
nullptr, 1, offsets, sizes, vector);
+ auto rowVector = makeRowVector({arrayVector, vector});
+ auto cb1 = std::make_shared<VeloxColumnarBatch>(rowVector);
+ ASSERT_NOT_OK(shuffleWriter->write(cb1, kMemLimit));
+ auto expectedSize = rowVector->retainedSize() - strBuffer_->capacity();
EXPECT_EQ(expectedSize, shuffleWriter->getInputColumnBytes());
- shuffleWriter->resetBatches();
+}
+
+TEST_F(VeloxRssSortShuffleWriterTest, sharedStringInMap) {
+ auto shuffleWriter =
std::dynamic_pointer_cast<VeloxRssSortShuffleWriter>(createShuffleWriter(10));
// Shared string buffer in MapVector.
- auto keys = vector1;
- auto values = vector2;
+ auto keys = makeSharedStringVector();
+ auto values = makeSharedStringVector();
auto mapVector = makeMapVector({0, 10, 20, 50}, keys, values);
- auto rowVector4 = makeRowVector({mapVector, vector3});
- cb1 = std::make_shared<VeloxColumnarBatch>(rowVector4);
- ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit));
- expectedSize = rowVector4->retainedSize() - strBuffer->capacity() * 2;
+ auto rowVector = makeRowVector({mapVector, makeSharedStringVector()});
+ auto cb = std::make_shared<VeloxColumnarBatch>(rowVector);
+ ASSERT_NOT_OK(shuffleWriter->write(cb, kMemLimit));
+ auto expectedSize = rowVector->retainedSize() - strBuffer_->capacity() * 2;
EXPECT_EQ(expectedSize, shuffleWriter->getInputColumnBytes());
- shuffleWriter->resetBatches();
+}
+
+TEST_F(VeloxRssSortShuffleWriterTest, sharedStringInRowVector) {
+ auto shuffleWriter =
std::dynamic_pointer_cast<VeloxRssSortShuffleWriter>(createShuffleWriter(10));
// Shared string buffer in RowVector.
- auto rowVector5 = makeRowVector({rowVector1, vector3});
- cb1 = std::make_shared<VeloxColumnarBatch>(rowVector5);
- ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit));
- expectedSize = rowVector5->retainedSize() - strBuffer->capacity() * 2;
+ auto rowVectorInner = makeRowVector({makeSharedStringVector(),
makeSharedStringVector()});
+ auto rowVectorOuter = makeRowVector({rowVectorInner,
makeSharedStringVector()});
+ auto cb = std::make_shared<VeloxColumnarBatch>(rowVectorOuter);
+ ASSERT_NOT_OK(shuffleWriter->write(cb, kMemLimit));
+ auto expectedSize = rowVectorOuter->retainedSize() - strBuffer_->capacity()
* 2;
EXPECT_EQ(expectedSize, shuffleWriter->getInputColumnBytes());
- shuffleWriter->resetBatches();
+}
+
+TEST_F(VeloxRssSortShuffleWriterTest, sharedStringInDictionary) {
+ auto shuffleWriter =
std::dynamic_pointer_cast<VeloxRssSortShuffleWriter>(createShuffleWriter(10));
// Vector is not flatten.
+ auto vector = makeSharedStringVector();
auto dictionaryVector = BaseVector::wrapInDictionary(
- BufferPtr(nullptr),
- makeIndices(vector1->size(), [](vector_size_t row) { return row; }),
- vector1->size(),
- vector1);
- auto rowVector6 = makeRowVector({dictionaryVector});
- cb1 = std::make_shared<VeloxColumnarBatch>(rowVector6);
- ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit));
- EXPECT_EQ(rowVector6->retainedSize(), shuffleWriter->getInputColumnBytes());
+ BufferPtr(nullptr), makeIndices(vector->size(), [](vector_size_t row) {
return row; }), vector->size(), vector);
+ auto rowVector = makeRowVector({dictionaryVector});
+ auto cb = std::make_shared<VeloxColumnarBatch>(rowVector);
+ ASSERT_NOT_OK(shuffleWriter->write(cb, kMemLimit));
+ EXPECT_EQ(rowVector->retainedSize(), shuffleWriter->getInputColumnBytes());
}
} // namespace gluten
\ No newline at end of file
diff --git a/cpp/velox/tests/VeloxShuffleWriterTest.cc
b/cpp/velox/tests/VeloxShuffleWriterTest.cc
index 999d5e7435..27bc342946 100644
--- a/cpp/velox/tests/VeloxShuffleWriterTest.cc
+++ b/cpp/velox/tests/VeloxShuffleWriterTest.cc
@@ -908,20 +908,9 @@ TEST_P(RoundRobinPartitioningShuffleWriterTest,
preAllocForceRealloc) {
// First spilt no null.
auto inputNoNull = inputVectorNoNull_;
-
// Second split has null. Continue filling current partition buffers.
- std::vector<VectorPtr> intHasNull = {
- makeNullableFlatVector<int8_t>({std::nullopt, 1}),
- makeNullableFlatVector<int8_t>({std::nullopt, -1}),
- makeNullableFlatVector<int32_t>({std::nullopt, 100}),
- makeNullableFlatVector<int64_t>({0, 1}),
- makeNullableFlatVector<float>({0, 0.142857}),
- makeNullableFlatVector<bool>({false, true}),
- makeNullableFlatVector<StringView>({"", "alice"}),
- makeNullableFlatVector<StringView>({"alice", ""}),
- };
-
- auto inputHasNull = makeRowVector(intHasNull);
+ auto inputHasNull = inputVectorIntHasNull_;
+
// Split first input no null.
ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull));
// Split second input, continue filling but update null.
@@ -970,17 +959,7 @@ TEST_P(RoundRobinPartitioningShuffleWriterTest,
preAllocForceReuse) {
// Second split has null int, null string and non-null string,
auto inputFixedWidthHasNull = inputVector1_;
// Third split has null string.
- std::vector<VectorPtr> stringHasNull = {
- makeNullableFlatVector<int8_t>({0, 1}),
- makeNullableFlatVector<int8_t>({0, -1}),
- makeNullableFlatVector<int32_t>({0, 100}),
- makeNullableFlatVector<int64_t>({0, 1}),
- makeNullableFlatVector<float>({0, 0.142857}),
- makeNullableFlatVector<bool>({false, true}),
- makeNullableFlatVector<StringView>({std::nullopt, std::nullopt}),
- makeNullableFlatVector<StringView>({std::nullopt, std::nullopt}),
- };
- auto inputStringHasNull = makeRowVector(stringHasNull);
+ auto inputStringHasNull = inputVectorStringHasNull_;
ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull));
// Split more data with null. Already filled + to be filled > buffer size,
Buffer is resized larger.
diff --git a/cpp/velox/tests/VeloxShuffleWriterTestBase.h
b/cpp/velox/tests/VeloxShuffleWriterTestBase.h
index 29aac9817b..717d17e001 100644
--- a/cpp/velox/tests/VeloxShuffleWriterTestBase.h
+++ b/cpp/velox/tests/VeloxShuffleWriterTestBase.h
@@ -106,6 +106,17 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
std::nullopt}),
makeNullableFlatVector<bool>(
{std::nullopt, true, false, std::nullopt, true, true, false, true,
std::nullopt, std::nullopt}),
+ makeNullableFlatVector<facebook::velox::Timestamp>(
+ {std::nullopt,
+ facebook::velox::Timestamp(5, 0),
+ std::nullopt,
+ std::nullopt,
+ facebook::velox::Timestamp(4, 0),
+ std::nullopt,
+ facebook::velox::Timestamp(2, 0),
+ facebook::velox::Timestamp(1, 0),
+ facebook::velox::Timestamp(0, 0),
+ std::nullopt}),
makeFlatVector<facebook::velox::StringView>(
{"a",
"bobbobbobooooooooooooooooooooooooooooob1",
@@ -137,6 +148,7 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
makeFlatVector<int64_t>({1, 1}),
makeFlatVector<float>({0.142857, -0.142857}),
makeFlatVector<bool>({true, false}),
+ makeNullableFlatVector<facebook::velox::Timestamp>({std::nullopt,
facebook::velox::Timestamp(5, 0)}),
makeFlatVector<facebook::velox::StringView>(
{"bob",
"alicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealice"}),
@@ -150,9 +162,37 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
makeFlatVector<int64_t>({0, 1}),
makeFlatVector<float>({0, 0.142857}),
makeFlatVector<bool>({false, true}),
+ makeFlatVector<facebook::velox::Timestamp>(
+ {facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(0,
0)}),
makeFlatVector<facebook::velox::StringView>({"", "alice"}),
makeFlatVector<facebook::velox::StringView>({"alice", ""}),
- };
+ facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), 2,
pool())};
+
+ childrenIntHasNull_ = {
+ makeNullableFlatVector<int8_t>({std::nullopt, 1}),
+ makeNullableFlatVector<int8_t>({std::nullopt, -1}),
+ makeNullableFlatVector<int32_t>({std::nullopt, 100}),
+ makeNullableFlatVector<int64_t>({0, 1}),
+ makeNullableFlatVector<float>({0, 0.142857}),
+ makeNullableFlatVector<bool>({false, true}),
+ makeNullableFlatVector<facebook::velox::Timestamp>(
+ {facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(4,
0)}),
+ makeNullableFlatVector<facebook::velox::StringView>({"", "alice"}),
+ makeNullableFlatVector<facebook::velox::StringView>({"alice", ""}),
+ facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), 2,
pool())};
+
+ childrenStringHasNull_ = {
+ makeNullableFlatVector<int8_t>({0, 1}),
+ makeNullableFlatVector<int8_t>({0, -1}),
+ makeNullableFlatVector<int32_t>({0, 100}),
+ makeNullableFlatVector<int64_t>({0, 1}),
+ makeNullableFlatVector<float>({0, 0.142857}),
+ makeNullableFlatVector<bool>({false, true}),
+ makeNullableFlatVector<facebook::velox::Timestamp>(
+ {facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(4,
0)}),
+ makeNullableFlatVector<facebook::velox::StringView>({std::nullopt,
std::nullopt}),
+ makeNullableFlatVector<facebook::velox::StringView>({std::nullopt,
std::nullopt}),
+ facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), 2,
pool())};
largeString1_ = makeString(1024);
int32_t numRows = 1024;
@@ -163,26 +203,30 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
makeFlatVector<int64_t>(std::vector<int64_t>(numRows, 0)),
makeFlatVector<float>(std::vector<float>(numRows, 0)),
makeFlatVector<bool>(std::vector<bool>(numRows, true)),
+ makeFlatVector<facebook::velox::Timestamp>(
+ std::vector<facebook::velox::Timestamp>(numRows,
facebook::velox::Timestamp(5, 0))),
makeNullableFlatVector<facebook::velox::StringView>(
std::vector<std::optional<facebook::velox::StringView>>(numRows,
largeString1_.c_str())),
makeNullableFlatVector<facebook::velox::StringView>(
std::vector<std::optional<facebook::velox::StringView>>(numRows,
std::nullopt)),
- };
+ facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(),
numRows, pool())};
largeString2_ = makeString(4096);
numRows = 2048;
- auto vectorToSpill = childrenLargeBinary2_ = {
+ childrenLargeBinary2_ = {
makeFlatVector<int8_t>(std::vector<int8_t>(numRows, 0)),
makeFlatVector<int8_t>(std::vector<int8_t>(numRows, 0)),
makeFlatVector<int32_t>(std::vector<int32_t>(numRows, 0)),
makeFlatVector<int64_t>(std::vector<int64_t>(numRows, 0)),
makeFlatVector<float>(std::vector<float>(numRows, 0)),
makeFlatVector<bool>(std::vector<bool>(numRows, true)),
+ makeFlatVector<facebook::velox::Timestamp>(
+ std::vector<facebook::velox::Timestamp>(numRows,
facebook::velox::Timestamp(5, 0))),
makeNullableFlatVector<facebook::velox::StringView>(
std::vector<std::optional<facebook::velox::StringView>>(numRows,
largeString2_.c_str())),
makeNullableFlatVector<facebook::velox::StringView>(
std::vector<std::optional<facebook::velox::StringView>>(numRows,
std::nullopt)),
- };
+ facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(),
numRows, pool())};
childrenComplex_ = {
makeNullableFlatVector<int32_t>({std::nullopt, 1}),
@@ -202,6 +246,8 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
inputVector1_ = makeRowVector(children1_);
inputVector2_ = makeRowVector(children2_);
inputVectorNoNull_ = makeRowVector(childrenNoNull_);
+ inputVectorIntHasNull_ = makeRowVector(childrenIntHasNull_);
+ inputVectorStringHasNull_ = makeRowVector(childrenStringHasNull_);
inputVectorLargeBinary1_ = makeRowVector(childrenLargeBinary1_);
inputVectorLargeBinary2_ = makeRowVector(childrenLargeBinary2_);
inputVectorComplex_ = makeRowVector(childrenComplex_);
@@ -242,6 +288,8 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
std::vector<facebook::velox::VectorPtr> children1_;
std::vector<facebook::velox::VectorPtr> children2_;
std::vector<facebook::velox::VectorPtr> childrenNoNull_;
+ std::vector<facebook::velox::VectorPtr> childrenIntHasNull_;
+ std::vector<facebook::velox::VectorPtr> childrenStringHasNull_;
std::vector<facebook::velox::VectorPtr> childrenLargeBinary1_;
std::vector<facebook::velox::VectorPtr> childrenLargeBinary2_;
std::vector<facebook::velox::VectorPtr> childrenComplex_;
@@ -249,6 +297,8 @@ class VeloxShuffleWriterTestBase : public
facebook::velox::test::VectorTestBase
facebook::velox::RowVectorPtr inputVector1_;
facebook::velox::RowVectorPtr inputVector2_;
facebook::velox::RowVectorPtr inputVectorNoNull_;
+ facebook::velox::RowVectorPtr inputVectorIntHasNull_;
+ facebook::velox::RowVectorPtr inputVectorStringHasNull_;
std::string largeString1_;
std::string largeString2_;
facebook::velox::RowVectorPtr inputVectorLargeBinary1_;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
