From 257ccee073890849ed8fa3a65988c693ccf42d37 Mon Sep 17 00:00:00 2001 From: lizhen-0710 Date: Fri, 31 Jul 2026 12:23:01 +0800 Subject: [PATCH 1/2] fix: avoid overlapping scan output buffers HiveDataSource keeps a backing reader output RowVector in output_ while returning a projected RowVector to the operator pipeline. The returned RowVector holds references to the child vectors from output_. If output_ is reused for the next splitReader::next() call, HiveDataSource also keeps the previous backing RowVector alive while the next batch allocates selective column values. This makes the previous output children and current output allocation overlap, which can temporarily double scan-side memory. Recreate the backing reader output before every splitReader::next() call so HiveDataSource drops its reference to the previous backing RowVector before the next batch allocation. The returned output still owns the previous child vectors, but HiveDataSource no longer adds an extra reference that extends their lifetime into the next read. Also bound SelectiveColumnReader values_ reuse to buffers no larger than 1.5x the current request. This preserves useful same-size reuse while allowing oversized buffers from larger batches to be released before smaller batches allocate replacement values. Add targeted tests covering the HiveDataSource backing RowVector lifetime and the SelectiveColumnReader reuse, oversize-release, and shared-buffer paths. Co-authored-by: TRAE CLI --- bolt/connectors/hive/HiveDataSource.cpp | 4 +- .../hive/tests/HiveConnectorTest.cpp | 141 ++++++++++++++++++ .../common/SelectiveColumnReaderInternal.h | 8 +- bolt/dwio/common/tests/ReaderTest.cpp | 116 ++++++++++++++ 4 files changed, 263 insertions(+), 6 deletions(-) diff --git a/bolt/connectors/hive/HiveDataSource.cpp b/bolt/connectors/hive/HiveDataSource.cpp index 1929c0c41..51281d7bd 100644 --- a/bolt/connectors/hive/HiveDataSource.cpp +++ b/bolt/connectors/hive/HiveDataSource.cpp @@ -639,9 +639,7 @@ std::optional HiveDataSource::next( return nullptr; } - if (!output_) { - output_ = BaseVector::create(readerOutputType_, 0, pool_); - } + output_ = BaseVector::create(readerOutputType_, 0, pool_); // TODO Check if remaining filter has a conjunct that doesn't depend on // any column, e.g. rand() < 0.1. Evaluate that conjunct first, then scan diff --git a/bolt/connectors/hive/tests/HiveConnectorTest.cpp b/bolt/connectors/hive/tests/HiveConnectorTest.cpp index 5b9ba92bc..766a12d4d 100644 --- a/bolt/connectors/hive/tests/HiveConnectorTest.cpp +++ b/bolt/connectors/hive/tests/HiveConnectorTest.cpp @@ -35,6 +35,7 @@ #include "bolt/connectors/hive/HiveConfig.h" #include "bolt/connectors/hive/HiveConnectorUtil.h" #include "bolt/connectors/hive/HiveDataSource.h" +#include "bolt/dwio/common/Statistics.h" #include "bolt/expression/ExprToSubfieldFilter.h" namespace bytedance::bolt::connector::hive { namespace { @@ -47,6 +48,101 @@ class HiveConnectorTest : public exec::test::HiveConnectorTestBase { memory::memoryManager()->addLeafPool(); }; +class TrackingSplitReader : public HiveSplitReaderBase { + public: + TrackingSplitReader(RowTypePtr rowType, memory::MemoryPool* pool) + : rowType_(std::move(rowType)), pool_(pool) {} + + uint64_t next(int64_t /*size*/, VectorPtr& output) override { + if (calls_ == 1) { + auto firstChild = firstChild_.lock(); + EXPECT_NE(firstChild, nullptr); + if (firstChild == nullptr) { + return 0; + } + // The temporary 'firstChild' reference and the previous batch returned + // by HiveDataSource should be the only remaining references. If + // HiveDataSource still keeps its previous backing output_, this count is + // higher and the next batch overlaps with previous child values. + EXPECT_EQ(firstChild.use_count(), 2); + } + if (calls_ >= 2) { + return 0; + } + + auto values = AlignedBuffer::allocate(1, pool_); + values->asMutable()[0] = calls_; + values->setSize(sizeof(int64_t)); + auto child = std::make_shared>( + pool_, BIGINT(), nullptr, 1, values, std::vector{}); + + if (calls_ == 0) { + firstChild_ = child; + } + + output = std::make_shared( + pool_, rowType_, nullptr, 1, std::vector{child}); + ++calls_; + return 1; + } + + bool allPrefetchIssued() const override { + return true; + } + + bool emptySplit() const override { + return false; + } + + void resetFilterCaches() override {} + + int64_t estimatedRowSize() const override { + return sizeof(int64_t); + } + + void updateRuntimeStats(dwio::common::RuntimeStatistics&) const override {} + + void resetSplit() override {} + + private: + RowTypePtr rowType_; + memory::MemoryPool* pool_; + int32_t calls_{0}; + std::weak_ptr firstChild_; +}; + +class TestingHiveDataSource : public HiveDataSource { + public: + TestingHiveDataSource( + const RowTypePtr& outputType, + const std::shared_ptr& tableHandle, + const std::unordered_map< + std::string, + std::shared_ptr>& columnHandles, + FileHandleFactory* fileHandleFactory, + const core::QueryConfig& queryConfig, + folly::Executor* executor, + const std::shared_ptr& connectorQueryCtx, + const std::shared_ptr& hiveConfig) + : HiveDataSource( + outputType, + tableHandle, + columnHandles, + fileHandleFactory, + queryConfig, + executor, + connectorQueryCtx, + hiveConfig) {} + + void setSplitReader(std::unique_ptr splitReader) { + splitReader_ = std::move(splitReader); + } + + void setSplit(std::shared_ptr split) { + split_ = std::move(split); + } +}; + void validateNullConstant(const ScanSpec& spec, const Type& type) { ASSERT_TRUE(spec.isConstant()); auto constant = spec.constantValue(); @@ -90,6 +186,51 @@ TEST_F(HiveConnectorTest, hiveConfig) { "UNKNOWN BEHAVIOR 100"); } +TEST_F(HiveConnectorTest, nextReleasesPreviousReaderOutput) { + auto rowType = ROW({"c0"}, {BIGINT()}); + auto tableHandle = makeTableHandle({}, nullptr, "test_table", rowType); + ColumnHandleMap assignments = {{"c0", regularColumn("c0", BIGINT())}}; + auto hiveConfig = + std::make_shared(std::make_shared( + std::unordered_map{})); + auto connectorQueryCtx = std::make_shared( + pool_.get(), + pool_.get(), + hiveConfig->config().get(), + nullptr, + nullptr, + nullptr, + nullptr, + "query.HiveConnectorTest", + "task.HiveConnectorTest", + "planNodeId.HiveConnectorTest", + 0); + + TestingHiveDataSource dataSource( + rowType, + tableHandle, + assignments, + nullptr, + core::QueryConfig(std::unordered_map{}), + nullptr, + connectorQueryCtx, + hiveConfig); + dataSource.setSplit(HiveConnectorSplitBuilder("unused") + .connectorId(kHiveConnectorId) + .build()); + dataSource.setSplitReader( + std::make_unique(rowType, pool_.get())); + + ContinueFuture future; + auto first = dataSource.next(1, future); + ASSERT_TRUE(first.has_value()); + ASSERT_NE(first.value(), nullptr); + + auto second = dataSource.next(1, future); + ASSERT_TRUE(second.has_value()); + ASSERT_NE(second.value(), nullptr); +} + TEST_F(HiveConnectorTest, makeScanSpec_requiredSubfields_multilevel) { auto columnType = ROW( {{"c0c0", BIGINT()}, diff --git a/bolt/dwio/common/SelectiveColumnReaderInternal.h b/bolt/dwio/common/SelectiveColumnReaderInternal.h index d9358bc46..d1e6ea242 100644 --- a/bolt/dwio/common/SelectiveColumnReaderInternal.h +++ b/bolt/dwio/common/SelectiveColumnReaderInternal.h @@ -49,11 +49,13 @@ bolt::common::AlwaysTrue& alwaysTrue(); template void SelectiveColumnReader::ensureValuesCapacity(vector_size_t numRows) { - if (values_ && values_->unique() && - values_->capacity() >= - BaseVector::byteSize(numRows) + simd::kPadding) { + const auto requiredBytes = BaseVector::byteSize(numRows) + simd::kPadding; + if (values_ && values_->unique() && values_->capacity() >= requiredBytes && + values_->capacity() <= requiredBytes + requiredBytes / 2) { return; } + values_.reset(); + rawValues_ = nullptr; values_ = AlignedBuffer::allocate( numRows + (simd::kPadding / sizeof(T)), &memoryPool_); rawValues_ = values_->asMutable(); diff --git a/bolt/dwio/common/tests/ReaderTest.cpp b/bolt/dwio/common/tests/ReaderTest.cpp index f07d6b527..c35a6208c 100644 --- a/bolt/dwio/common/tests/ReaderTest.cpp +++ b/bolt/dwio/common/tests/ReaderTest.cpp @@ -30,6 +30,7 @@ #include "bolt/dwio/common/Reader.h" #include "bolt/common/base/tests/GTestUtils.h" +#include "bolt/dwio/common/SelectiveColumnReaderInternal.h" #include "bolt/type/Subfield.h" #include "bolt/vector/tests/utils/VectorTestBase.h" @@ -45,6 +46,87 @@ class ReaderTest : public testing::Test, public test::VectorTestBase { } }; +class TestFormatData : public FormatData { + public: + void readNulls( + vector_size_t /*numValues*/, + const uint64_t* FOLLY_NULLABLE /*incomingNulls*/, + BufferPtr& nulls, + bool /*nullsOnly*/ = false) override { + nulls = nullptr; + } + + uint64_t skipNulls(uint64_t numValues, bool /*nullsOnly*/ = false) override { + return numValues; + } + + uint64_t skip(uint64_t numValues) override { + return numValues; + } + + bool hasNulls() const override { + return false; + } + + PositionProvider seekToRowGroup(int64_t /*index*/) override { + static const std::vector kEmptyPositions; + return PositionProvider(kEmptyPositions); + } + + void filterRowGroups( + const bolt::common::ScanSpec& /*scanSpec*/, + uint64_t /*rowsPerRowGroup*/, + const StatsContext& /*writerContext*/, + FilterRowGroupsResult& /*result*/, + BufferedInput& /*input*/) override {} +}; + +class TestFormatParams : public FormatParams { + public: + TestFormatParams(memory::MemoryPool& pool, ColumnReaderStatistics& stats) + : FormatParams(pool, stats) {} + + std::unique_ptr toFormatData( + const std::shared_ptr& /*type*/, + const bolt::common::ScanSpec& /*scanSpec*/) override { + return std::make_unique(); + } +}; + +class TestSelectiveColumnReader : public SelectiveColumnReader { + public: + TestSelectiveColumnReader( + const TypePtr& type, + FormatParams& params, + bolt::common::ScanSpec& scanSpec) + : SelectiveColumnReader( + type, + TypeWithId::create(type), + params, + scanSpec) {} + + void read( + int64_t /*offset*/, + const RowSet& /*rows*/, + const uint64_t* /*incomingNulls*/) override {} + + void getValues(const RowSet& /*rows*/, VectorPtr* FOLLY_NONNULL /*result*/) + override {} + + template + void ensureCapacity(vector_size_t numRows) { + ensureValuesCapacity(numRows); + } + + Buffer* values() const { + return values_.get(); + } + + BufferPtr retainedValues() const { + return values_; + } +}; + TEST_F(ReaderTest, projectColumnsFilterStruct) { constexpr int kSize = 10; auto input = makeRowVector({ @@ -64,6 +146,40 @@ TEST_F(ReaderTest, projectColumnsFilterStruct) { test::assertEqualVectors(expected, actual); } +TEST_F(ReaderTest, selectiveValuesReuseBound) { + ColumnReaderStatistics stats; + bolt::common::ScanSpec scanSpec("c0"); + TestFormatParams params(*pool_, stats); + TestSelectiveColumnReader reader(BIGINT(), params, scanSpec); + + reader.ensureCapacity(100); + auto* values = reader.values(); + ASSERT_NE(values, nullptr); + + reader.ensureCapacity(100); + EXPECT_EQ(reader.values(), values); + + reader.ensureCapacity(99); + EXPECT_EQ(reader.values(), values); + auto maxReusableBytes = BaseVector::byteSize(99) + simd::kPadding; + maxReusableBytes += maxReusableBytes / 2; + EXPECT_LE(reader.values()->capacity(), maxReusableBytes); + + const auto bytesWithOversizedValues = pool_->currentBytes(); + reader.ensureCapacity(60); + EXPECT_LT(pool_->currentBytes(), bytesWithOversizedValues); + maxReusableBytes = BaseVector::byteSize(60) + simd::kPadding; + maxReusableBytes += maxReusableBytes / 2; + EXPECT_LE(reader.values()->capacity(), maxReusableBytes); + + values = reader.values(); + auto sharedValues = reader.retainedValues(); + EXPECT_FALSE(sharedValues->unique()); + reader.ensureCapacity(60); + EXPECT_TRUE(sharedValues->unique()); + EXPECT_NE(reader.values(), values); +} + TEST_F(ReaderTest, projectColumnsFilterArray) { constexpr int kSize = 10; auto input = makeRowVector({ From f17eb3e40e1dad8f05e18e181e905aa293622ce2 Mon Sep 17 00:00:00 2001 From: lizhen-0710 Date: Sun, 2 Aug 2026 12:00:05 +0800 Subject: [PATCH 2/2] fix: reuse scan output shell safely Refine the scan output lifetime fix by reusing the top-level HiveDataSource backing RowVector shell instead of rebuilding a full empty vector tree before every next(). This keeps the previous batch child backing from overlapping with the next batch allocation while reducing per-batch empty vector allocations. The backing output now uses the minimal shape required by selective readers: non-ROW children are reset to null so leaf, array, and map readers recreate their outputs; ROW children keep empty RowVector shells because struct readers require a non-null result. The initial backing output uses the same minimal skeleton, so the first batch and type-change paths avoid full recursive empty-vector construction too. Return an independent zero-column RowVector for empty scan output so the reusable backing shell is never exposed to downstream operators. Also allow synthesized row-index output to start from a null child slot before filling values. Add coverage for backing shell reuse, child reset shape, zero-column output isolation, row-index output, and real scan paths with leaf, string, ROW, ARRAY, MAP, and nested complex types. Co-authored-by: TRAE CLI --- bolt/connectors/hive/HiveDataSource.cpp | 55 +++++- bolt/connectors/hive/HiveDataSource.h | 2 + .../hive/tests/HiveConnectorTest.cpp | 178 ++++++++++++++++-- .../common/SelectiveStructColumnReader.cpp | 25 ++- bolt/exec/tests/TableScanTest.cpp | 57 ++++++ 5 files changed, 296 insertions(+), 21 deletions(-) diff --git a/bolt/connectors/hive/HiveDataSource.cpp b/bolt/connectors/hive/HiveDataSource.cpp index 51281d7bd..c2b424f7c 100644 --- a/bolt/connectors/hive/HiveDataSource.cpp +++ b/bolt/connectors/hive/HiveDataSource.cpp @@ -600,6 +600,37 @@ void HiveDataSource::addSplit( splitReader_ = createConfiguredSplitReader(split, false); } +namespace { +void resetRowVectorChildren( + const RowType& rowType, + memory::MemoryPool* pool, + std::vector& children) { + BOLT_CHECK_EQ(children.size(), rowType.size()); + for (auto i = 0; i < rowType.size(); ++i) { + const auto& childType = rowType.childAt(i); + if (childType->isRow()) { + // Struct readers expect a non-null RowVector result. Other readers can + // recreate their result from a null child slot. + std::vector innerChildren(childType->size()); + resetRowVectorChildren(childType->asRow(), pool, innerChildren); + children[i] = std::make_shared( + pool, childType, nullptr, 0, std::move(innerChildren)); + } else { + children[i].reset(); + } + } +} + +RowVectorPtr createBackingOutput( + const RowTypePtr& rowType, + memory::MemoryPool* pool) { + std::vector children(rowType->size()); + resetRowVectorChildren(rowType->asRow(), pool, children); + return std::make_shared( + pool, rowType, nullptr, 0, std::move(children)); +} +} // namespace + vector_size_t getLength(std::shared_ptr& split) { auto paimonConnectorSplit = std::dynamic_pointer_cast(split); @@ -620,6 +651,20 @@ vector_size_t getLength(std::shared_ptr& split) { BOLT_FAIL("Unsupported split type for getting length"); } +void HiveDataSource::resetBackingOutput() { + if (!output_ || *output_->type() != *readerOutputType_) { + output_ = createBackingOutput(readerOutputType_, pool_); + return; + } + + auto rowOutput = output_->as(); + BOLT_CHECK_NOT_NULL(rowOutput); + resetRowVectorChildren( + readerOutputType_->asRow(), pool_, rowOutput->children()); + rowOutput->unsafeResize(0); + rowOutput->updateContainsLazyNotLoaded(); +} + std::optional HiveDataSource::next( uint64_t size, bolt::ContinueFuture& /*future*/) { @@ -639,7 +684,7 @@ std::optional HiveDataSource::next( return nullptr; } - output_ = BaseVector::create(readerOutputType_, 0, pool_); + resetBackingOutput(); // TODO Check if remaining filter has a conjunct that doesn't depend on // any column, e.g. rand() < 0.1. Evaluate that conjunct first, then scan @@ -662,7 +707,7 @@ std::optional HiveDataSource::next( auto rowsRemaining = output_->size(); if (rowsRemaining == 0) { // no rows passed the pushed down filters. - output_->prepareForReuse(); + resetBackingOutput(); return getEmptyOutput(); } @@ -678,7 +723,7 @@ std::optional HiveDataSource::next( BOLT_CHECK_LE(rowsRemaining, rowsScanned); if (rowsRemaining == 0) { // No rows passed the remaining filter. - output_->prepareForReuse(); + resetBackingOutput(); return getEmptyOutput(); } @@ -689,8 +734,8 @@ std::optional HiveDataSource::next( } if (outputType_->size() == 0) { - return exec::wrapAndCombineDict( - rowsRemaining, remainingIndices, rowVector); + return std::make_shared( + pool_, outputType_, nullptr, rowsRemaining, std::vector{}); } std::vector outputColumns; diff --git a/bolt/connectors/hive/HiveDataSource.h b/bolt/connectors/hive/HiveDataSource.h index 4757d5a30..0fd2f1337 100644 --- a/bolt/connectors/hive/HiveDataSource.h +++ b/bolt/connectors/hive/HiveDataSource.h @@ -145,6 +145,8 @@ class HiveDataSource : public DataSource { std::shared_ptr ioStats_; private: + void resetBackingOutput(); + // Evaluates remainingFilter_ on the specified vector. Returns number of // rows passed. Populates filterEvalCtx_.selectedIndices and selectedBits // if only some rows passed the filter. If none or all rows passed diff --git a/bolt/connectors/hive/tests/HiveConnectorTest.cpp b/bolt/connectors/hive/tests/HiveConnectorTest.cpp index 766a12d4d..d0b892c89 100644 --- a/bolt/connectors/hive/tests/HiveConnectorTest.cpp +++ b/bolt/connectors/hive/tests/HiveConnectorTest.cpp @@ -50,11 +50,53 @@ class HiveConnectorTest : public exec::test::HiveConnectorTestBase { class TrackingSplitReader : public HiveSplitReaderBase { public: - TrackingSplitReader(RowTypePtr rowType, memory::MemoryPool* pool) - : rowType_(std::move(rowType)), pool_(pool) {} + TrackingSplitReader( + RowTypePtr rowType, + memory::MemoryPool* pool, + std::vector outputSizes = {1, 1}) + : rowType_(std::move(rowType)), + pool_(pool), + outputSizes_(std::move(outputSizes)) {} uint64_t next(int64_t /*size*/, VectorPtr& output) override { + auto rowOutput = std::dynamic_pointer_cast(output); + EXPECT_NE(rowOutput, nullptr); + if (rowOutput == nullptr) { + return 0; + } + EXPECT_EQ(rowOutput->children().size(), rowType_->size()); + + if (rowType_->size() == 0) { + if (calls_ == 1) { + EXPECT_EQ(rowOutput.get(), backingOutput_); + EXPECT_EQ(rowOutput->size(), 0); + } else { + backingOutput_ = rowOutput.get(); + } + if (calls_ >= outputSizes_.size()) { + return 0; + } + const auto outputSize = outputSizes_[calls_]; + rowOutput->unsafeResize(outputSize); + ++calls_; + return outputSize; + } + if (calls_ == 1) { + EXPECT_EQ(rowOutput.get(), backingOutput_); + EXPECT_EQ(rowOutput->size(), 0); + EXPECT_EQ(rowOutput->childAt(0), nullptr); + auto nestedRow = + std::dynamic_pointer_cast(rowOutput->childAt(1)); + EXPECT_NE(nestedRow, nullptr); + if (nestedRow == nullptr) { + return 0; + } + EXPECT_EQ(nestedRow->size(), 0); + EXPECT_EQ(nestedRow->childAt(0), nullptr); + EXPECT_EQ(rowOutput->childAt(2), nullptr); + EXPECT_EQ(rowOutput->childAt(3), nullptr); + auto firstChild = firstChild_.lock(); EXPECT_NE(firstChild, nullptr); if (firstChild == nullptr) { @@ -65,25 +107,71 @@ class TrackingSplitReader : public HiveSplitReaderBase { // HiveDataSource still keeps its previous backing output_, this count is // higher and the next batch overlaps with previous child values. EXPECT_EQ(firstChild.use_count(), 2); + } else { + backingOutput_ = rowOutput.get(); } - if (calls_ >= 2) { + if (calls_ >= outputSizes_.size()) { return 0; } + const auto outputSize = outputSizes_[calls_]; - auto values = AlignedBuffer::allocate(1, pool_); - values->asMutable()[0] = calls_; - values->setSize(sizeof(int64_t)); - auto child = std::make_shared>( - pool_, BIGINT(), nullptr, 1, values, std::vector{}); + auto makeChild = [&](int64_t value) { + auto values = AlignedBuffer::allocate(outputSize, pool_); + for (auto i = 0; i < outputSize; ++i) { + values->asMutable()[i] = value + i; + } + values->setSize(outputSize * sizeof(int64_t)); + return std::make_shared>( + pool_, + BIGINT(), + nullptr, + outputSize, + values, + std::vector{}); + }; + + auto child = makeChild(calls_); + auto nestedChild = makeChild(calls_ + 10); + auto arrayChild = makeChild(calls_ + 20); + auto mapKeyChild = makeChild(calls_ + 30); + auto mapValueChild = makeChild(calls_ + 40); if (calls_ == 0) { firstChild_ = child; } - output = std::make_shared( - pool_, rowType_, nullptr, 1, std::vector{child}); + rowOutput->unsafeResize(outputSize); + rowOutput->childAt(0) = child; + rowOutput->childAt(1) = std::make_shared( + pool_, + rowType_->childAt(1), + nullptr, + outputSize, + std::vector{nestedChild}); + auto offsets = AlignedBuffer::allocate(outputSize, pool_); + auto* rawOffsets = offsets->asMutable(); + std::iota(rawOffsets, rawOffsets + outputSize, 0); + auto sizes = AlignedBuffer::allocate(outputSize, pool_, 1); + rowOutput->childAt(2) = std::make_shared( + pool_, + rowType_->childAt(2), + nullptr, + outputSize, + offsets, + sizes, + arrayChild); + rowOutput->childAt(3) = std::make_shared( + pool_, + rowType_->childAt(3), + nullptr, + outputSize, + offsets, + sizes, + mapKeyChild, + mapValueChild); + rowOutput->updateContainsLazyNotLoaded(); ++calls_; - return 1; + return outputSize; } bool allPrefetchIssued() const override { @@ -107,7 +195,9 @@ class TrackingSplitReader : public HiveSplitReaderBase { private: RowTypePtr rowType_; memory::MemoryPool* pool_; + std::vector outputSizes_; int32_t calls_{0}; + RowVector* backingOutput_{nullptr}; std::weak_ptr firstChild_; }; @@ -186,10 +276,20 @@ TEST_F(HiveConnectorTest, hiveConfig) { "UNKNOWN BEHAVIOR 100"); } -TEST_F(HiveConnectorTest, nextReleasesPreviousReaderOutput) { - auto rowType = ROW({"c0"}, {BIGINT()}); +TEST_F( + HiveConnectorTest, + backingOutputResetReleasesPreviousChildBeforeNextRead) { + auto rowType = + ROW({"c0", "c1", "c2", "c3"}, + {BIGINT(), + ROW({"n0"}, {BIGINT()}), + ARRAY(BIGINT()), + MAP(BIGINT(), BIGINT())}); auto tableHandle = makeTableHandle({}, nullptr, "test_table", rowType); - ColumnHandleMap assignments = {{"c0", regularColumn("c0", BIGINT())}}; + ColumnHandleMap assignments; + for (const auto& name : rowType->names()) { + assignments[name] = regularColumn(name, rowType->findChild(name)); + } auto hiveConfig = std::make_shared(std::make_shared( std::unordered_map{})); @@ -231,6 +331,56 @@ TEST_F(HiveConnectorTest, nextReleasesPreviousReaderOutput) { ASSERT_NE(second.value(), nullptr); } +TEST_F(HiveConnectorTest, emptyOutputDoesNotExposeBackingOutput) { + auto rowType = ROW({}, {}); + auto tableHandle = makeTableHandle({}, nullptr, "test_table", rowType); + ColumnHandleMap assignments; + auto hiveConfig = + std::make_shared(std::make_shared( + std::unordered_map{})); + auto connectorQueryCtx = std::make_shared( + pool_.get(), + pool_.get(), + hiveConfig->config().get(), + nullptr, + nullptr, + nullptr, + nullptr, + "query.HiveConnectorTest", + "task.HiveConnectorTest", + "planNodeId.HiveConnectorTest", + 0); + + TestingHiveDataSource dataSource( + ROW({}, {}), + tableHandle, + assignments, + nullptr, + core::QueryConfig(std::unordered_map{}), + nullptr, + connectorQueryCtx, + hiveConfig); + dataSource.setSplit(HiveConnectorSplitBuilder("unused") + .connectorId(kHiveConnectorId) + .build()); + dataSource.setSplitReader(std::make_unique( + rowType, pool_.get(), std::vector{2, 3})); + + ContinueFuture future; + auto first = dataSource.next(1, future); + ASSERT_TRUE(first.has_value()); + ASSERT_NE(first.value(), nullptr); + EXPECT_EQ(first.value()->childrenSize(), 0); + EXPECT_EQ(first.value()->size(), 2); + + auto second = dataSource.next(1, future); + ASSERT_TRUE(second.has_value()); + ASSERT_NE(second.value(), nullptr); + EXPECT_EQ(second.value()->childrenSize(), 0); + EXPECT_EQ(second.value()->size(), 3); + EXPECT_EQ(first.value()->size(), 2); +} + TEST_F(HiveConnectorTest, makeScanSpec_requiredSubfields_multilevel) { auto columnType = ROW( {{"c0c0", BIGINT()}, diff --git a/bolt/dwio/common/SelectiveStructColumnReader.cpp b/bolt/dwio/common/SelectiveStructColumnReader.cpp index e5a57407f..957304677 100644 --- a/bolt/dwio/common/SelectiveStructColumnReader.cpp +++ b/bolt/dwio/common/SelectiveStructColumnReader.cpp @@ -34,6 +34,19 @@ namespace bytedance::bolt::dwio::common { namespace { +void prepareRowIndexField( + vector_size_t size, + VectorPtr& field, + const TypePtr& type, + memory::MemoryPool* pool) { + BOLT_CHECK_EQ(type->kind(), TypeKind::BIGINT); + if (field) { + BaseVector::prepareForReuse(field, size); + } else { + field = BaseVector::create(type, size, pool); + } +} + bool testFilterOnConstant(const bolt::common::ScanSpec& spec) { if (spec.isConstant() && !spec.constantValue()->isNullAt(0)) { // Non-null constant is known value during split scheduling and filters on @@ -105,7 +118,11 @@ void SelectiveStructColumnReaderBase::next( if (childSpec->isRowIndex()) { auto channel = childSpec->channel(); auto& childResult = resultRowVector->childAt(channel); - BaseVector::prepareForReuse(childResult, numValues); + prepareRowIndexField( + numValues, + childResult, + resultRowVector->type()->asRow().childAt(channel), + resultRowVector->pool()); auto* rowIdxs = childResult->as>()->mutableRawValues(); int64_t startIdx = @@ -410,7 +427,11 @@ void SelectiveStructColumnReaderBase::getValues( // it if (childSpec->isRowIndex() && !fileType_->containsChild(childSpec->fieldName())) { - BaseVector::prepareForReuse(childResult, rows.size()); + prepareRowIndexField( + rows.size(), + childResult, + rowType.childAt(channel), + resultRow->pool()); auto* flatVector = childResult->as>()->mutableRawValues(); int64_t rowIndexBase = childSpec->getRowIndexBase() + rowGroupOffset_ + diff --git a/bolt/exec/tests/TableScanTest.cpp b/bolt/exec/tests/TableScanTest.cpp index eeb974ba9..bb9ea8d04 100644 --- a/bolt/exec/tests/TableScanTest.cpp +++ b/bolt/exec/tests/TableScanTest.cpp @@ -1573,6 +1573,63 @@ TEST_F(TableScanTest, batchSize) { } } +TEST_F(TableScanTest, readsAfterBackingOutputChildrenReset) { + const vector_size_t size = 10; + auto c0 = makeFlatVector(size, folly::identity); + auto c1 = makeFlatVector( + size, [](auto row) { return fmt::format("value-{}", row); }); + auto c2 = makeRowVector({ + makeFlatVector(size, [](auto row) { return row * 3; }), + }); + auto offsets = [&]() { + std::vector values(size); + std::iota(values.begin(), values.end(), 0); + return values; + }; + auto keys = makeFlatVector(size, folly::identity); + auto ints = makeFlatVector(size, [](auto row) { return row * 11; }); + auto nestedRows = makeRowVector({ + makeFlatVector(size, [](auto row) { return row * 7; }), + }); + auto arrayInts = makeArrayVector(offsets(), ints); + auto mapInts = makeMapVector(offsets(), keys, ints); + auto arrayRows = makeArrayVector(offsets(), nestedRows); + auto mapRows = makeMapVector(offsets(), keys, nestedRows); + auto arrayArrayInts = makeArrayVector(offsets(), arrayInts); + auto mapArrayInts = makeMapVector(offsets(), keys, arrayInts); + auto arrayMapRows = makeArrayVector(offsets(), mapRows); + auto mapArrayRows = makeMapVector(offsets(), keys, arrayRows); + auto mapArrayMapRows = makeMapVector(offsets(), keys, arrayMapRows); + auto arrayMapArrayRows = makeArrayVector(offsets(), mapArrayRows); + auto data = makeRowVector(std::vector{ + c0, + c1, + c2, + arrayInts, + mapInts, + arrayRows, + mapRows, + arrayArrayInts, + mapArrayInts, + arrayMapRows, + mapArrayRows, + mapArrayMapRows, + arrayMapArrayRows, + }); + + auto filePath = TempFilePath::create(); + writeToFile(filePath->path, data); + + auto plan = tableScanNode(asRowType(data->type())); + auto task = AssertQueryBuilder(plan) + .splits(makeHiveConnectorSplits({filePath})) + .config(QueryConfig::kMaxOutputBatchRows, "3") + .assertResults(data); + + const auto opStats = task->taskStats().pipelineStats[0].operatorStats[0]; + EXPECT_GT(opStats.outputVectors, 1); +} + // Test that adding the same split with the same sequence id does not cause // double read and the 2nd split is ignored. TEST_F(TableScanTest, sequentialSplitNoDoubleRead) {