Support Per Stream String Buffers in Stream Data (facebookincubator#391)

Summary:

Adds per-stream string buffer support to Nimble's StreamData infrastructure. Previously, string data buffers were managed externally, but now each MutableStreamData instance can own its own Buffer for string content, enabling better memory isolation and lifecycle management with chunking.

Differential Revision: D89933055

Author: macvincent

dwio/nimble/velox/StreamChunker.h

@@ -194,6 +194,7 @@ class ContentStreamChunker final : public StreamChunker {
 
     // Move and clear existing buffer
     auto tempData = std::move(currentData);
+    auto tempBuffer = streamData_->takeStringBuffer();
     streamData_->reset();
     NIMBLE_DCHECK(
         streamData_->empty(), "StreamData should be empty after reset");
@@ -215,6 +216,11 @@ class ContentStreamChunker final : public StreamChunker {
         tempData.begin() + dataElementOffset_,
         remainingDataCount,
         mutableData.begin());
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      if (tempBuffer) {
+        streamData_->compactStringBuffer(mutableData);
+      }
+    }
     dataElementOffset_ = 0;
     streamData_->extraMemory() = extraMemory_;
 
@@ -461,6 +467,7 @@ class NullableContentStreamChunker final : public StreamChunker {
     // Move and clear existing buffers
     auto tempNonNulls = std::move(currentNonNulls);
     auto tempData = std::move(currentData);
+    auto tempBuffer = streamData_->takeStringBuffer();
     const bool hasNulls = streamData_->hasNulls();
     streamData_->reset();
     NIMBLE_CHECK(
@@ -505,6 +512,11 @@ class NullableContentStreamChunker final : public StreamChunker {
           remainingNonNullsCount,
           mutableNonNulls.begin());
     }
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      if (tempBuffer) {
+        streamData_->compactStringBuffer(mutableData);
+      }
+    }
     streamData_->extraMemory() = extraMemory_;
     dataElementOffset_ = 0;
     nonNullsOffset_ = 0;

dwio/nimble/velox/StreamData.h

@@ -19,6 +19,7 @@
 #include <span>
 #include <string_view>
 
+#include "dwio/nimble/common/Buffer.h"
 #include "dwio/nimble/common/Vector.h"
 #include "dwio/nimble/velox/BufferGrowthPolicy.h"
 #include "dwio/nimble/velox/SchemaBuilder.h"
@@ -65,9 +66,12 @@ class StreamData {
 class MutableStreamData : public StreamData {
  protected:
   MutableStreamData(
+      velox::memory::MemoryPool& memoryPool,
       const StreamDescriptorBuilder& descriptor,
       const InputBufferGrowthPolicy& growthPolicy)
-      : StreamData(descriptor), growthPolicy_{growthPolicy} {}
+      : StreamData(descriptor),
+        growthPolicy_{growthPolicy},
+        memoryPool_{memoryPool} {}
 
   virtual ~MutableStreamData() = default;
 
@@ -80,8 +84,44 @@ class MutableStreamData : public StreamData {
     }
   }
 
+  void resetStringBuffer() {
+    if (stringBuffer_) {
+      stringBuffer_.reset();
+    }
+  }
+
+ public:
+  /// Returns a reference to the internal string buffer.
+  /// Creates the string buffer lazily if it doesn't exist.
+  Buffer& stringBuffer() {
+    if (!stringBuffer_) {
+      stringBuffer_ = std::make_unique<Buffer>(memoryPool_);
+    }
+    return *stringBuffer_;
+  }
+
+  /// Takes ownership of the buffer from this stream.
+  std::unique_ptr<Buffer> takeStringBuffer() {
+    return std::move(stringBuffer_);
+  }
+
+  /// Compacts the string buffer by copying remaining strings to a new buffer
+  /// and updating the string_views in data.
+  void compactStringBuffer(Vector<std::string_view>& data) {
+    if (data.empty()) {
+      return;
+    }
+    auto newStringBuffer = std::make_unique<Buffer>(memoryPool_);
+    for (size_t i = 0; i < data.size(); ++i) {
+      data[i] = newStringBuffer->writeString(data[i]);
+    }
+    stringBuffer_ = std::move(newStringBuffer);
+  }
+
  private:
   const InputBufferGrowthPolicy& growthPolicy_;
+  velox::memory::MemoryPool& memoryPool_;
+  std::unique_ptr<Buffer> stringBuffer_;
 };
 
 /// Provides a lightweight, non-owning view into a portion of stream data,
@@ -156,7 +196,7 @@ class ContentStreamData final : public MutableStreamData {
       velox::memory::MemoryPool& pool,
       const StreamDescriptorBuilder& descriptor,
       const InputBufferGrowthPolicy& growthPolicy)
-      : MutableStreamData(descriptor, growthPolicy), data_{&pool} {}
+      : MutableStreamData(pool, descriptor, growthPolicy), data_{&pool} {}
 
   inline uint32_t rowCount() const override {
     return data_.size();
@@ -198,6 +238,7 @@ class ContentStreamData final : public MutableStreamData {
   inline virtual void reset() override {
     data_.clear();
     extraMemory_ = 0;
+    resetStringBuffer();
   }
 
  private:
@@ -218,7 +259,7 @@ class NullsStreamData : public MutableStreamData {
       velox::memory::MemoryPool& pool,
       const StreamDescriptorBuilder& descriptor,
       const InputBufferGrowthPolicy& growthPolicy)
-      : MutableStreamData(descriptor, growthPolicy), nonNulls_{&pool} {}
+      : MutableStreamData(pool, descriptor, growthPolicy), nonNulls_{&pool} {}
 
   inline uint32_t rowCount() const override {
     return bufferedCount_;
@@ -252,6 +293,7 @@ class NullsStreamData : public MutableStreamData {
     nonNulls_.clear();
     hasNulls_ = false;
     bufferedCount_ = 0;
+    resetStringBuffer();
   }
 
   void materialize() override {

dwio/nimble/velox/tests/StreamChunkerTests.cpp

@@ -35,6 +35,18 @@ void populateData(Vector<T>& vec, const std::vector<T>& data) {
   }
 }
 
+uint64_t populateStringData(
+    Buffer& stringBuffer,
+    Vector<std::string_view>& data,
+    const std::vector<std::string_view>& testData) {
+  uint64_t extraMemory = 0;
+  for (const auto& str : testData) {
+    extraMemory += str.size();
+    data.emplace_back(stringBuffer.writeString(str));
+  }
+  return extraMemory;
+}
+
 template <typename T>
 std::vector<T> toVector(std::string_view data) {
   const T* chunkData = reinterpret_cast<const T*>(data.data());
@@ -508,13 +520,8 @@ TEST_F(StreamChunkerTestsBase, contentStreamStringChunking) {
   {
     std::vector<std::string_view> testData = {
         "short", "a_longer_string", "x", "medium_size", "tiny"};
-    auto& data = stream.mutableData();
-    populateData(data, testData);
-
-    // Calculate extra memory for string content
-    for (const auto& str : testData) {
-      stream.extraMemory() += str.size();
-    }
+    stream.extraMemory() += populateStringData(
+        stream.stringBuffer(), stream.mutableData(), testData);
 
     // Total string content sizes:
     // "short"(5) + "a_longer_string"(15) + "x"(1) + "medium_size"(11) +
@@ -582,16 +589,11 @@ TEST_F(StreamChunkerTestsBase, contentStreamStringChunking) {
 
   // Test 4: We can reuse a stream post compaction.
   {
-    auto& data = stream.mutableData();
     stream.extraMemory() = 0; // Reset extra memory
     std::vector<std::string_view> testData = {
         "hello", "world", "hello", "world"};
-    populateData(data, testData);
-
-    // Calculate extra memory for string content
-    for (const auto& str : testData) {
-      stream.extraMemory() += str.size();
-    }
+    stream.extraMemory() += populateStringData(
+        stream.stringBuffer(), stream.mutableData(), testData);
 
     auto chunker = getStreamChunker(
         stream,
@@ -619,17 +621,12 @@ TEST_F(StreamChunkerTestsBase, contentStreamStringChunking) {
   {
     std::vector<std::string_view> testData = {
         "a", "short", "really really large string", "small"};
-    auto& data = stream.mutableData();
-    populateData(data, testData);
 
     // Size of "really really large string" and string view minus 1.
     maxChunkSize = 24 + sizeof(std::string_view);
 
-    // Calculate extra memory for string content.
-    for (const auto& str : testData) {
-      stream.extraMemory() += str.size();
-    }
-
+    stream.extraMemory() += populateStringData(
+        stream.stringBuffer(), stream.mutableData(), testData);
     // "a"(1), "short"(5), "really really large string"(25), "small"(5) = 37
     // bytes of extra memory.
     ASSERT_EQ(stream.extraMemory(), 37);
@@ -888,9 +885,7 @@ TEST_F(StreamChunkerTestsBase, nullableContentStreamStringChunking) {
   // insertion
   stream.ensureAdditionalNullsCapacity(
       /*mayHaveNulls=*/true, static_cast<uint32_t>(nonNullsData.size()));
-  auto& data = stream.mutableData();
   auto& nonNulls = stream.mutableNonNulls();
-  populateData(data, testData);
   populateData(nonNulls, nonNullsData);
 
   // Set extra memory for string overhead
@@ -900,9 +895,9 @@ TEST_F(StreamChunkerTestsBase, nullableContentStreamStringChunking) {
   // size of nulls = 11 * 1 = 11 bytes
   // total size of stored string data = 36 bytes
   // total size of stream data = 56 + 56 + 11 + 36 = 159 bytes
-  for (const auto& entry : testData) {
-    stream.extraMemory() += entry.size();
-  }
+  stream.extraMemory() +=
+      populateStringData(stream.stringBuffer(), stream.mutableData(), testData);
+
   ASSERT_EQ(stream.memoryUsed(), 159);
 
   // Test 1: Not last chunk
@@ -982,16 +977,13 @@ TEST_F(StreamChunkerTestsBase, nullableContentStreamStringChunking) {
     stream.ensureAdditionalNullsCapacity(
         /*mayHaveNulls=*/false,
         static_cast<uint32_t>(smallNonNullsData.size()));
-    auto& smallData = stream.mutableData();
     auto& smallNonNulls = stream.mutableNonNulls();
-    populateData(smallData, smallTestData);
     populateData(smallNonNulls, smallNonNullsData);
 
     // Reset extra memory for new test data
     stream.extraMemory() = 0;
-    for (const auto& entry : smallTestData) {
-      stream.extraMemory() += entry.size();
-    }
+    stream.extraMemory() += populateStringData(
+        stream.stringBuffer(), stream.mutableData(), smallTestData);
 
     const uint64_t minChunkSize =
         smallTestData.at(0).size() + sizeof(std::string_view);
@@ -1025,7 +1017,6 @@ TEST_F(StreamChunkerTestsBase, nullableContentStreamStringChunking) {
   // Test 5: Single string exceeds maxChunkSize.
   {
     testData = {"a", "short", "really really large string", "small"};
-    populateData(data, testData);
 
     nonNullsData = {true, false, true, true, false, true, false};
     stream.ensureAdditionalNullsCapacity(
@@ -1037,9 +1028,8 @@ TEST_F(StreamChunkerTestsBase, nullableContentStreamStringChunking) {
     const uint64_t maxChunkSize = 24 + sizeof(std::string_view) + sizeof(bool);
 
     // Calculate extra memory for string content.
-    for (const auto& str : testData) {
-      stream.extraMemory() += str.size();
-    }
+    stream.extraMemory() += populateStringData(
+        stream.stringBuffer(), stream.mutableData(), testData);
 
     // "a"(1), "short"(5), "really really large string"(25), "small"(5) =
     // 37 bytes of extra memory.
@@ -1113,11 +1103,8 @@ TEST_F(StreamChunkerTestsBase, contentStreamStringChunkerRowCount) {
   // Add strings with varying sizes.
   std::vector<std::string_view> testData = {
       "a", "bb", "ccc", "dddd", "eeeee", "ffffff"};
-  auto& data = stream.mutableData();
-  populateData(data, testData);
-  for (const auto& str : testData) {
-    stream.extraMemory() += str.size();
-  }
+  stream.extraMemory() +=
+      populateStringData(stream.stringBuffer(), stream.mutableData(), testData);
 
   // Max chunk size that fits ~2-3 strings per chunk.
   auto chunker = getStreamChunker(
@@ -1228,13 +1215,10 @@ TEST_F(StreamChunkerTestsBase, nullableContentStreamStringChunkerRowCount) {
   std::vector<bool> nonNullsData = {true, false, true, true, false, true};
   stream.ensureAdditionalNullsCapacity(
       /*mayHaveNulls=*/true, static_cast<uint32_t>(nonNullsData.size()));
-  auto& data = stream.mutableData();
   auto& nonNulls = stream.mutableNonNulls();
-  populateData(data, testData);
   populateData(nonNulls, nonNullsData);
-  for (const auto& str : testData) {
-    stream.extraMemory() += str.size();
-  }
+  stream.extraMemory() +=
+      populateStringData(stream.stringBuffer(), stream.mutableData(), testData);
 
   auto chunker = getStreamChunker(
       stream,