Support Per Stream String Buffers in Writer (facebookincubator#415)

Summary:

This change enables Nimble to handle string data with per-stream buffers rather than a shared buffer, improving memory locality and enabling independent chunking of string columns. It follows the existing chunker pattern while adding string-specific optimizations for buffer management. Integration with the Velox Writer happens in the next diff.

Reviewed By: helfman

Differential Revision: D90411755

Author: macvincent

dwio/nimble/velox/StreamChunker.cpp

@@ -37,6 +37,11 @@ std::unique_ptr<StreamChunker> getStreamChunkerTyped(
     }
     return std::make_unique<NullableContentStreamChunker<T>>(
         *nullableContentStreamData, options);
+  } else if (
+      auto* nullableContentStringStreamData =
+          dynamic_cast<NullableContentStringStreamData*>(&streamData)) {
+    return std::make_unique<NullableContentStringStreamChunker>(
+        *nullableContentStringStreamData, options);
   } else if (
       auto* nullsStreamData = dynamic_cast<NullsStreamData*>(&streamData)) {
     return std::make_unique<NullsStreamChunker>(*nullsStreamData, options);

dwio/nimble/velox/StreamChunker.h

@@ -21,14 +21,15 @@
 namespace facebook::nimble {
 namespace detail {
 inline bool shouldOmitDataStream(
-    const StreamData& streamData,
+    uint64_t dataVectorSize,
     uint64_t minChunkSize,
+    bool allNulls,
     bool isFirstChunk) {
-  if (streamData.data().size() > minChunkSize) {
+  if (dataVectorSize > minChunkSize) {
     return false;
   }
   // When all values are null, the values stream is omitted.
-  return isFirstChunk || streamData.nonNulls().empty();
+  return isFirstChunk || allNulls;
 }
 
 inline bool shouldOmitNullStream(
@@ -184,7 +185,7 @@ class ContentStreamChunker final : public StreamChunker {
   }
 
   void compact() override {
-    // No changes made to stream data, nothing to compact.
+    // No chunks consumed from stream data, we should not compact.
     if (dataElementOffset_ == 0) {
       return;
     }
@@ -300,7 +301,7 @@ class NullsStreamChunker final : public StreamChunker {
 
  private:
   void compact() override {
-    // No changes made to stream data, nothing to compact.
+    // No chunks consumed from stream data, we should not compact.
     if (nonNullsOffset_ == 0) {
       return;
     }
@@ -357,8 +358,9 @@ class NullableContentStreamChunker final : public StreamChunker {
         minChunkSize_{options.minChunkSize},
         maxChunkSize_{options.maxChunkSize},
         omitStream_{detail::shouldOmitDataStream(
-            streamData,
+            streamData.data().size(),
             options.minChunkSize,
+            streamData.nonNulls().empty(),
             options.isFirstChunk)},
         ensureFullChunks_{options.ensureFullChunks},
         extraMemory_{streamData_->extraMemory()} {
@@ -448,7 +450,7 @@ class NullableContentStreamChunker final : public StreamChunker {
   }
 
   void compact() override {
-    // No changes made to stream data, nothing to compact.
+    // No chunks consumed from stream data, we should not compact.
     if (nonNullsOffset_ == 0) {
       return;
     }
@@ -573,4 +575,201 @@ NullableContentStreamChunker<std::string_view>::nextChunkSize() {
   }
   return chunkSize;
 }
+
+class NullableContentStringStreamChunker final : public StreamChunker {
+ public:
+  explicit NullableContentStringStreamChunker(
+      NullableContentStringStreamData& streamData,
+      const StreamChunkerOptions& options)
+      : streamData_{&streamData},
+        minChunkSize_{options.minChunkSize},
+        maxChunkSize_{options.maxChunkSize},
+        omitStream_{detail::shouldOmitDataStream(
+            streamData.bufferSize(),
+            options.minChunkSize,
+            streamData.nonNulls().empty(),
+            options.isFirstChunk)},
+        ensureFullChunks_{options.ensureFullChunks} {
+    static_assert(sizeof(bool) == 1);
+  }
+
+  std::optional<StreamDataView> next() override {
+    if (omitStream_) {
+      return std::nullopt;
+    }
+    const auto& chunkSize = nextChunkSize();
+    if (chunkSize.rollingChunkSize == 0) {
+      return std::nullopt;
+    }
+
+    // Content
+    auto& output = streamData_->mutableStringViews();
+    output.resize(0);
+    output.reserve(chunkSize.dataElementCount);
+    auto mutableData = streamData_->mutableData();
+    auto& mutableLengths = mutableData.lengths;
+    auto& mutableBuffer = mutableData.buffer;
+    auto currentBufferOffset = bufferOffset_;
+    for (size_t i = 0; i < chunkSize.dataElementCount; ++i) {
+      const auto currentLength = mutableLengths[lengthsOffset_ + i];
+      output.emplace_back(
+          std::string_view{
+              mutableBuffer.data() + currentBufferOffset, currentLength});
+      currentBufferOffset += currentLength;
+    }
+    std::string_view dataChunk = {
+        reinterpret_cast<const char*>(output.data()),
+        chunkSize.dataElementCount * sizeof(std::string_view)};
+
+    // Nulls
+    std::span<const bool> nonNullsChunk(
+        streamData_->mutableNonNulls().data() + nonNullsOffset_,
+        chunkSize.nullElementCount);
+
+    lengthsOffset_ += chunkSize.dataElementCount;
+    nonNullsOffset_ += chunkSize.nullElementCount;
+    bufferOffset_ += chunkSize.extraMemory;
+
+    if (chunkSize.nullElementCount > chunkSize.dataElementCount) {
+      return StreamDataView{
+          streamData_->descriptor(),
+          dataChunk,
+          static_cast<uint32_t>(chunkSize.nullElementCount),
+          nonNullsChunk};
+    }
+    NIMBLE_DCHECK_EQ(chunkSize.dataElementCount, chunkSize.nullElementCount);
+    return StreamDataView{
+        streamData_->descriptor(),
+        dataChunk,
+        static_cast<uint32_t>(chunkSize.dataElementCount)};
+  }
+
+ private:
+  ChunkSize nextChunkSize() {
+    const auto& stringLengths = streamData_->mutableData().lengths;
+    const auto& nonNulls = streamData_->mutableNonNulls();
+    ChunkSize chunkSize;
+    bool fullChunk{false};
+    // Calculate how many entries we can fit in the chunk
+    for (size_t idx = nonNullsOffset_; idx < nonNulls.size(); ++idx) {
+      uint64_t currentTotalSize{sizeof(bool)};
+      uint32_t currentDataCount{0};
+      size_t currentExtraMemory{0};
+      if (nonNulls[idx]) {
+        currentExtraMemory =
+            stringLengths[lengthsOffset_ + chunkSize.dataElementCount];
+        currentTotalSize += currentExtraMemory + sizeof(uint64_t);
+        ++currentDataCount;
+      }
+
+      if (chunkSize.rollingChunkSize == 0 && currentTotalSize > maxChunkSize_) {
+        // Allow a single oversized string as its own chunk.
+        fullChunk = true;
+        chunkSize.extraMemory += currentExtraMemory;
+        chunkSize.dataElementCount += currentDataCount;
+        chunkSize.rollingChunkSize += currentTotalSize;
+        ++chunkSize.nullElementCount;
+        break;
+      }
+
+      if (chunkSize.rollingChunkSize + currentTotalSize > maxChunkSize_) {
+        fullChunk = true;
+        break;
+      }
+
+      chunkSize.extraMemory += currentExtraMemory;
+      chunkSize.dataElementCount += currentDataCount;
+      chunkSize.rollingChunkSize += currentTotalSize;
+      ++chunkSize.nullElementCount;
+    }
+
+    fullChunk = fullChunk || (chunkSize.rollingChunkSize == maxChunkSize_);
+    if ((ensureFullChunks_ && !fullChunk) ||
+        (chunkSize.rollingChunkSize < minChunkSize_)) {
+      chunkSize = ChunkSize{};
+    }
+    return chunkSize;
+  }
+
+  void compact() override {
+    // Clear existing outputvector before beginning compaction.
+    streamData_->mutableStringViews().clear();
+
+    // No chunks consumed from stream data, we should not compact.
+    if (nonNullsOffset_ == 0) {
+      return;
+    }
+
+    const bool hasNulls = streamData_->hasNulls();
+    // Move and clear existing buffers
+    auto tempBuffer = std::move(streamData_->mutableData().buffer);
+    auto tempLengths = std::move(streamData_->mutableData().lengths);
+    auto tempNonNulls = std::move(streamData_->mutableNonNulls());
+    streamData_->reset();
+    NIMBLE_DCHECK(
+        streamData_->empty(), "StreamData should be empty after reset");
+
+    {
+      const auto remainingDataCount = tempLengths.size() - lengthsOffset_;
+      auto& mutableDataLength = streamData_->mutableData().lengths;
+      mutableDataLength.resize(remainingDataCount);
+      NIMBLE_DCHECK_EQ(
+          mutableDataLength.size(),
+          remainingDataCount,
+          "Data length size should be equal to remaining data count");
+
+      std::copy_n(
+          tempLengths.begin() + lengthsOffset_,
+          remainingDataCount,
+          mutableDataLength.begin());
+    }
+
+    {
+      const auto remainingDataBytes = tempBuffer.size() - bufferOffset_;
+      auto& mutableDataBuffer = streamData_->mutableData().buffer;
+      mutableDataBuffer.resize(remainingDataBytes);
+      NIMBLE_DCHECK_EQ(
+          mutableDataBuffer.size(),
+          remainingDataBytes,
+          "Data buffer size should be equal to remaining data bytes");
+
+      std::copy_n(
+          tempBuffer.begin() + bufferOffset_,
+          remainingDataBytes,
+          mutableDataBuffer.begin());
+    }
+
+    {
+      auto& mutableNonNulls = streamData_->mutableNonNulls();
+      const auto remainingNonNullsCount = tempNonNulls.size() - nonNullsOffset_;
+      streamData_->ensureAdditionalNullsCapacity(
+          hasNulls, static_cast<uint32_t>(remainingNonNullsCount));
+      if (hasNulls) {
+        mutableNonNulls.resize(remainingNonNullsCount);
+        NIMBLE_DCHECK_EQ(
+            mutableNonNulls.size(),
+            remainingNonNullsCount,
+            "NonNulls buffer size should be equal to remaining non-nulls count");
+
+        std::copy_n(
+            tempNonNulls.begin() + nonNullsOffset_,
+            remainingNonNullsCount,
+            mutableNonNulls.begin());
+      }
+    }
+    lengthsOffset_ = 0;
+    nonNullsOffset_ = 0;
+    bufferOffset_ = 0;
+  }
+
+  NullableContentStringStreamData* const streamData_;
+  const uint64_t minChunkSize_;
+  const uint64_t maxChunkSize_;
+  const bool omitStream_;
+  const bool ensureFullChunks_;
+
+  size_t lengthsOffset_{0};
+  size_t nonNullsOffset_{0};
+  uint64_t bufferOffset_{0};
+};
 } // namespace facebook::nimble

dwio/nimble/velox/StreamData.h

@@ -325,4 +325,113 @@ class NullableContentStreamData final : public NullsStreamData {
   uint64_t extraMemory_;
 };
 
+struct StringBuffer {
+  Vector<char>& buffer;
+  Vector<size_t>& lengths;
+};
+
+class NullableContentStringStreamData final : public NullsStreamData {
+ public:
+  NullableContentStringStreamData(
+      velox::memory::MemoryPool& memoryPool,
+      const StreamDescriptorBuilder& descriptor,
+      const InputBufferGrowthPolicy& lengthsGrowthPolicy,
+      const InputBufferGrowthPolicy& bufferGrowthPolicy)
+      : NullsStreamData(memoryPool, descriptor, lengthsGrowthPolicy),
+        buffer_{&memoryPool},
+        lengths_{&memoryPool},
+        stringViews_{&memoryPool},
+        bufferGrowthPolicy_{bufferGrowthPolicy},
+        lengthGrowthPolicy_{lengthsGrowthPolicy} {}
+
+  inline std::string_view data() const override {
+    NIMBLE_DCHECK(materialized_, "Data must be materialized before access");
+    return std::string_view{
+        reinterpret_cast<const char*>(stringViews_.data()),
+        stringViews_.size() * sizeof(std::string_view)};
+  }
+
+  inline bool empty() const override {
+    return NullsStreamData::empty() && lengths_.empty();
+  }
+
+  inline uint64_t memoryUsed() const override {
+    return NullsStreamData::memoryUsed() + buffer_.size() +
+        lengths_.size() * sizeof(uint64_t) +
+        stringViews_.size() * sizeof(std::string_view);
+  }
+
+  inline uint64_t bufferSize() const {
+    return lengths_.size() * sizeof(uint64_t) + buffer_.size();
+  }
+
+  inline StringBuffer mutableData() {
+    // When mutableData() is called, it returns direct references to buffer_ and
+    // lengths_, allowing external code to modify them. Since dataChunkVector_
+    // contains string_view objects pointing into buffer_, any buffer
+    // modifications invalidate these views. Therefore, materialize() must be
+    // called before the next access to rebuild dataChunkVector_ with valid
+    // string_view references to the updated buffer content.
+    materialized_ = false;
+    return {buffer_, lengths_};
+  }
+
+  inline Vector<std::string_view>& mutableStringViews() {
+    materialized_ = false;
+    return stringViews_;
+  }
+
+  void materialize() override {
+    stringViews_.resize(lengths_.size());
+    auto runningOffset = 0;
+    for (auto i = 0; i < lengths_.size(); ++i) {
+      stringViews_[i] =
+          std::string_view(buffer_.data() + runningOffset, lengths_[i]);
+      runningOffset += lengths_[i];
+    }
+    NullsStreamData::materialize();
+    materialized_ = true;
+  }
+
+  void ensureStringBufferCapacity(
+      uint64_t additionalStrings,
+      uint64_t additionalBytes) {
+    if (additionalStrings == 0) {
+      return;
+    }
+
+    const auto newLengthsSize = lengths_.size() + additionalStrings;
+    if (newLengthsSize > lengths_.capacity()) {
+      const auto newCapacity = lengthGrowthPolicy_.getExtendedCapacity(
+          newLengthsSize, lengths_.capacity());
+      lengths_.reserve(newCapacity);
+    }
+
+    if (additionalBytes > 0) {
+      const auto newBufferSize = buffer_.size() + additionalBytes;
+      if (newBufferSize > buffer_.capacity()) {
+        const auto newCapacity = bufferGrowthPolicy_.getExtendedCapacity(
+            newBufferSize, buffer_.capacity());
+        buffer_.reserve(newCapacity);
+      }
+    }
+  }
+
+  inline void reset() override {
+    NullsStreamData::reset();
+    buffer_.clear();
+    lengths_.clear();
+    stringViews_.clear();
+  }
+
+ private:
+  Vector<char> buffer_;
+  Vector<uint64_t> lengths_;
+  Vector<std::string_view> stringViews_;
+  const InputBufferGrowthPolicy& bufferGrowthPolicy_;
+  const InputBufferGrowthPolicy& lengthGrowthPolicy_;
+
+  bool materialized_{false};
+};
+
 } // namespace facebook::nimble