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({