Skip to content

[VL] Add timestamp to shuffle test #12503

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marin-ma merged 3 commits into from
Jul 14, 2026
+138 −218
Merged

[VL] Add timestamp to shuffle test #12503

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3e6c2dd
add timestamp
marin-ma Jul 13, 2026
66aeebc
fix
marin-ma Jul 13, 2026
becdb35
remove duplicate tests
marin-ma Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion cpp/velox/shuffle/VeloxHashShuffleWriter.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
2 changes: 2 additions & 0 deletions cpp/velox/shuffle/VeloxRssSortShuffleWriter.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
2 changes: 0 additions & 2 deletions cpp/velox/shuffle/VeloxRssSortShuffleWriter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_;

Expand Down
2 changes: 2 additions & 0 deletions cpp/velox/shuffle/VeloxShuffleWriter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,8 @@ class VeloxShuffleWriter : public ShuffleWriter {

int32_t maxBatchSize_{0};

facebook::velox::RowTypePtr rowType_{nullptr};

enum EvictState { kEvictable, kUnevictable };

// stat
Expand Down
4 changes: 3 additions & 1 deletion cpp/velox/shuffle/VeloxSortShuffleWriter.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()));
}
}

Expand Down
1 change: 0 additions & 1 deletion cpp/velox/shuffle/VeloxSortShuffleWriter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_;
Expand Down
4 changes: 4 additions & 0 deletions cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -603,6 +603,8 @@ TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, preAllocForceRealloc) {
makeNullableFlatVector<int64_t>({0, 1}),
Comment thread
marin-ma marked this conversation as resolved.
Outdated
makeNullableFlatVector<float>({0, 0.142857}),
makeNullableFlatVector<bool>({false, true}),
makeNullableFlatVector<facebook::velox::Timestamp>(
{facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(4, 0)}),
makeNullableFlatVector<StringView>({"", "alice"}),
Comment thread
marin-ma marked this conversation as resolved.
Outdated
makeNullableFlatVector<StringView>({"alice", ""}),
};
Expand Down Expand Up @@ -663,6 +665,8 @@ TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, preAllocForceReuse) {
makeNullableFlatVector<int64_t>({0, 1}),
Comment thread
marin-ma marked this conversation as resolved.
Outdated
makeNullableFlatVector<float>({0, 0.142857}),
makeNullableFlatVector<bool>({false, true}),
makeNullableFlatVector<facebook::velox::Timestamp>(
{facebook::velox::Timestamp(5, 0), facebook::velox::Timestamp(4, 0)}),
Comment thread
marin-ma marked this conversation as resolved.
Outdated
makeNullableFlatVector<StringView>({std::nullopt, std::nullopt}),
makeNullableFlatVector<StringView>({std::nullopt, std::nullopt}),
};
Expand Down
27 changes: 3 additions & 24 deletions cpp/velox/tests/VeloxShuffleWriterTest.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
58 changes: 54 additions & 4 deletions cpp/velox/tests/VeloxShuffleWriterTestBase.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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"}),
Expand All @@ -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;
Expand All @@ -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}),
Expand All @@ -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_);
Expand Down Expand Up @@ -242,13 +288,17 @@ 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_;

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_;
Expand Down
Loading