This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 48b5fca64d [GLUTEN-12985][VL] Handle mid-page EOS in
VeloxRssSortShuffleReaderDeserializer (#12986)
48b5fca64d is described below
commit 48b5fca64d72b863960bca8129d11cbfdfae03d8
Author: Lingfeng Zhang <[email protected]>
AuthorDate: Thu Sep 17 21:48:07 2026 +0800
[GLUTEN-12985][VL] Handle mid-page EOS in
VeloxRssSortShuffleReaderDeserializer (#12986)
* [GLUTEN-12985][VL] Handle mid-page EOS in
VeloxRssSortShuffleReaderDeserializer
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()
* Raise read errors via GLUTEN_ASSIGN_OR_THROW
---------
Co-authored-by: Yuan <[email protected]>
---
cpp/velox/shuffle/VeloxShuffleReader.cc | 17 ++-
cpp/velox/tests/CMakeLists.txt | 2 +
cpp/velox/tests/VeloxShuffleReaderTest.cc | 211 ++++++++++++++++++++++++++++++
3 files changed, 223 insertions(+), 7 deletions(-)
diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc
b/cpp/velox/shuffle/VeloxShuffleReader.cc
index d10a0ca305..6c2e101b40 100644
--- a/cpp/velox/shuffle/VeloxShuffleReader.cc
+++ b/cpp/velox/shuffle/VeloxShuffleReader.cc
@@ -845,7 +845,7 @@
VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::VeloxInputStream(
std::shared_ptr<arrow::io::InputStream> 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,14 @@ bool
VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::hasNext() {
void VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::next(bool
throwIfPastEnd) {
const uint32_t readBytes = buffer_->capacity();
- offset_ = in_->Read(readBytes, buffer_->asMutable<char>()).ValueOr(0);
- if (offset_ > 0) {
- int32_t realBytes = offset_;
- VELOX_CHECK_LT(0, realBytes, "Reading past end of file.");
- setRange({buffer_->asMutable<uint8_t>(), realBytes, 0});
+ offset_ = 0;
+ GLUTEN_ASSIGN_OR_THROW(int64_t realBytes, in_->Read(readBytes,
buffer_->asMutable<char>()));
+ if (realBytes > 0) {
+ offset_ = realBytes;
+ setRange({buffer_->asMutable<uint8_t>(), static_cast<int32_t>(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 136884503a..b3b368703a 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 0000000000..7cfab92c78
--- /dev/null
+++ b/cpp/velox/tests/VeloxShuffleReaderTest.cc
@@ -0,0 +1,211 @@
+/*
+ * 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 an error status;
+//
+
+#include <gtest/gtest.h>
+
+#include <arrow/buffer.h>
+#include <arrow/io/interfaces.h>
+#include <arrow/memory_pool.h>
+
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#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 `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
+// 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 errorRead =
false)
+ : payload_(std::move(payload)), errorRead_(errorRead) {}
+
+ 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 (errorRead_) {
+ return arrow::Status::IOError("fake upstream read failure");
+ }
+ int64_t toRead = std::min<int64_t>(nbytes,
static_cast<int64_t>(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 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");
+ }
+ return toRead; // 0 == EOS when payload exhausted
+ }
+
+ arrow::Result<std::shared_ptr<arrow::Buffer>> 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<uint8_t> payload_;
+ int64_t pos_{0};
+ bool errorRead_{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 <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, /*numRows=*/1);
+ appendLe<int8_t>(out, /*pageCodecMarker=*/1); // kCompressedBitMask, no
checksum
+ appendLe<int32_t>(out, /*uncompressedSize=*/compressedSize + 64);
+ appendLe<int32_t>(out, compressedSize);
+ appendLe<int64_t>(out, /*checksum=*/0);
+ 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 SetUpTestCase() {
+ if (!isRegisteredNamedVectorSerde("Presto")) {
+ serializer::presto::PrestoVectorSerde::registerNamedVectorSerde();
+ }
+ auto listener = std::make_unique<TestAllocationListener>();
+ std::unordered_map<std::string, std::string>
conf{{kMemoryReservationBlockSize, "1"}, {kDebugModeEnabled, "true"}};
+ VeloxBackend::create(std::move(listener), conf);
+ }
+
+ static void TearDownTestCase() {
+ VeloxBackend::get()->tearDown();
+ }
+
+ std::shared_ptr<VeloxRssSortShuffleReaderDeserializer>
makeDeserializer(std::shared_ptr<arrow::io::InputStream> in) {
+ auto streamReader = std::make_shared<TestStreamReader>(std::move(in));
+ return std::make_shared<VeloxRssSortShuffleReaderDeserializer>(
+ 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<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, EosMidPageThrows) {
+ 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 an error status. The reader must
+// propagate it instead of swallowing it.
+TEST_F(VeloxShuffleReaderTest, ErrorReadThrows) {
+ auto deserializer =
makeDeserializer(std::make_shared<FakeInputStream>(std::vector<uint8_t>{},
/*errorRead=*/true));
+
+ EXPECT_THROW((void)deserializer->next(), GlutenException);
+}
+
+} // namespace gluten
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
