fix: avoid overlapping scan output buffers

bolt/connectors/hive/HiveDataSource.cpp

@@ -639,9 +639,7 @@ std::optional<RowVectorPtr> 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

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<int64_t>(1, pool_);
+    values->asMutable<int64_t>()[0] = calls_;
+    values->setSize(sizeof(int64_t));
+    auto child = std::make_shared<FlatVector<int64_t>>(
+        pool_, BIGINT(), nullptr, 1, values, std::vector<BufferPtr>{});
+
+    if (calls_ == 0) {
+      firstChild_ = child;
+    }
+
+    output = std::make_shared<RowVector>(
+        pool_, rowType_, nullptr, 1, std::vector<VectorPtr>{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<BaseVector> firstChild_;
+};
+
+class TestingHiveDataSource : public HiveDataSource {
+ public:
+  TestingHiveDataSource(
+      const RowTypePtr& outputType,
+      const std::shared_ptr<connector::ConnectorTableHandle>& tableHandle,
+      const std::unordered_map<
+          std::string,
+          std::shared_ptr<connector::ColumnHandle>>& columnHandles,
+      FileHandleFactory* fileHandleFactory,
+      const core::QueryConfig& queryConfig,
+      folly::Executor* executor,
+      const std::shared_ptr<ConnectorQueryCtx>& connectorQueryCtx,
+      const std::shared_ptr<HiveConfig>& hiveConfig)
+      : HiveDataSource(
+            outputType,
+            tableHandle,
+            columnHandles,
+            fileHandleFactory,
+            queryConfig,
+            executor,
+            connectorQueryCtx,
+            hiveConfig) {}
+
+  void setSplitReader(std::unique_ptr<HiveSplitReaderBase> splitReader) {
+    splitReader_ = std::move(splitReader);
+  }
+
+  void setSplit(std::shared_ptr<ConnectorSplit> 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<HiveConfig>(std::make_shared<config::ConfigBase>(
+          std::unordered_map<std::string, std::string>{}));
+  auto connectorQueryCtx = std::make_shared<ConnectorQueryCtx>(
+      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<std::string, std::string>{}),
+      nullptr,
+      connectorQueryCtx,
+      hiveConfig);
+  dataSource.setSplit(HiveConnectorSplitBuilder("unused")
+                          .connectorId(kHiveConnectorId)
+                          .build());
+  dataSource.setSplitReader(
+      std::make_unique<TrackingSplitReader>(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()},

bolt/dwio/common/SelectiveColumnReaderInternal.h

@@ -49,11 +49,13 @@ bolt::common::AlwaysTrue& alwaysTrue();
 
 template <typename T>
 void SelectiveColumnReader::ensureValuesCapacity(vector_size_t numRows) {
-  if (values_ && values_->unique() &&
-      values_->capacity() >=
-          BaseVector::byteSize<T>(numRows) + simd::kPadding) {
+  const auto requiredBytes = BaseVector::byteSize<T>(numRows) + simd::kPadding;
+  if (values_ && values_->unique() && values_->capacity() >= requiredBytes &&
+      values_->capacity() <= requiredBytes + requiredBytes / 2) {
     return;
   }
+  values_.reset();
+  rawValues_ = nullptr;
   values_ = AlignedBuffer::allocate<T>(
       numRows + (simd::kPadding / sizeof(T)), &memoryPool_);
   rawValues_ = values_->asMutable<char>();

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<uint64_t> 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<FormatData> toFormatData(
+      const std::shared_ptr<const TypeWithId>& /*type*/,
+      const bolt::common::ScanSpec& /*scanSpec*/) override {
+    return std::make_unique<TestFormatData>();
+  }
+};
+
+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 <typename T>
+  void ensureCapacity(vector_size_t numRows) {
+    ensureValuesCapacity<T>(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<int64_t>(100);
+  auto* values = reader.values();
+  ASSERT_NE(values, nullptr);
+
+  reader.ensureCapacity<int64_t>(100);
+  EXPECT_EQ(reader.values(), values);
+
+  reader.ensureCapacity<int64_t>(99);
+  EXPECT_EQ(reader.values(), values);
+  auto maxReusableBytes = BaseVector::byteSize<int64_t>(99) + simd::kPadding;
+  maxReusableBytes += maxReusableBytes / 2;
+  EXPECT_LE(reader.values()->capacity(), maxReusableBytes);
+
+  const auto bytesWithOversizedValues = pool_->currentBytes();
+  reader.ensureCapacity<int64_t>(60);
+  EXPECT_LT(pool_->currentBytes(), bytesWithOversizedValues);
+  maxReusableBytes = BaseVector::byteSize<int64_t>(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<int64_t>(60);
+  EXPECT_TRUE(sharedValues->unique());
+  EXPECT_NE(reader.values(), values);
+}
+
 TEST_F(ReaderTest, projectColumnsFilterArray) {
   constexpr int kSize = 10;
   auto input = makeRowVector({