From e5c015fd123eb20c3d8710b4a3b04c598959f0c0 Mon Sep 17 00:00:00 2001 From: Lingfeng Zhang Date: Mon, 17 Aug 2026 03:24:31 +0000 Subject: [PATCH 1/2] [GLUTEN-12985][VL] Handle mid-page EOS in VeloxRssSortShuffleReaderDeserializer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GlutenByteInputStream::readBytes() drives next(true) in a for(;;) loop. When a page header declares more bytes than the stream actually holds (truncated partition data), the reader hits EOS mid-page — but VeloxInputStream::next() ignores its throwIfPastEnd argument and silently returns on EOS, so the loop never exits. This PR impl throwIfPassEnd arg of VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::next() --- cpp/velox/shuffle/VeloxShuffleReader.cc | 18 +- cpp/velox/tests/CMakeLists.txt | 2 + cpp/velox/tests/VeloxShuffleReaderTest.cc | 215 ++++++++++++++++++++++ 3 files changed, 228 insertions(+), 7 deletions(-) create mode 100644 cpp/velox/tests/VeloxShuffleReaderTest.cc diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc b/cpp/velox/shuffle/VeloxShuffleReader.cc index d10a0ca3054..86796bda5f2 100644 --- a/cpp/velox/shuffle/VeloxShuffleReader.cc +++ b/cpp/velox/shuffle/VeloxShuffleReader.cc @@ -845,7 +845,7 @@ VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::VeloxInputStream( std::shared_ptr input, facebook::velox::BufferPtr buffer) : in_(std::move(input)), buffer_(std::move(buffer)) { - next(true); + next(false); } bool VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::hasNext() { @@ -853,7 +853,7 @@ bool VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::hasNext() { return false; } if (ranges()[0].position >= ranges()[0].size) { - next(true); + next(false); return offset_ != 0; } return true; @@ -861,11 +861,15 @@ bool VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::hasNext() { void VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::next(bool throwIfPastEnd) { const uint32_t readBytes = buffer_->capacity(); - offset_ = in_->Read(readBytes, buffer_->asMutable()).ValueOr(0); - if (offset_ > 0) { - int32_t realBytes = offset_; - VELOX_CHECK_LT(0, realBytes, "Reading past end of file."); - setRange({buffer_->asMutable(), realBytes, 0}); + offset_ = 0; + int64_t realBytes = in_->Read(readBytes, buffer_->asMutable()).ValueOr(0); + VELOX_CHECK_LE(0, realBytes, "Read returned negative value: {}", realBytes); + if (realBytes > 0) { + offset_ = realBytes; + setRange({buffer_->asMutable(), static_cast(realBytes), 0}); + } else if (throwIfPastEnd) { + VELOX_FAIL( + "Reading past end of VeloxRssSortShuffleReaderDeserializer::VeloxInputStream, real bytes = {}", realBytes); } } diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt index 136884503ab..b3b368703a1 100644 --- a/cpp/velox/tests/CMakeLists.txt +++ b/cpp/velox/tests/CMakeLists.txt @@ -114,6 +114,8 @@ add_velox_test(velox_rss_sort_shuffle_writer_test SOURCES add_velox_test(velox_sort_shuffle_writer_test SOURCES VeloxSortShuffleWriterTest.cc) +add_velox_test(velox_shuffle_reader_test SOURCES VeloxShuffleReaderTest.cc) + # TODO: ORC is not well supported. add_velox_test(orc_test SOURCES OrcTest.cc) add_velox_test( velox_operators_test diff --git a/cpp/velox/tests/VeloxShuffleReaderTest.cc b/cpp/velox/tests/VeloxShuffleReaderTest.cc new file mode 100644 index 00000000000..22d64cf6d29 --- /dev/null +++ b/cpp/velox/tests/VeloxShuffleReaderTest.cc @@ -0,0 +1,215 @@ +/* + * 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; +// + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "compute/VeloxBackend.h" +#include "config/GlutenConfig.h" +#include "memory/VeloxMemoryManager.h" +#include "shuffle/VeloxShuffleReader.h" +#include "tests/utils/TestAllocationListener.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/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. +// +// 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 payload = {}, bool negativeRead = false) + : payload_(std::move(payload)), negativeRead_(negativeRead) {} + + arrow::Status Close() override { + closed_ = true; + return arrow::Status::OK(); + } + arrow::Result Tell() const override { + return pos_; + } + bool closed() const override { + return closed_; + } + + arrow::Result Read(int64_t nbytes, void* out) override { + if (negativeRead_) { + return static_cast(-1); + } + int64_t toRead = std::min(nbytes, static_cast(payload_.size()) - pos_); + 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> Read(int64_t nbytes) override { + GLUTEN_ASSIGN_OR_THROW(auto buffer, arrow::AllocateResizableBuffer(nbytes, arrow::default_memory_pool())); + GLUTEN_ASSIGN_OR_THROW(int64_t bytesRead, Read(nbytes, buffer->mutable_data())); + GLUTEN_THROW_NOT_OK(buffer->Resize(bytesRead, false)); + buffer->ZeroPadding(); + return std::move(buffer); + } + + private: + static constexpr int32_t kMaxConsecutiveEosReads = 100; + + std::vector payload_; + int64_t pos_{0}; + bool negativeRead_{false}; + 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 +void appendLe(std::vector& out, T value) { + T v = value; + const auto* p = reinterpret_cast(&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 buildTruncatedCompressedPage(int32_t compressedSize, int32_t bodyBytes) { + std::vector out; + out.reserve(21 + bodyBytes); + appendLe(out, /*numRows=*/1); + appendLe(out, /*pageCodecMarker=*/1); // kCompressedBitMask, no checksum + appendLe(out, /*uncompressedSize=*/compressedSize + 64); + appendLe(out, compressedSize); + appendLe(out, /*checksum=*/0); + for (int i = 0; i < bodyBytes; ++i) { + out.push_back(static_cast(i & 0xFF)); + } + return out; +} +} // namespace + +class VeloxShuffleReaderTest : public ::testing::Test, public test::VectorTestBase { + protected: + static void SetUpTestCase() { + if (!isRegisteredNamedVectorSerde("Presto")) { + serializer::presto::PrestoVectorSerde::registerNamedVectorSerde(); + } + auto listener = std::make_unique(); + std::unordered_map conf{{kMemoryReservationBlockSize, "1"}, {kDebugModeEnabled, "true"}}; + VeloxBackend::create(std::move(listener), conf); + } + + static void TearDownTestCase() { + VeloxBackend::get()->tearDown(); + } + + std::shared_ptr makeDeserializer(std::shared_ptr in) { + auto streamReader = std::make_shared(std::move(in)); + return std::make_shared( + streamReader, + getDefaultMemoryManager(), + ROW({"c0"}, {INTEGER()}), + /*batchSize=*/1024, + common::CompressionKind_NONE, + deserializeTime_); + } + + int64_t deserializeTime_{0}; +}; + +// 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()); + 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, EosMidPageThrows) { + auto payload = buildTruncatedCompressedPage(/*compressedSize=*/1000, /*bodyBytes=*/8); + auto deserializer = makeDeserializer(std::make_shared(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) { + // A buggy upstream whose Read returns a negative byte count + auto deserializer = + makeDeserializer(std::make_shared(std::vector{}, /*negativeRead=*/true)); + VELOX_ASSERT_THROW(deserializer->next(), "Read returned negative value"); +} + +} // namespace gluten From 5fa48808f98e960eee5fcd581d9c12a6e50dd629 Mon Sep 17 00:00:00 2001 From: Lingfeng Zhang Date: Thu, 17 Sep 2026 16:50:54 +0800 Subject: [PATCH 2/2] Raise read errors via GLUTEN_ASSIGN_OR_THROW --- cpp/velox/shuffle/VeloxShuffleReader.cc | 3 +- cpp/velox/tests/VeloxShuffleReaderTest.cc | 38 ++++++++++------------- 2 files changed, 18 insertions(+), 23 deletions(-) diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc b/cpp/velox/shuffle/VeloxShuffleReader.cc index 86796bda5f2..6c2e101b40a 100644 --- a/cpp/velox/shuffle/VeloxShuffleReader.cc +++ b/cpp/velox/shuffle/VeloxShuffleReader.cc @@ -862,8 +862,7 @@ bool VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::hasNext() { void VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::next(bool throwIfPastEnd) { const uint32_t readBytes = buffer_->capacity(); offset_ = 0; - int64_t realBytes = in_->Read(readBytes, buffer_->asMutable()).ValueOr(0); - VELOX_CHECK_LE(0, realBytes, "Read returned negative value: {}", realBytes); + GLUTEN_ASSIGN_OR_THROW(int64_t realBytes, in_->Read(readBytes, buffer_->asMutable())); if (realBytes > 0) { offset_ = realBytes; setRange({buffer_->asMutable(), static_cast(realBytes), 0}); diff --git a/cpp/velox/tests/VeloxShuffleReaderTest.cc b/cpp/velox/tests/VeloxShuffleReaderTest.cc index 22d64cf6d29..7cfab92c788 100644 --- a/cpp/velox/tests/VeloxShuffleReaderTest.cc +++ b/cpp/velox/tests/VeloxShuffleReaderTest.cc @@ -22,7 +22,7 @@ // 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; +// - a buggy upstream whose Read() returns an error status; // #include @@ -56,9 +56,8 @@ 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. +// the payload is exhausted, Read returns 0 (EOS). With `errorRead`, Read +// returns an IOError instead, modeling a buggy upstream that fails the read. // // To keep a possible reader-side infinite loop (readBytes -> next() -> EOS -> // silently return -> spin) from hanging the test until the CI timeout, Read @@ -66,8 +65,8 @@ namespace { // 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 payload = {}, bool negativeRead = false) - : payload_(std::move(payload)), negativeRead_(negativeRead) {} + explicit FakeInputStream(std::vector payload = {}, bool errorRead = false) + : payload_(std::move(payload)), errorRead_(errorRead) {} arrow::Status Close() override { closed_ = true; @@ -81,8 +80,8 @@ class FakeInputStream final : public arrow::io::InputStream { } arrow::Result Read(int64_t nbytes, void* out) override { - if (negativeRead_) { - return static_cast(-1); + if (errorRead_) { + return arrow::Status::IOError("fake upstream read failure"); } int64_t toRead = std::min(nbytes, static_cast(payload_.size()) - pos_); if (toRead > 0) { @@ -90,9 +89,9 @@ class FakeInputStream final : public arrow::io::InputStream { 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 a plain C++ exception as a loop guard: the EOS contract itself + // must not be an error (a clean 0 return is legitimate), so this is the + // only way to cut a possible spin short. throw std::runtime_error( "possible infinite loop: Read() returned 0 for " + std::to_string(kMaxConsecutiveEosReads) + " consecutive calls"); @@ -113,7 +112,7 @@ class FakeInputStream final : public arrow::io::InputStream { std::vector payload_; int64_t pos_{0}; - bool negativeRead_{false}; + bool errorRead_{false}; int32_t consecutiveEosReads_{0}; bool closed_{false}; }; @@ -201,15 +200,12 @@ TEST_F(VeloxShuffleReaderTest, EosMidPageThrows) { 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) { - // A buggy upstream whose Read returns a negative byte count - auto deserializer = - makeDeserializer(std::make_shared(std::vector{}, /*negativeRead=*/true)); - VELOX_ASSERT_THROW(deserializer->next(), "Read returned negative value"); +// A buggy upstream whose Read returns an error status. The reader must +// propagate it instead of swallowing it. +TEST_F(VeloxShuffleReaderTest, ErrorReadThrows) { + auto deserializer = makeDeserializer(std::make_shared(std::vector{}, /*errorRead=*/true)); + + EXPECT_THROW((void)deserializer->next(), GlutenException); } } // namespace gluten