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

@@ -41,6 +41,28 @@ inline bool shouldOmitNullStream(
   }
   return isFirstChunk || streamData.empty();
 }
+
+template <typename T>
+inline void compactStringBuffer(
+    velox::memory::MemoryPool* pool,
+    StringBuffer& stringBuffer,
+    Vector<T>& mutableData,
+    size_t dataElementOffset,
+    uint64_t extraMemory) {
+  StringBuffer newStringBuffer(*pool);
+  auto& newBuffer = newStringBuffer.mutableBuffer();
+  newBuffer.reserve(extraMemory);
+
+  auto& newLengths = newStringBuffer.mutableLengths();
+  newLengths.reserve(mutableData.size() - dataElementOffset);
+
+  for (auto i = dataElementOffset; i < mutableData.size(); ++i) {
+    auto str = mutableData[i];
+    newBuffer.insert(newBuffer.end(), str.begin(), str.end());
+    newLengths.push_back(str.size());
+  }
+  stringBuffer = std::move(newStringBuffer);
+}
 } // namespace detail
 
 /**
@@ -106,6 +128,9 @@ class ContentStreamChunker final : public StreamChunker {
         maxChunkSize_,
         sizeof(T),
         "MaxChunkSize must be at least the size of a single data element.");
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      streamData_->materializeStringBuffer();
+    }
   }
 
   std::optional<StreamDataView> next() override {
@@ -194,27 +219,36 @@ class ContentStreamChunker final : public StreamChunker {
 
     // Move and clear existing buffer
     auto tempData = std::move(currentData);
+    auto tempBuffer = std::move(streamData_->mutableStringBuffer());
     streamData_->reset();
     NIMBLE_DCHECK(
         streamData_->empty(), "StreamData should be empty after reset");
 
     auto& mutableData = streamData_->mutableData();
-    mutableData.reserve(remainingDataCount);
-    NIMBLE_DCHECK_GE(
-        mutableData.capacity(),
-        remainingDataCount,
-        "Data buffer capacity should be at least new capacity");
-
-    mutableData.resize(remainingDataCount);
-    NIMBLE_DCHECK_EQ(
-        mutableData.size(),
-        remainingDataCount,
-        "Data buffer size should be equal to remaining data count");
-
-    std::copy_n(
-        tempData.begin() + dataElementOffset_,
-        remainingDataCount,
-        mutableData.begin());
+
+    bool compactedStringBuffer{false};
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      if (!tempBuffer.empty()) {
+        detail::compactStringBuffer<T>(
+            mutableData.memoryPool(),
+            streamData_->mutableStringBuffer(),
+            tempData,
+            dataElementOffset_,
+            extraMemory_);
+        compactedStringBuffer = true;
+      }
+    }
+    if (!compactedStringBuffer) {
+      mutableData.resize(remainingDataCount);
+      NIMBLE_DCHECK_EQ(
+          mutableData.size(),
+          remainingDataCount,
+          "Data buffer size should be equal to remaining data count");
+      std::copy_n(
+          tempData.begin() + dataElementOffset_,
+          remainingDataCount,
+          mutableData.begin());
+    }
     dataElementOffset_ = 0;
     streamData_->extraMemory() = extraMemory_;
 
@@ -314,16 +348,10 @@ class NullsStreamChunker final : public StreamChunker {
     NIMBLE_CHECK(
         streamData_->empty(), "StreamData should be empty after reset");
 
-    auto& mutableNonNulls = streamData_->mutableNonNulls();
-    mutableNonNulls.reserve(remainingNonNullsCount);
-    NIMBLE_DCHECK_GE(
-        mutableNonNulls.capacity(),
-        remainingNonNullsCount,
-        "NonNulls buffer capacity should be at least new capacity");
-
     streamData_->ensureAdditionalNullsCapacity(
         hasNulls, static_cast<uint32_t>(remainingNonNullsCount));
     if (hasNulls) {
+      auto& mutableNonNulls = streamData_->mutableNonNulls();
       mutableNonNulls.resize(remainingNonNullsCount);
       NIMBLE_CHECK_EQ(
           mutableNonNulls.size(),
@@ -372,6 +400,9 @@ class NullableContentStreamChunker final : public StreamChunker {
         "MaxChunkSize must be at least the size of a single element.");
 
     streamData.materialize();
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      streamData_->materializeStringBuffer();
+    }
   }
 
   std::optional<StreamDataView> next() override {
@@ -461,39 +492,41 @@ class NullableContentStreamChunker final : public StreamChunker {
     // Move and clear existing buffers
     auto tempNonNulls = std::move(currentNonNulls);
     auto tempData = std::move(currentData);
+    auto tempBuffer = std::move(streamData_->mutableStringBuffer());
     const bool hasNulls = streamData_->hasNulls();
     streamData_->reset();
     NIMBLE_CHECK(
         streamData_->empty(), "StreamData should be empty after reset");
 
     auto& mutableData = streamData_->mutableData();
-    mutableData.reserve(remainingDataCount);
-    NIMBLE_DCHECK_GE(
-        mutableData.capacity(),
-        remainingDataCount,
-        "Data buffer capacity should be at least new capacity");
-
-    mutableData.resize(remainingDataCount);
-    NIMBLE_CHECK_EQ(
-        mutableData.size(),
-        remainingDataCount,
-        "Data buffer size should be equal to remaining data count");
-
-    std::copy_n(
-        tempData.begin() + dataElementOffset_,
-        remainingDataCount,
-        mutableData.begin());
-
-    auto& mutableNonNulls = streamData_->mutableNonNulls();
-    mutableNonNulls.reserve(remainingNonNullsCount);
-    NIMBLE_DCHECK_GE(
-        mutableNonNulls.capacity(),
-        remainingNonNullsCount,
-        "NonNulls buffer capacity should be at least new capacity");
+    bool compactedStringBuffer{false};
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      if (!tempBuffer.empty()) {
+        detail::compactStringBuffer<T>(
+            mutableData.memoryPool(),
+            streamData_->mutableStringBuffer(),
+            tempData,
+            dataElementOffset_,
+            extraMemory_);
+        compactedStringBuffer = true;
+      }
+    }
+    if (!compactedStringBuffer) {
+      mutableData.resize(remainingDataCount);
+      NIMBLE_DCHECK_EQ(
+          mutableData.size(),
+          remainingDataCount,
+          "Data buffer size should be equal to remaining data count");
+      std::copy_n(
+          tempData.begin() + dataElementOffset_,
+          remainingDataCount,
+          mutableData.begin());
+    }
 
     streamData_->ensureAdditionalNullsCapacity(
         hasNulls, static_cast<uint32_t>(remainingNonNullsCount));
     if (hasNulls) {
+      auto& mutableNonNulls = streamData_->mutableNonNulls();
       mutableNonNulls.resize(remainingNonNullsCount);
       NIMBLE_CHECK_EQ(
           mutableNonNulls.size(),

dwio/nimble/velox/StreamData.h

@@ -54,6 +54,8 @@ class StreamData {
 
   virtual void materialize() {}
 
+  virtual void materializeStringBuffer() {}
+
   const StreamDescriptorBuilder& descriptor() const {
     return descriptor_;
   }
@@ -84,6 +86,50 @@ class MutableStreamData : public StreamData {
   const InputBufferGrowthPolicy& growthPolicy_;
 };
 
+class StringBuffer {
+ public:
+  explicit StringBuffer(velox::memory::MemoryPool& pool)
+      : buffer{&pool}, lengths{&pool} {}
+
+  Vector<char>& mutableBuffer() {
+    return buffer;
+  }
+
+  Vector<size_t>& mutableLengths() {
+    return lengths;
+  }
+
+  bool empty() const {
+    return lengths.empty();
+  }
+
+  uint64_t size() const {
+    return lengths.size();
+  }
+
+  void clear() {
+    buffer.clear();
+    lengths.clear();
+  }
+
+  void materialize(Vector<std::string_view>& data) {
+    if (empty()) {
+      return;
+    }
+    data.clear();
+    data.resize(lengths.size());
+    size_t runningOffset = 0;
+    for (size_t i = 0; i < lengths.size(); ++i) {
+      data[i] = std::string_view(buffer.data() + runningOffset, lengths[i]);
+      runningOffset += lengths[i];
+    }
+  }
+
+ private:
+  Vector<char> buffer;
+  Vector<size_t> lengths;
+};
+
 /// Provides a lightweight, non-owning view into a portion of stream data,
 /// containing references to the data content and null indicators for efficient
 /// processing of large streams without copying data.
@@ -156,7 +202,9 @@ class ContentStreamData final : public MutableStreamData {
       velox::memory::MemoryPool& pool,
       const StreamDescriptorBuilder& descriptor,
       const InputBufferGrowthPolicy& growthPolicy)
-      : MutableStreamData(descriptor, growthPolicy), data_{&pool} {}
+      : MutableStreamData(descriptor, growthPolicy),
+        data_{&pool},
+        stringBuffer_{pool} {}
 
   inline uint32_t rowCount() const override {
     return data_.size();
@@ -180,13 +228,18 @@ class ContentStreamData final : public MutableStreamData {
   }
 
   inline uint64_t memoryUsed() const override {
-    return (data_.size() * sizeof(T)) + extraMemory_;
+    return (data_.size() * sizeof(T)) + extraMemory_ +
+        (stringBuffer_.size() * sizeof(size_t));
   }
 
   inline Vector<T>& mutableData() {
     return data_;
   }
 
+  inline StringBuffer& mutableStringBuffer() {
+    return stringBuffer_;
+  }
+
   void ensureMutableDataCapacity(uint64_t newSize) {
     this->ensureDataCapacity(data_, newSize);
   }
@@ -197,11 +250,19 @@ class ContentStreamData final : public MutableStreamData {
 
   inline virtual void reset() override {
     data_.clear();
+    stringBuffer_.clear();
     extraMemory_ = 0;
   }
 
+  void materializeStringBuffer() override {
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      stringBuffer_.materialize(data_);
+    }
+  }
+
  private:
   Vector<T> data_;
+  StringBuffer stringBuffer_;
   uint64_t extraMemory_{0};
 };
 
@@ -286,6 +347,7 @@ class NullableContentStreamData final : public NullsStreamData {
       const InputBufferGrowthPolicy& growthPolicy)
       : NullsStreamData(memoryPool, descriptor, growthPolicy),
         data_{&memoryPool},
+        stringBuffer_{memoryPool},
         extraMemory_{0} {}
 
   inline std::string_view data() const override {
@@ -298,14 +360,19 @@ class NullableContentStreamData final : public NullsStreamData {
   }
 
   inline uint64_t memoryUsed() const override {
-    return (data_.size() * sizeof(T)) + extraMemory_ +
-        NullsStreamData::memoryUsed();
+    return (data_.size() * sizeof(T)) +
+        (stringBuffer_.size() * sizeof(size_t)) +
+        NullsStreamData::memoryUsed() + extraMemory_;
   }
 
   inline Vector<T>& mutableData() {
     return data_;
   }
 
+  inline StringBuffer& mutableStringBuffer() {
+    return stringBuffer_;
+  }
+
   void ensureMutableDataCapacity(uint64_t newSize) {
     this->ensureDataCapacity(data_, newSize);
   }
@@ -317,11 +384,19 @@ class NullableContentStreamData final : public NullsStreamData {
   inline void reset() override {
     NullsStreamData::reset();
     data_.clear();
+    stringBuffer_.clear();
     extraMemory_ = 0;
   }
 
+  void materializeStringBuffer() override {
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      stringBuffer_.materialize(data_);
+    }
+  }
+
  private:
   Vector<T> data_;
+  StringBuffer stringBuffer_;
   uint64_t extraMemory_;
 };