diff --git a/velox/dwio/common/OnDemandUnitLoader.cpp b/velox/dwio/common/OnDemandUnitLoader.cpp index 5a45e396f7d5b..ee21a2b442338 100644 --- a/velox/dwio/common/OnDemandUnitLoader.cpp +++ b/velox/dwio/common/OnDemandUnitLoader.cpp @@ -67,7 +67,7 @@ class OnDemandUnitLoader : public UnitLoader { void onSeek(uint32_t unit, uint64_t rowOffsetInUnit) override { VELOX_CHECK_LT(unit, loadUnits_.size(), "Unit out of range"); - VELOX_CHECK_LT( + VELOX_CHECK_LE( rowOffsetInUnit, loadUnits_[unit]->getNumRows(), "Row out of range"); } diff --git a/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp b/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp index 082bc18733b85..a491fe07ff07c 100644 --- a/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp +++ b/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp @@ -121,19 +121,37 @@ TEST(OnDemandUnitLoaderTests, CanSeek) { EXPECT_EQ(blockedOnIoCount, 3); } -TEST(OnDemandUnitLoaderTests, SeekOutOfRange) { +TEST(OnDemandUnitLoaderTests, SeekOutOfRangeReaderError) { size_t blockedOnIoCount = 0; OnDemandUnitLoaderFactory factory([&](auto) { ++blockedOnIoCount; }); ReaderMock readerMock{{10, 20, 30}, {0, 0, 0}, factory}; EXPECT_EQ(readerMock.unitsLoaded(), std::vector({false, false, false})); EXPECT_EQ(blockedOnIoCount, 0); + readerMock.seek(60); + + EXPECT_THAT( + [&]() { readerMock.seek(61); }, + Throws(Property( + &facebook::velox::VeloxRuntimeError::message, + HasSubstr("Can't seek to possition 61 in file. Must be up to 60.")))); +} + +TEST(OnDemandUnitLoaderTests, SeekOutOfRange) { + OnDemandUnitLoaderFactory factory(nullptr); + std::vector unitsLoaded(getUnitsLoadedWithFalse(1)); + std::vector> units; + units.push_back(std::make_unique(10, 0, unitsLoaded, 0)); + + auto unitLoader = factory.create(std::move(units)); + + unitLoader->onSeek(0, 10); + EXPECT_THAT( - [&]() { readerMock.seek(60); }, + [&]() { unitLoader->onSeek(0, 11); }, Throws(Property( &facebook::velox::VeloxRuntimeError::message, - HasSubstr( - "Can't seek to possition 60 in file. Must be less than: 60.")))); + HasSubstr("Row out of range")))); } TEST(OnDemandUnitLoaderTests, UnitOutOfRange) { diff --git a/velox/dwio/common/tests/utils/UnitLoaderTestTools.cpp b/velox/dwio/common/tests/utils/UnitLoaderTestTools.cpp index 4b49c7c3d3751..ed1e9417d48bf 100644 --- a/velox/dwio/common/tests/utils/UnitLoaderTestTools.cpp +++ b/velox/dwio/common/tests/utils/UnitLoaderTestTools.cpp @@ -56,8 +56,10 @@ void ReaderMock::seek(uint64_t rowNumber) { rowsLeft -= rowCount; totalRows += rowCount; } - VELOX_FAIL( - "Can't seek to possition {} in file. Must be less than: {}.", + VELOX_CHECK_EQ( + rowsLeft, + 0, + "Can't seek to possition {} in file. Must be up to {}.", rowNumber, totalRows); }