[VL] Add timestamp to shuffle test

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();
 }

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();
 }

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_;
 

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

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()));
   }
 }
 

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_;

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,

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