diff --git a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc index 157d421644e..3cec6e4d41e 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 0f9297ac5f9..fd8af57aa7e 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(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 dbb5051bbd1..e2914f2322d 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 batch_; std::unique_ptr bufferOutputStream_; diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.h b/cpp/velox/shuffle/VeloxShuffleWriter.h index 45276e3cd41..5a478a77558 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 7a4066153b0..0a31dd3d7d3 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 a30f0edb83b..59541d24962 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 row2Partition_; - std::shared_ptr rowType_; std::optional fixedRowSize_; std::vector rowSize_; std::vector rowSizePrefixSum_; diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc index 7853ad418a4..cc846eb4eb6 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(); - 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 intHasNull = { - makeNullableFlatVector({std::nullopt, 1}), - makeNullableFlatVector({std::nullopt, -1}), - makeNullableFlatVector({std::nullopt, 100}), - makeNullableFlatVector({0, 1}), - makeNullableFlatVector({0, 0.142857}), - makeNullableFlatVector({false, true}), - makeNullableFlatVector({"", "alice"}), - makeNullableFlatVector({"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(); - 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 stringHasNull = { - makeNullableFlatVector({0, 1}), - makeNullableFlatVector({0, -1}), - makeNullableFlatVector({0, 100}), - makeNullableFlatVector({0, 1}), - makeNullableFlatVector({0, 0.142857}), - makeNullableFlatVector({false, true}), - makeNullableFlatVector({std::nullopt, std::nullopt}), - makeNullableFlatVector({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 32fe57b2418..129097106ce 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(200, pool()); } std::shared_ptr 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> makeNonSharedStringVector() const { + auto strBuffer = AlignedBuffer::allocate(200, pool()); + auto vector = BaseVector::create>(VARCHAR(), 100, pool()); + vector->setStringBuffers({std::move(strBuffer)}); + return vector; + } + + std::shared_ptr> makeSharedStringVector() { + auto vector = BaseVector::create>(VARCHAR(), 100, pool()); + vector->setStringBuffers({strBuffer_}); + return vector; + } + + BufferPtr strBuffer_; + + static constexpr int64_t kMemLimit = std::numeric_limits::max(); }; TEST_F(VeloxRssSortShuffleWriterTest, calculateBatchesSize) { auto shuffleWriter = std::dynamic_pointer_cast(createShuffleWriter(10)); - // Do not trigger resetBatches by shuffle writer. - const int64_t memLimit = INT64_MAX; - // Shared string buffer in FlatVector. - BufferPtr strBuffer = AlignedBuffer::allocate(200, pool()); - auto vector1 = BaseVector::create>(VARCHAR(), 100, pool()); - vector1->setStringBuffers({strBuffer}); - auto vector2 = BaseVector::create>(VARCHAR(), 100, pool()); - vector2->setStringBuffers({strBuffer}); - auto vector3 = BaseVector::create>(VARCHAR(), 100, pool()); - vector3->setStringBuffers({strBuffer}); - auto vector4 = BaseVector::create>(INTEGER(), 100, pool()); - - auto rowVector1 = makeRowVector({vector1, vector2}); - auto rowVector2 = makeRowVector({vector3, vector4}); - std::shared_ptr cb1 = std::make_shared(rowVector1); - std::shared_ptr cb2 = std::make_shared(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(rowVector1); + auto cb2 = std::make_shared(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(createShuffleWriter(10)); + + // Shared string buffer in FlatVector. + auto vector = makeSharedStringVector(); // Shared string buffer in ArrayVector. - BufferPtr offsets = allocateOffsets(1, vector1->pool()); - BufferPtr sizes = allocateOffsets(1, vector1->pool()); - sizes->asMutable()[0] = vector1->size(); + BufferPtr offsets = allocateOffsets(1, vector->pool()); + BufferPtr sizes = allocateOffsets(1, vector->pool()); + sizes->asMutable()[0] = vector->size(); auto arrayVector = - std::make_shared(pool(), ARRAY(VARCHAR()), nullptr, 1, offsets, sizes, vector1); - auto rowVector3 = makeRowVector({arrayVector, vector1}); - cb1 = std::make_shared(rowVector3); - ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit)); - expectedSize = rowVector3->retainedSize() - strBuffer->capacity(); + std::make_shared(pool(), ARRAY(VARCHAR()), nullptr, 1, offsets, sizes, vector); + auto rowVector = makeRowVector({arrayVector, vector}); + auto cb1 = std::make_shared(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(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(rowVector4); - ASSERT_NOT_OK(shuffleWriter->write(cb1, memLimit)); - expectedSize = rowVector4->retainedSize() - strBuffer->capacity() * 2; + auto rowVector = makeRowVector({mapVector, makeSharedStringVector()}); + auto cb = std::make_shared(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(createShuffleWriter(10)); // Shared string buffer in RowVector. - auto rowVector5 = makeRowVector({rowVector1, vector3}); - cb1 = std::make_shared(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(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(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(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(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 999d5e74354..27bc3429462 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 intHasNull = { - makeNullableFlatVector({std::nullopt, 1}), - makeNullableFlatVector({std::nullopt, -1}), - makeNullableFlatVector({std::nullopt, 100}), - makeNullableFlatVector({0, 1}), - makeNullableFlatVector({0, 0.142857}), - makeNullableFlatVector({false, true}), - makeNullableFlatVector({"", "alice"}), - makeNullableFlatVector({"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 stringHasNull = { - makeNullableFlatVector({0, 1}), - makeNullableFlatVector({0, -1}), - makeNullableFlatVector({0, 100}), - makeNullableFlatVector({0, 1}), - makeNullableFlatVector({0, 0.142857}), - makeNullableFlatVector({false, true}), - makeNullableFlatVector({std::nullopt, std::nullopt}), - makeNullableFlatVector({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 29aac9817b9..717d17e0016 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( {std::nullopt, true, false, std::nullopt, true, true, false, true, std::nullopt, std::nullopt}), + makeNullableFlatVector( + {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( {"a", "bobbobbobooooooooooooooooooooooooooooob1", @@ -137,6 +148,7 @@ class VeloxShuffleWriterTestBase : public facebook::velox::test::VectorTestBase makeFlatVector({1, 1}), makeFlatVector({0.142857, -0.142857}), makeFlatVector({true, false}), + makeNullableFlatVector({std::nullopt, facebook::velox::Timestamp(5, 0)}), makeFlatVector( {"bob", "alicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealicealice"}), @@ -150,9 +162,37 @@ class VeloxShuffleWriterTestBase : public facebook::velox::test::VectorTestBase makeFlatVector({0, 1}), makeFlatVector({0, 0.142857}), makeFlatVector({false, true}), + makeFlatVector( + {facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(0, 0)}), makeFlatVector({"", "alice"}), makeFlatVector({"alice", ""}), - }; + facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), 2, pool())}; + + childrenIntHasNull_ = { + makeNullableFlatVector({std::nullopt, 1}), + makeNullableFlatVector({std::nullopt, -1}), + makeNullableFlatVector({std::nullopt, 100}), + makeNullableFlatVector({0, 1}), + makeNullableFlatVector({0, 0.142857}), + makeNullableFlatVector({false, true}), + makeNullableFlatVector( + {facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(4, 0)}), + makeNullableFlatVector({"", "alice"}), + makeNullableFlatVector({"alice", ""}), + facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), 2, pool())}; + + childrenStringHasNull_ = { + makeNullableFlatVector({0, 1}), + makeNullableFlatVector({0, -1}), + makeNullableFlatVector({0, 100}), + makeNullableFlatVector({0, 1}), + makeNullableFlatVector({0, 0.142857}), + makeNullableFlatVector({false, true}), + makeNullableFlatVector( + {facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(4, 0)}), + makeNullableFlatVector({std::nullopt, std::nullopt}), + makeNullableFlatVector({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(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, true)), + makeFlatVector( + std::vector(numRows, facebook::velox::Timestamp(5, 0))), makeNullableFlatVector( std::vector>(numRows, largeString1_.c_str())), makeNullableFlatVector( std::vector>(numRows, std::nullopt)), - }; + facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), numRows, pool())}; largeString2_ = makeString(4096); numRows = 2048; - auto vectorToSpill = childrenLargeBinary2_ = { + childrenLargeBinary2_ = { makeFlatVector(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, 0)), makeFlatVector(std::vector(numRows, true)), + makeFlatVector( + std::vector(numRows, facebook::velox::Timestamp(5, 0))), makeNullableFlatVector( std::vector>(numRows, largeString2_.c_str())), makeNullableFlatVector( std::vector>(numRows, std::nullopt)), - }; + facebook::velox::BaseVector::create(facebook::velox::UNKNOWN(), numRows, pool())}; childrenComplex_ = { makeNullableFlatVector({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 children1_; std::vector children2_; std::vector childrenNoNull_; + std::vector childrenIntHasNull_; + std::vector childrenStringHasNull_; std::vector childrenLargeBinary1_; std::vector childrenLargeBinary2_; std::vector 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_;