Copilot commented on code in PR #12983: URL: https://github.com/apache/gluten/pull/12983#discussion_r3967098587
########## cpp/velox/tests/VeloxShuffleReaderTest.cc: ########## @@ -0,0 +1,401 @@ +/* + * 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. + */ + +// Regression tests for the rss_sort shuffle reader, driven through the +// public VeloxRssSortShuffleReaderDeserializer API via controllable fake +// arrow::io::InputStreams. +// +// Covers: +// - graceful EOS on an empty stream (e.g. an empty Celeborn partition); +// - EOS hit mid-page on a truncated compressed page; +// - a buggy upstream whose Read() returns a negative byte count; +// - uncompressed Presto pages spanning multiple read windows (nested struct +// pre-scan, header crossing a refill boundary, checksummed pages) and the +// single-window zero-copy fast path. +// +// NOTE on the two bug-probe cases (EosMidPageTerminates, NegativeReadThrows): +// on code where VeloxInputStream::next() ignores its throwIfPastEnd argument +// and silently returns on EOS, the mid-page case enters the readBytes for(;;) +// loop that never exits. FakeInputStream counts consecutive EOS reads and +// throws after kMaxConsecutiveEosReads so the loop is cut short in +// milliseconds instead of hanging until the test timeout; the assertion then +// fails on the message mismatch and surfaces "possible infinite loop" as the +// actual exception. Both cases turn green once the EOS handling fix lands. + +#include <gtest/gtest.h> + +#include <arrow/buffer.h> +#include <arrow/io/interfaces.h> +#include <arrow/result.h> +#include <arrow/status.h> + +#include <cstdint> +#include <cstring> +#include <memory> +#include <sstream> +#include <stdexcept> Review Comment: `FakeInputStream::Read` uses `std::min` and the error path uses `std::to_string`, but `<algorithm>` and `<string>` aren’t included here. This can fail to compile on stricter standard library configurations. Add the missing standard headers explicitly. ########## cpp/velox/tests/VeloxShuffleReaderTest.cc: ########## @@ -0,0 +1,401 @@ +/* + * 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. + */ + +// Regression tests for the rss_sort shuffle reader, driven through the +// public VeloxRssSortShuffleReaderDeserializer API via controllable fake +// arrow::io::InputStreams. +// +// Covers: +// - graceful EOS on an empty stream (e.g. an empty Celeborn partition); +// - EOS hit mid-page on a truncated compressed page; +// - a buggy upstream whose Read() returns a negative byte count; +// - uncompressed Presto pages spanning multiple read windows (nested struct +// pre-scan, header crossing a refill boundary, checksummed pages) and the +// single-window zero-copy fast path. +// +// NOTE on the two bug-probe cases (EosMidPageTerminates, NegativeReadThrows): +// on code where VeloxInputStream::next() ignores its throwIfPastEnd argument +// and silently returns on EOS, the mid-page case enters the readBytes for(;;) +// loop that never exits. FakeInputStream counts consecutive EOS reads and +// throws after kMaxConsecutiveEosReads so the loop is cut short in +// milliseconds instead of hanging until the test timeout; the assertion then +// fails on the message mismatch and surfaces "possible infinite loop" as the +// actual exception. Both cases turn green once the EOS handling fix lands. + +#include <gtest/gtest.h> + +#include <arrow/buffer.h> +#include <arrow/io/interfaces.h> +#include <arrow/result.h> +#include <arrow/status.h> + +#include <cstdint> +#include <cstring> +#include <memory> +#include <sstream> +#include <stdexcept> +#include <vector> + +#include "compute/VeloxBackend.h" +#include "memory/VeloxColumnarBatch.h" +#include "memory/VeloxMemoryManager.h" +#include "shuffle/GlutenByteStream.h" +#include "shuffle/VeloxShuffleReader.h" +#include "tests/utils/TestStreamReader.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/serializers/PrestoSerializer.h" +#include "velox/type/Type.h" +#include "velox/vector/VectorStream.h" +#include "velox/vector/tests/utils/VectorTestBase.h" + +using namespace facebook::velox; +using namespace facebook::velox::test; + +namespace gluten { + +namespace { +// A minimal arrow::io::InputStream backed by a fixed in-memory payload. Once +// the payload is exhausted, Read returns 0 (EOS). With `negativeRead`, Read +// always returns -1 instead, modeling a buggy upstream that reports EOF as a +// negative byte count. With `firstReadLimit` >= 0, only the FIRST Read is +// capped to that many bytes (subsequent reads are unbounded), modeling an +// upstream (e.g. a network stream) whose first chunk ends mid-page. +// +// To keep a possible reader-side infinite loop (readBytes -> next() -> EOS -> +// silently return -> spin) from hanging the test until the CI timeout, Read +// throws after kMaxConsecutiveEosReads consecutive EOS returns. Well-behaved +// readers probe EOS only a couple of times, so the cap never trips for them. +class FakeInputStream final : public arrow::io::InputStream { + public: + explicit FakeInputStream( + std::vector<uint8_t> payload = {}, + bool negativeRead = false, + int64_t firstReadLimit = -1) + : payload_(std::move(payload)), + negativeRead_(negativeRead), + firstReadLimit_(firstReadLimit) {} + + arrow::Status Close() override { + closed_ = true; + return arrow::Status::OK(); + } + arrow::Result<int64_t> Tell() const override { + return pos_; + } + bool closed() const override { + return closed_; + } + + arrow::Result<int64_t> Read(int64_t nbytes, void* out) override { + if (negativeRead_) { + return static_cast<int64_t>(-1); + } + int64_t toRead = std::min<int64_t>(nbytes, static_cast<int64_t>(payload_.size()) - pos_); + if (firstRead_ && firstReadLimit_ >= 0) { + toRead = std::min<int64_t>(toRead, firstReadLimit_); + firstRead_ = false; + } + if (toRead > 0) { + std::memcpy(out, payload_.data() + pos_, toRead); + pos_ += toRead; + consecutiveEosReads_ = 0; + } else if (++consecutiveEosReads_ > kMaxConsecutiveEosReads) { + // Throw a plain C++ exception: the reader wraps Read() in + // arrow::Result and drops arrow errors via .ValueOr(0), so an + // arrow::Status::IOError would be swallowed and the loop would spin on. + throw std::runtime_error( + "possible infinite loop: Read() returned 0 for " + std::to_string(kMaxConsecutiveEosReads) + + " consecutive calls"); + } + return toRead; // 0 == EOS when payload exhausted + } + + arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override { + ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(nbytes)); + ARROW_ASSIGN_OR_RAISE(int64_t bytesRead, Read(nbytes, buffer->mutable_data())); + ARROW_RETURN_NOT_OK(buffer->Resize(bytesRead, false)); + buffer->ZeroPadding(); + return std::move(buffer); + } Review Comment: If `negativeRead_` is enabled, `Read(nbytes, ...)` returns `-1`, and this overload then calls `buffer->Resize(-1, ...)`, which is invalid and will likely fail in an unexpected way (or mask what the test is trying to model). Add an explicit guard for `bytesRead < 0` and fail in a controlled way (e.g., throw `std::runtime_error` like the EOS-loop breaker, or return an Arrow error) so the behavior is deterministic if this overload is used. ########## cpp/velox/tests/VeloxShuffleReaderTest.cc: ########## @@ -0,0 +1,401 @@ +/* + * 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. + */ + +// Regression tests for the rss_sort shuffle reader, driven through the +// public VeloxRssSortShuffleReaderDeserializer API via controllable fake +// arrow::io::InputStreams. +// +// Covers: +// - graceful EOS on an empty stream (e.g. an empty Celeborn partition); +// - EOS hit mid-page on a truncated compressed page; +// - a buggy upstream whose Read() returns a negative byte count; +// - uncompressed Presto pages spanning multiple read windows (nested struct +// pre-scan, header crossing a refill boundary, checksummed pages) and the +// single-window zero-copy fast path. +// +// NOTE on the two bug-probe cases (EosMidPageTerminates, NegativeReadThrows): +// on code where VeloxInputStream::next() ignores its throwIfPastEnd argument +// and silently returns on EOS, the mid-page case enters the readBytes for(;;) +// loop that never exits. FakeInputStream counts consecutive EOS reads and +// throws after kMaxConsecutiveEosReads so the loop is cut short in +// milliseconds instead of hanging until the test timeout; the assertion then +// fails on the message mismatch and surfaces "possible infinite loop" as the +// actual exception. Both cases turn green once the EOS handling fix lands. + +#include <gtest/gtest.h> + +#include <arrow/buffer.h> +#include <arrow/io/interfaces.h> +#include <arrow/result.h> +#include <arrow/status.h> + +#include <cstdint> +#include <cstring> +#include <memory> +#include <sstream> +#include <stdexcept> +#include <vector> + +#include "compute/VeloxBackend.h" +#include "memory/VeloxColumnarBatch.h" +#include "memory/VeloxMemoryManager.h" +#include "shuffle/GlutenByteStream.h" +#include "shuffle/VeloxShuffleReader.h" +#include "tests/utils/TestStreamReader.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/serializers/PrestoSerializer.h" +#include "velox/type/Type.h" +#include "velox/vector/VectorStream.h" +#include "velox/vector/tests/utils/VectorTestBase.h" + +using namespace facebook::velox; +using namespace facebook::velox::test; + +namespace gluten { + +namespace { +// A minimal arrow::io::InputStream backed by a fixed in-memory payload. Once +// the payload is exhausted, Read returns 0 (EOS). With `negativeRead`, Read +// always returns -1 instead, modeling a buggy upstream that reports EOF as a +// negative byte count. With `firstReadLimit` >= 0, only the FIRST Read is +// capped to that many bytes (subsequent reads are unbounded), modeling an +// upstream (e.g. a network stream) whose first chunk ends mid-page. +// +// To keep a possible reader-side infinite loop (readBytes -> next() -> EOS -> +// silently return -> spin) from hanging the test until the CI timeout, Read +// throws after kMaxConsecutiveEosReads consecutive EOS returns. Well-behaved +// readers probe EOS only a couple of times, so the cap never trips for them. +class FakeInputStream final : public arrow::io::InputStream { + public: + explicit FakeInputStream( + std::vector<uint8_t> payload = {}, + bool negativeRead = false, + int64_t firstReadLimit = -1) + : payload_(std::move(payload)), + negativeRead_(negativeRead), + firstReadLimit_(firstReadLimit) {} + + arrow::Status Close() override { + closed_ = true; + return arrow::Status::OK(); + } + arrow::Result<int64_t> Tell() const override { + return pos_; + } + bool closed() const override { + return closed_; + } + + arrow::Result<int64_t> Read(int64_t nbytes, void* out) override { + if (negativeRead_) { + return static_cast<int64_t>(-1); + } + int64_t toRead = std::min<int64_t>(nbytes, static_cast<int64_t>(payload_.size()) - pos_); + if (firstRead_ && firstReadLimit_ >= 0) { + toRead = std::min<int64_t>(toRead, firstReadLimit_); + firstRead_ = false; + } + if (toRead > 0) { + std::memcpy(out, payload_.data() + pos_, toRead); + pos_ += toRead; + consecutiveEosReads_ = 0; + } else if (++consecutiveEosReads_ > kMaxConsecutiveEosReads) { + // Throw a plain C++ exception: the reader wraps Read() in + // arrow::Result and drops arrow errors via .ValueOr(0), so an + // arrow::Status::IOError would be swallowed and the loop would spin on. + throw std::runtime_error( + "possible infinite loop: Read() returned 0 for " + std::to_string(kMaxConsecutiveEosReads) + + " consecutive calls"); + } + return toRead; // 0 == EOS when payload exhausted + } + + arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override { + ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(nbytes)); + ARROW_ASSIGN_OR_RAISE(int64_t bytesRead, Read(nbytes, buffer->mutable_data())); + ARROW_RETURN_NOT_OK(buffer->Resize(bytesRead, false)); + buffer->ZeroPadding(); + return std::move(buffer); + } + + private: + static constexpr int32_t kMaxConsecutiveEosReads = 100; + + std::vector<uint8_t> payload_; + int64_t pos_{0}; + bool negativeRead_{false}; + int64_t firstReadLimit_{-1}; + bool firstRead_{true}; + int32_t consecutiveEosReads_{0}; + bool closed_{false}; +}; + +// Append a little-endian POD value to `out` (Presto page header fields are +// machine byte order / little-endian on x86). +template <typename T> +void appendLe(std::vector<uint8_t>& out, T value) { + T v = value; + const auto* p = reinterpret_cast<const uint8_t*>(&v); + out.insert(out.end(), p, p + sizeof(T)); +} Review Comment: `appendLe` claims little-endian encoding, but it currently writes native-endian bytes. On big-endian architectures this will generate invalid Presto headers and make the tests non-portable. Encode to little-endian explicitly (and consider `static_assert(std::is_trivially_copyable_v<T>)` to prevent accidental misuse with non-POD types). ########## cpp/velox/tests/VeloxShuffleReaderTest.cc: ########## @@ -0,0 +1,401 @@ +/* + * 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. + */ + +// Regression tests for the rss_sort shuffle reader, driven through the +// public VeloxRssSortShuffleReaderDeserializer API via controllable fake +// arrow::io::InputStreams. +// +// Covers: +// - graceful EOS on an empty stream (e.g. an empty Celeborn partition); +// - EOS hit mid-page on a truncated compressed page; +// - a buggy upstream whose Read() returns a negative byte count; +// - uncompressed Presto pages spanning multiple read windows (nested struct +// pre-scan, header crossing a refill boundary, checksummed pages) and the +// single-window zero-copy fast path. +// +// NOTE on the two bug-probe cases (EosMidPageTerminates, NegativeReadThrows): +// on code where VeloxInputStream::next() ignores its throwIfPastEnd argument +// and silently returns on EOS, the mid-page case enters the readBytes for(;;) +// loop that never exits. FakeInputStream counts consecutive EOS reads and +// throws after kMaxConsecutiveEosReads so the loop is cut short in +// milliseconds instead of hanging until the test timeout; the assertion then +// fails on the message mismatch and surfaces "possible infinite loop" as the +// actual exception. Both cases turn green once the EOS handling fix lands. + +#include <gtest/gtest.h> + +#include <arrow/buffer.h> +#include <arrow/io/interfaces.h> +#include <arrow/result.h> +#include <arrow/status.h> + +#include <cstdint> +#include <cstring> +#include <memory> +#include <sstream> +#include <stdexcept> +#include <vector> + +#include "compute/VeloxBackend.h" +#include "memory/VeloxColumnarBatch.h" +#include "memory/VeloxMemoryManager.h" +#include "shuffle/GlutenByteStream.h" +#include "shuffle/VeloxShuffleReader.h" +#include "tests/utils/TestStreamReader.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/serializers/PrestoSerializer.h" +#include "velox/type/Type.h" +#include "velox/vector/VectorStream.h" +#include "velox/vector/tests/utils/VectorTestBase.h" + +using namespace facebook::velox; +using namespace facebook::velox::test; + +namespace gluten { + +namespace { +// A minimal arrow::io::InputStream backed by a fixed in-memory payload. Once +// the payload is exhausted, Read returns 0 (EOS). With `negativeRead`, Read +// always returns -1 instead, modeling a buggy upstream that reports EOF as a +// negative byte count. With `firstReadLimit` >= 0, only the FIRST Read is +// capped to that many bytes (subsequent reads are unbounded), modeling an +// upstream (e.g. a network stream) whose first chunk ends mid-page. +// +// To keep a possible reader-side infinite loop (readBytes -> next() -> EOS -> +// silently return -> spin) from hanging the test until the CI timeout, Read +// throws after kMaxConsecutiveEosReads consecutive EOS returns. Well-behaved +// readers probe EOS only a couple of times, so the cap never trips for them. +class FakeInputStream final : public arrow::io::InputStream { + public: + explicit FakeInputStream( + std::vector<uint8_t> payload = {}, + bool negativeRead = false, + int64_t firstReadLimit = -1) + : payload_(std::move(payload)), + negativeRead_(negativeRead), + firstReadLimit_(firstReadLimit) {} + + arrow::Status Close() override { + closed_ = true; + return arrow::Status::OK(); + } + arrow::Result<int64_t> Tell() const override { + return pos_; + } + bool closed() const override { + return closed_; + } + + arrow::Result<int64_t> Read(int64_t nbytes, void* out) override { + if (negativeRead_) { + return static_cast<int64_t>(-1); + } + int64_t toRead = std::min<int64_t>(nbytes, static_cast<int64_t>(payload_.size()) - pos_); + if (firstRead_ && firstReadLimit_ >= 0) { + toRead = std::min<int64_t>(toRead, firstReadLimit_); + firstRead_ = false; + } + if (toRead > 0) { + std::memcpy(out, payload_.data() + pos_, toRead); + pos_ += toRead; + consecutiveEosReads_ = 0; + } else if (++consecutiveEosReads_ > kMaxConsecutiveEosReads) { + // Throw a plain C++ exception: the reader wraps Read() in + // arrow::Result and drops arrow errors via .ValueOr(0), so an + // arrow::Status::IOError would be swallowed and the loop would spin on. + throw std::runtime_error( + "possible infinite loop: Read() returned 0 for " + std::to_string(kMaxConsecutiveEosReads) + + " consecutive calls"); + } + return toRead; // 0 == EOS when payload exhausted + } + + arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override { + ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(nbytes)); + ARROW_ASSIGN_OR_RAISE(int64_t bytesRead, Read(nbytes, buffer->mutable_data())); + ARROW_RETURN_NOT_OK(buffer->Resize(bytesRead, false)); + buffer->ZeroPadding(); + return std::move(buffer); + } + + private: + static constexpr int32_t kMaxConsecutiveEosReads = 100; + + std::vector<uint8_t> payload_; + int64_t pos_{0}; + bool negativeRead_{false}; + int64_t firstReadLimit_{-1}; + bool firstRead_{true}; + int32_t consecutiveEosReads_{0}; + bool closed_{false}; +}; + +// Append a little-endian POD value to `out` (Presto page header fields are +// machine byte order / little-endian on x86). +template <typename T> +void appendLe(std::vector<uint8_t>& out, T value) { + T v = value; + const auto* p = reinterpret_cast<const uint8_t*>(&v); + out.insert(out.end(), p, p + sizeof(T)); +} + +// Build a truncated Presto compressed page: a valid 21-byte header declaring +// compressedSize bytes of body, but only `bodyBytes` bytes follow. The +// reader's compressed branch calls source->readBytes(buf, compressedSize); +// when EOS is hit mid-drain, GlutenByteInputStream::readBytes loops to +// next(true) which must VELOX_FAIL instead of spinning. +// +// Header layout (PrestoHeader.cpp): numRows:int32, pageCodecMarker:int8, +// uncompressedSize:int32, compressedSize:int32, checksum:int64 == 21 bytes. +// pageCodecMarker = kCompressedBitMask (1), no checksum bit -> actualCheckSum +// stays 0 and matches header.checksum = 0 (PrestoSerializer.cpp:159). +std::vector<uint8_t> buildTruncatedCompressedPage(int32_t compressedSize, int32_t bodyBytes) { + std::vector<uint8_t> out; + out.reserve(21 + bodyBytes); + appendLe<int32_t>(out, 1); // numRows + appendLe<int8_t>(out, 1); // pageCodecMarker = kCompressedBitMask, no checksum + appendLe<int32_t>(out, compressedSize + 64); // uncompressedSize + appendLe<int32_t>(out, compressedSize); + appendLe<int64_t>(out, 0); // checksum + for (int i = 0; i < bodyBytes; ++i) { + out.push_back(static_cast<uint8_t>(i & 0xFF)); + } + return out; +} +} // namespace + +class VeloxShuffleReaderTest : public ::testing::Test, public test::VectorTestBase { + protected: + static void SetUpTestSuite() { + VeloxBackend::create(AllocationListener::noop(), {}); + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + } + + static void TearDownTestSuite() { + VeloxBackend::get()->tearDown(); + } + + std::shared_ptr<VeloxRssSortShuffleReaderDeserializer> makeDeserializer( + const std::shared_ptr<arrow::io::InputStream>& in, + const RowTypePtr& rowType = ROW({"c0"}, {INTEGER()})) { + int64_t deserializeTime = 0; + return std::make_shared<VeloxRssSortShuffleReaderDeserializer>( + std::make_shared<TestStreamReader>(in), + getDefaultMemoryManager(), + rowType, + /*batchSize=*/1024, + common::CompressionKind_NONE, + deserializeTime); + } + + // Serializes `rowVector` into a single uncompressed Presto page + // (21-byte header + payload), the exact wire format the rss-sort writer + // produces. With `withChecksum`, a PrestoOutputStreamListener is attached so + // the writer fills in the checksum bit and CRC (same mechanism as + // VeloxHashShuffleWriter's complex-type flush). NOTE: must not use + // gluten::BufferOutputStream here — its write() ignores the listener. + std::vector<uint8_t> serializePage(const RowVectorPtr& rowVector, bool withChecksum = false) { + serializer::presto::PrestoVectorSerde::PrestoOptions options; + options.compressionKind = common::CompressionKind_NONE; + auto serde = std::make_unique<serializer::presto::PrestoVectorSerde>(); + VectorStreamGroup group(pool(), serde.get()); + group.createStreamTree(asRowType(rowVector->type()), rowVector->size(), &options); + group.append(rowVector); + serializer::presto::PrestoOutputStreamListener listener; + std::stringstream out; + facebook::velox::OStreamOutputStream os(&out, withChecksum ? &listener : nullptr); + group.flush(&os); + const auto str = out.str(); + return std::vector<uint8_t>(str.begin(), str.end()); + } + + // ROW<c0: ARRAY<ROW<a: INTEGER>>> with `numArrays` arrays of + // `elementsPerArray` dense elements each. The nested struct triggers the + // Presto serde's pre-scan (tellp -> scan page -> seekp back). + RowVectorPtr makeNestedArraysRowVector(vector_size_t numArrays, vector_size_t elementsPerArray) { + const vector_size_t numElements = numArrays * elementsPerArray; + auto offsets = AlignedBuffer::allocate<vector_size_t>(numArrays, pool()); + auto sizes = AlignedBuffer::allocate<vector_size_t>(numArrays, pool()); + auto* rawOffsets = offsets->asMutable<vector_size_t>(); + auto* rawSizes = sizes->asMutable<vector_size_t>(); + for (vector_size_t i = 0; i < numArrays; ++i) { + rawOffsets[i] = i * elementsPerArray; + rawSizes[i] = elementsPerArray; + } + auto elements = makeRowVector({makeFlatVector<int32_t>( + numElements, [](vector_size_t row) { return row % 1024; })}); + auto arrayVector = std::make_shared<ArrayVector>( + pool(), + ARRAY(ROW({"a"}, {INTEGER()})), + BufferPtr(nullptr), + numArrays, + offsets, + sizes, + elements); + return makeRowVector({arrayVector}); + } +}; + +// Empty stream (e.g. an empty Celeborn partition): construction must NOT +// throw; next() returns nullptr (graceful EOS). +TEST_F(VeloxShuffleReaderTest, EmptyStreamGracefulEos) { + auto deserializer = makeDeserializer(std::make_shared<FakeInputStream>()); + EXPECT_EQ(deserializer->next(), nullptr); +} + +// Truncated compressed page: header reads fine and construction succeeds, +// but next() drives PrestoVectorSerde::deserialize's compressed branch -> +// readBytes(compressedSize) -> next(true) on EOS -> VELOX_FAIL. Pre-fix this +// is the documented infinite loop; post-fix it throws. See the file header +// note for how FakeInputStream cuts the pre-fix loop short. +TEST_F(VeloxShuffleReaderTest, EosMidPageTerminates) { + auto payload = buildTruncatedCompressedPage(/*compressedSize=*/1000, /*bodyBytes=*/8); + auto deserializer = makeDeserializer(std::make_shared<FakeInputStream>(std::move(payload))); + + VELOX_ASSERT_THROW( + deserializer->next(), "Reading past end of VeloxRssSortShuffleReaderDeserializer::VeloxInputStream"); +} + +// A buggy upstream whose Read returns a negative byte count. Without a +// signed-result guard the negative value would implicitly convert to a huge +// uint64_t offset_ and corrupt setRange / loop forever. With the guard it +// fails fast during the probe-read. +TEST_F(VeloxShuffleReaderTest, NegativeReadThrows) { + auto deserializer = makeDeserializer( + std::make_shared<FakeInputStream>(std::vector<uint8_t>{}, /*negativeRead=*/true)); + VELOX_ASSERT_THROW(deserializer->next(), "Read returned negative value"); +} + +// EOS bug regression: uncompressed + nested struct + page spanning multiple +// read windows is the exact combination that reproduced the original failure +// (corrupted data / spurious "Reading past end" EOS). With the fix the page +// deserializes correctly. +TEST_F(VeloxShuffleReaderTest, UncompressedNestedStructPageSpansWindows) { + constexpr vector_size_t kNumArrays = 20000; + constexpr vector_size_t kElementsPerArray = 30; + + auto rowVector = makeNestedArraysRowVector(kNumArrays, kElementsPerArray); + auto payload = serializePage(rowVector); + // The page must be larger than the reader's read window (~1MB) to exercise + // the multi-window slow path. + ASSERT_GT(payload.size(), 1 << 20); Review Comment: This test constructs a very large nested vector (`20000 * 30` elements) just to exceed ~1MB, which can make the test slow and memory-heavy in CI. Consider generating the smallest input that reliably crosses the window boundary (e.g., start smaller and scale until `payload.size() > threshold`, or use a more direct way to create a payload slightly above the window size). This keeps runtime stable while still testing the multi-window path. ########## cpp/velox/tests/VeloxShuffleReaderTest.cc: ########## @@ -0,0 +1,401 @@ +/* + * 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. + */ + +// Regression tests for the rss_sort shuffle reader, driven through the +// public VeloxRssSortShuffleReaderDeserializer API via controllable fake +// arrow::io::InputStreams. +// +// Covers: +// - graceful EOS on an empty stream (e.g. an empty Celeborn partition); +// - EOS hit mid-page on a truncated compressed page; +// - a buggy upstream whose Read() returns a negative byte count; +// - uncompressed Presto pages spanning multiple read windows (nested struct +// pre-scan, header crossing a refill boundary, checksummed pages) and the +// single-window zero-copy fast path. +// +// NOTE on the two bug-probe cases (EosMidPageTerminates, NegativeReadThrows): +// on code where VeloxInputStream::next() ignores its throwIfPastEnd argument +// and silently returns on EOS, the mid-page case enters the readBytes for(;;) +// loop that never exits. FakeInputStream counts consecutive EOS reads and +// throws after kMaxConsecutiveEosReads so the loop is cut short in +// milliseconds instead of hanging until the test timeout; the assertion then +// fails on the message mismatch and surfaces "possible infinite loop" as the +// actual exception. Both cases turn green once the EOS handling fix lands. + +#include <gtest/gtest.h> + +#include <arrow/buffer.h> +#include <arrow/io/interfaces.h> +#include <arrow/result.h> +#include <arrow/status.h> + +#include <cstdint> +#include <cstring> +#include <memory> +#include <sstream> +#include <stdexcept> +#include <vector> + +#include "compute/VeloxBackend.h" +#include "memory/VeloxColumnarBatch.h" +#include "memory/VeloxMemoryManager.h" +#include "shuffle/GlutenByteStream.h" +#include "shuffle/VeloxShuffleReader.h" +#include "tests/utils/TestStreamReader.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/serializers/PrestoSerializer.h" +#include "velox/type/Type.h" +#include "velox/vector/VectorStream.h" +#include "velox/vector/tests/utils/VectorTestBase.h" + +using namespace facebook::velox; +using namespace facebook::velox::test; + +namespace gluten { + +namespace { +// A minimal arrow::io::InputStream backed by a fixed in-memory payload. Once +// the payload is exhausted, Read returns 0 (EOS). With `negativeRead`, Read +// always returns -1 instead, modeling a buggy upstream that reports EOF as a +// negative byte count. With `firstReadLimit` >= 0, only the FIRST Read is +// capped to that many bytes (subsequent reads are unbounded), modeling an +// upstream (e.g. a network stream) whose first chunk ends mid-page. +// +// To keep a possible reader-side infinite loop (readBytes -> next() -> EOS -> +// silently return -> spin) from hanging the test until the CI timeout, Read +// throws after kMaxConsecutiveEosReads consecutive EOS returns. Well-behaved +// readers probe EOS only a couple of times, so the cap never trips for them. +class FakeInputStream final : public arrow::io::InputStream { + public: + explicit FakeInputStream( + std::vector<uint8_t> payload = {}, + bool negativeRead = false, + int64_t firstReadLimit = -1) + : payload_(std::move(payload)), + negativeRead_(negativeRead), + firstReadLimit_(firstReadLimit) {} + + arrow::Status Close() override { + closed_ = true; + return arrow::Status::OK(); + } + arrow::Result<int64_t> Tell() const override { + return pos_; + } + bool closed() const override { + return closed_; + } + + arrow::Result<int64_t> Read(int64_t nbytes, void* out) override { + if (negativeRead_) { + return static_cast<int64_t>(-1); + } + int64_t toRead = std::min<int64_t>(nbytes, static_cast<int64_t>(payload_.size()) - pos_); + if (firstRead_ && firstReadLimit_ >= 0) { + toRead = std::min<int64_t>(toRead, firstReadLimit_); + firstRead_ = false; + } + if (toRead > 0) { + std::memcpy(out, payload_.data() + pos_, toRead); + pos_ += toRead; + consecutiveEosReads_ = 0; + } else if (++consecutiveEosReads_ > kMaxConsecutiveEosReads) { + // Throw a plain C++ exception: the reader wraps Read() in + // arrow::Result and drops arrow errors via .ValueOr(0), so an + // arrow::Status::IOError would be swallowed and the loop would spin on. + throw std::runtime_error( + "possible infinite loop: Read() returned 0 for " + std::to_string(kMaxConsecutiveEosReads) + + " consecutive calls"); + } + return toRead; // 0 == EOS when payload exhausted + } + + arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override { + ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(nbytes)); + ARROW_ASSIGN_OR_RAISE(int64_t bytesRead, Read(nbytes, buffer->mutable_data())); + ARROW_RETURN_NOT_OK(buffer->Resize(bytesRead, false)); + buffer->ZeroPadding(); + return std::move(buffer); + } + + private: + static constexpr int32_t kMaxConsecutiveEosReads = 100; + + std::vector<uint8_t> payload_; + int64_t pos_{0}; + bool negativeRead_{false}; + int64_t firstReadLimit_{-1}; + bool firstRead_{true}; + int32_t consecutiveEosReads_{0}; + bool closed_{false}; +}; + +// Append a little-endian POD value to `out` (Presto page header fields are +// machine byte order / little-endian on x86). +template <typename T> +void appendLe(std::vector<uint8_t>& out, T value) { + T v = value; + const auto* p = reinterpret_cast<const uint8_t*>(&v); + out.insert(out.end(), p, p + sizeof(T)); +} + +// Build a truncated Presto compressed page: a valid 21-byte header declaring +// compressedSize bytes of body, but only `bodyBytes` bytes follow. The +// reader's compressed branch calls source->readBytes(buf, compressedSize); +// when EOS is hit mid-drain, GlutenByteInputStream::readBytes loops to +// next(true) which must VELOX_FAIL instead of spinning. +// +// Header layout (PrestoHeader.cpp): numRows:int32, pageCodecMarker:int8, +// uncompressedSize:int32, compressedSize:int32, checksum:int64 == 21 bytes. +// pageCodecMarker = kCompressedBitMask (1), no checksum bit -> actualCheckSum +// stays 0 and matches header.checksum = 0 (PrestoSerializer.cpp:159). +std::vector<uint8_t> buildTruncatedCompressedPage(int32_t compressedSize, int32_t bodyBytes) { + std::vector<uint8_t> out; + out.reserve(21 + bodyBytes); + appendLe<int32_t>(out, 1); // numRows + appendLe<int8_t>(out, 1); // pageCodecMarker = kCompressedBitMask, no checksum + appendLe<int32_t>(out, compressedSize + 64); // uncompressedSize + appendLe<int32_t>(out, compressedSize); + appendLe<int64_t>(out, 0); // checksum + for (int i = 0; i < bodyBytes; ++i) { + out.push_back(static_cast<uint8_t>(i & 0xFF)); + } + return out; +} +} // namespace + +class VeloxShuffleReaderTest : public ::testing::Test, public test::VectorTestBase { + protected: + static void SetUpTestSuite() { + VeloxBackend::create(AllocationListener::noop(), {}); + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + } + + static void TearDownTestSuite() { + VeloxBackend::get()->tearDown(); + } + + std::shared_ptr<VeloxRssSortShuffleReaderDeserializer> makeDeserializer( + const std::shared_ptr<arrow::io::InputStream>& in, + const RowTypePtr& rowType = ROW({"c0"}, {INTEGER()})) { + int64_t deserializeTime = 0; + return std::make_shared<VeloxRssSortShuffleReaderDeserializer>( + std::make_shared<TestStreamReader>(in), + getDefaultMemoryManager(), + rowType, + /*batchSize=*/1024, + common::CompressionKind_NONE, + deserializeTime); + } + + // Serializes `rowVector` into a single uncompressed Presto page + // (21-byte header + payload), the exact wire format the rss-sort writer + // produces. With `withChecksum`, a PrestoOutputStreamListener is attached so + // the writer fills in the checksum bit and CRC (same mechanism as + // VeloxHashShuffleWriter's complex-type flush). NOTE: must not use + // gluten::BufferOutputStream here — its write() ignores the listener. + std::vector<uint8_t> serializePage(const RowVectorPtr& rowVector, bool withChecksum = false) { + serializer::presto::PrestoVectorSerde::PrestoOptions options; + options.compressionKind = common::CompressionKind_NONE; + auto serde = std::make_unique<serializer::presto::PrestoVectorSerde>(); + VectorStreamGroup group(pool(), serde.get()); + group.createStreamTree(asRowType(rowVector->type()), rowVector->size(), &options); + group.append(rowVector); + serializer::presto::PrestoOutputStreamListener listener; + std::stringstream out; + facebook::velox::OStreamOutputStream os(&out, withChecksum ? &listener : nullptr); + group.flush(&os); + const auto str = out.str(); + return std::vector<uint8_t>(str.begin(), str.end()); + } + + // ROW<c0: ARRAY<ROW<a: INTEGER>>> with `numArrays` arrays of + // `elementsPerArray` dense elements each. The nested struct triggers the + // Presto serde's pre-scan (tellp -> scan page -> seekp back). + RowVectorPtr makeNestedArraysRowVector(vector_size_t numArrays, vector_size_t elementsPerArray) { + const vector_size_t numElements = numArrays * elementsPerArray; + auto offsets = AlignedBuffer::allocate<vector_size_t>(numArrays, pool()); + auto sizes = AlignedBuffer::allocate<vector_size_t>(numArrays, pool()); + auto* rawOffsets = offsets->asMutable<vector_size_t>(); + auto* rawSizes = sizes->asMutable<vector_size_t>(); + for (vector_size_t i = 0; i < numArrays; ++i) { + rawOffsets[i] = i * elementsPerArray; + rawSizes[i] = elementsPerArray; + } + auto elements = makeRowVector({makeFlatVector<int32_t>( + numElements, [](vector_size_t row) { return row % 1024; })}); + auto arrayVector = std::make_shared<ArrayVector>( + pool(), + ARRAY(ROW({"a"}, {INTEGER()})), + BufferPtr(nullptr), + numArrays, + offsets, + sizes, + elements); + return makeRowVector({arrayVector}); + } +}; + +// Empty stream (e.g. an empty Celeborn partition): construction must NOT +// throw; next() returns nullptr (graceful EOS). +TEST_F(VeloxShuffleReaderTest, EmptyStreamGracefulEos) { + auto deserializer = makeDeserializer(std::make_shared<FakeInputStream>()); + EXPECT_EQ(deserializer->next(), nullptr); +} + +// Truncated compressed page: header reads fine and construction succeeds, +// but next() drives PrestoVectorSerde::deserialize's compressed branch -> +// readBytes(compressedSize) -> next(true) on EOS -> VELOX_FAIL. Pre-fix this +// is the documented infinite loop; post-fix it throws. See the file header +// note for how FakeInputStream cuts the pre-fix loop short. +TEST_F(VeloxShuffleReaderTest, EosMidPageTerminates) { + auto payload = buildTruncatedCompressedPage(/*compressedSize=*/1000, /*bodyBytes=*/8); + auto deserializer = makeDeserializer(std::make_shared<FakeInputStream>(std::move(payload))); + + VELOX_ASSERT_THROW( + deserializer->next(), "Reading past end of VeloxRssSortShuffleReaderDeserializer::VeloxInputStream"); Review Comment: The assertion matches a very specific exception message that includes a fully-qualified type name. This can be brittle if the wording changes upstream (even when behavior is still correct). Prefer matching a stable substring (e.g., just the core 'Reading past end' phrase) or using a less specific pattern if the test framework supports it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
