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

macvincent · facebook-github-bot · commit 2aa9df83fa50 · 2026-01-06T09:32:33.000-08:00
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
diff --git a/dwio/nimble/velox/StreamChunker.cpp b/dwio/nimble/velox/StreamChunker.cpp
@@ -21,6 +21,7 @@ template <typename T>
 std::unique_ptr<StreamChunker> getStreamChunkerTyped(
     StreamData& streamData,
     const StreamChunkerOptions& options) {
+  streamData.materializeStrings();
   if (auto* contentStreamChunker =
           dynamic_cast<ContentStreamData<T>*>(&streamData)) {
     return std::make_unique<ContentStreamChunker<T>>(
diff --git a/dwio/nimble/velox/StreamChunker.h b/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
 
 /**
@@ -94,6 +116,7 @@ class ContentStreamChunker final : public StreamChunker {
         minChunkSize_{options.minChunkSize},
         maxChunkSize_{options.maxChunkSize},
         ensureFullChunks_{options.ensureFullChunks},
+        useStreamStringBuffer_{!streamData_->mutableStringBuffer().empty()},
         extraMemory_{streamData_->extraMemory()} {
     static_assert(
         std::is_same_v<StreamDataT, ContentStreamData<T>> ||
@@ -194,27 +217,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());
+
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      if (useStreamStringBuffer_) {
+        detail::compactStringBuffer<T>(
+            mutableData.memoryPool(),
+            streamData_->mutableStringBuffer(),
+            tempData,
+            dataElementOffset_,
+            extraMemory_);
+      }
+    }
+
+    if (!useStreamStringBuffer_) {
+      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_;
 
@@ -230,6 +262,7 @@ class ContentStreamChunker final : public StreamChunker {
   const uint64_t minChunkSize_;
   const uint64_t maxChunkSize_;
   const bool ensureFullChunks_;
+  const bool useStreamStringBuffer_;
 
   size_t dataElementOffset_{0};
   size_t extraMemory_{0};
@@ -314,16 +347,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(),
@@ -361,6 +388,7 @@ class NullableContentStreamChunker final : public StreamChunker {
             options.minChunkSize,
             options.isFirstChunk)},
         ensureFullChunks_{options.ensureFullChunks},
+        useStreamStringBuffer_{!streamData_->mutableStringBuffer().empty()},
         extraMemory_{streamData_->extraMemory()} {
     static_assert(sizeof(bool) == 1);
     NIMBLE_CHECK(
@@ -461,39 +489,40 @@ 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");
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      if (useStreamStringBuffer_) {
+        detail::compactStringBuffer<T>(
+            mutableData.memoryPool(),
+            streamData_->mutableStringBuffer(),
+            tempData,
+            dataElementOffset_,
+            extraMemory_);
+      }
+    }
+    if (!useStreamStringBuffer_) {
+      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(),
@@ -515,6 +544,7 @@ class NullableContentStreamChunker final : public StreamChunker {
   const uint64_t maxChunkSize_;
   const bool omitStream_;
   const bool ensureFullChunks_;
+  const bool useStreamStringBuffer_;
 
   size_t dataElementOffset_{0};
   size_t nonNullsOffset_{0};
diff --git a/dwio/nimble/velox/StreamData.h b/dwio/nimble/velox/StreamData.h
@@ -54,6 +54,8 @@ class StreamData {
 
   virtual void materialize() {}
 
+  virtual void materializeStrings() {}
+
   const StreamDescriptorBuilder& descriptor() const {
     return descriptor_;
   }
@@ -84,6 +86,53 @@ 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();
+  }
+
+  uint64_t memoryUsed() const {
+    return buffer_.size() + lengths_.size() * sizeof(size_t);
+  }
+
+  void clear() {
+    buffer_.clear();
+    lengths_.clear();
+  }
+
+  void materialize(Vector<std::string_view>& data) {
+    if (empty()) {
+      return;
+    }
+    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 +205,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 +231,18 @@ class ContentStreamData final : public MutableStreamData {
   }
 
   inline uint64_t memoryUsed() const override {
-    return (data_.size() * sizeof(T)) + extraMemory_;
+    return (data_.size() * sizeof(T)) +
+        (stringBuffer_.empty() ? extraMemory_ : stringBuffer_.memoryUsed());
   }
 
   inline Vector<T>& mutableData() {
     return data_;
   }
 
+  inline StringBuffer& mutableStringBuffer() {
+    return stringBuffer_;
+  }
+
   void ensureMutableDataCapacity(uint64_t newSize) {
     this->ensureDataCapacity(data_, newSize);
   }
@@ -197,11 +253,19 @@ class ContentStreamData final : public MutableStreamData {
 
   inline virtual void reset() override {
     data_.clear();
+    stringBuffer_.clear();
     extraMemory_ = 0;
   }
 
+  void materializeStrings() 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 +350,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 +363,18 @@ class NullableContentStreamData final : public NullsStreamData {
   }
 
   inline uint64_t memoryUsed() const override {
-    return (data_.size() * sizeof(T)) + extraMemory_ +
-        NullsStreamData::memoryUsed();
+    return (data_.size() * sizeof(T)) + NullsStreamData::memoryUsed() +
+        (stringBuffer_.empty() ? extraMemory_ : stringBuffer_.memoryUsed());
   }
 
   inline Vector<T>& mutableData() {
     return data_;
   }
 
+  inline StringBuffer& mutableStringBuffer() {
+    return stringBuffer_;
+  }
+
   void ensureMutableDataCapacity(uint64_t newSize) {
     this->ensureDataCapacity(data_, newSize);
   }
@@ -317,11 +386,19 @@ class NullableContentStreamData final : public NullsStreamData {
   inline void reset() override {
     NullsStreamData::reset();
     data_.clear();
+    stringBuffer_.clear();
     extraMemory_ = 0;
   }
 
+  void materializeStrings() override {
+    if constexpr (std::is_same_v<T, std::string_view>) {
+      stringBuffer_.materialize(data_);
+    }
+  }
+
  private:
   Vector<T> data_;
+  StringBuffer stringBuffer_;
   uint64_t extraMemory_;
 };
 
diff --git a/dwio/nimble/velox/tests/StreamChunkerTests.cpp b/dwio/nimble/velox/tests/StreamChunkerTests.cpp
