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.

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(
@@ -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()} {
@@ -573,4 +575,204 @@ 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.elementCount(),
+            options.minChunkSize,
+            streamData.nonNulls().empty(),
+            options.isFirstChunk)},
+        ensureFullChunks_{options.ensureFullChunks} {
+    static_assert(sizeof(bool) == 1);
+    streamData.materialize();
+  }
+
+  std::optional<StreamDataView> next() override {
+    if (omitStream_) {
+      return std::nullopt;
+    }
+    const auto& chunkSize = nextChunkSize();
+    if (chunkSize.rollingChunkSize == 0) {
+      return std::nullopt;
+    }
+
+    // Content
+    auto& dataChunkVector = streamData_->mutableStringViews();
+    dataChunkVector.resize(0);
+    dataChunkVector.reserve(chunkSize.dataElementCount);
+    auto mutableData = streamData_->mutableData();
+    auto& mutableLengths = mutableData.lengths;
+    auto& mutableBuffer = mutableData.buffer;
+    auto runningOffset = extraMemoryOffset_;
+    auto endOffset = dataElementOffset_ + chunkSize.dataElementCount;
+    for (size_t i = dataElementOffset_; i < endOffset; ++i) {
+      dataChunkVector.emplace_back(
+          std::string_view{
+              mutableBuffer.data() + runningOffset, mutableLengths[i]});
+      runningOffset += mutableLengths[i];
+    }
+    std::string_view dataChunk = {
+        reinterpret_cast<const char*>(dataChunkVector.data()),
+        chunkSize.dataElementCount * sizeof(std::string_view)};
+
+    // Nulls
+    std::span<const bool> nonNullsChunk(
+        streamData_->mutableNonNulls().data() + nonNullsOffset_,
+        chunkSize.nullElementCount);
+
+    dataElementOffset_ += chunkSize.dataElementCount;
+    nonNullsOffset_ += chunkSize.nullElementCount;
+    extraMemoryOffset_ += chunkSize.extraMemory;
+
+    if (chunkSize.nullElementCount > chunkSize.dataElementCount) {
+      return StreamDataView{
+          streamData_->descriptor(),
+          dataChunk,
+          static_cast<uint32_t>(chunkSize.nullElementCount),
+          nonNullsChunk};
+    }
+    NIMBLE_CHECK_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();
+    const auto bufferedCount = nonNulls.size();
+    ChunkSize chunkSize;
+    bool fullChunk{false};
+    // Calculate how many entries we can fit in the chunk
+    for (size_t nonNullsIndex = nonNullsOffset_; nonNullsIndex < bufferedCount;
+         ++nonNullsIndex) {
+      uint64_t elementTotalSize{sizeof(bool)};
+      uint32_t elementDataCount{0};
+      size_t elementExtraMemory{0};
+      if (nonNulls[nonNullsIndex]) {
+        elementExtraMemory =
+            stringLengths[dataElementOffset_ + chunkSize.dataElementCount];
+        elementTotalSize += elementExtraMemory + sizeof(uint64_t);
+        ++elementDataCount;
+      }
+
+      if (chunkSize.rollingChunkSize == 0 && elementTotalSize > maxChunkSize_) {
+        // Allow a single oversized string as its own chunk.
+        fullChunk = true;
+        chunkSize.extraMemory += elementExtraMemory;
+        chunkSize.dataElementCount += elementDataCount;
+        chunkSize.rollingChunkSize += elementTotalSize;
+        ++chunkSize.nullElementCount;
+        break;
+      }
+
+      if (chunkSize.rollingChunkSize + elementTotalSize > maxChunkSize_) {
+        fullChunk = true;
+        break;
+      }
+
+      chunkSize.extraMemory += elementExtraMemory;
+      chunkSize.dataElementCount += elementDataCount;
+      chunkSize.rollingChunkSize += elementTotalSize;
+      ++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 changes made to stream data, nothing to 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_CHECK(
+        streamData_->empty(), "StreamData should be empty after reset");
+
+    {
+      const auto remainingDataCount = tempLengths.size() - dataElementOffset_;
+      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() + dataElementOffset_,
+          remainingDataCount,
+          mutableDataLength.begin());
+    }
+
+    {
+      const auto remainingDataBytes = tempBuffer.size() - extraMemoryOffset_;
+      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() + extraMemoryOffset_,
+          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_CHECK_EQ(
+            mutableNonNulls.size(),
+            remainingNonNullsCount,
+            "NonNulls buffer size should be equal to remaining non-nulls count");
+
+        std::copy_n(
+            tempNonNulls.begin() + nonNullsOffset_,
+            remainingNonNullsCount,
+            mutableNonNulls.begin());
+      }
+    }
+    dataElementOffset_ = 0;
+    nonNullsOffset_ = 0;
+    extraMemoryOffset_ = 0;
+  }
+
+  NullableContentStringStreamData* const streamData_;
+  const uint64_t minChunkSize_;
+  const uint64_t maxChunkSize_;
+  const bool omitStream_;
+  const bool ensureFullChunks_;
+
+  size_t dataElementOffset_{0};
+  size_t nonNullsOffset_{0};
+  uint64_t extraMemoryOffset_{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& dataGrowthPolicy,
+      const InputBufferGrowthPolicy& bufferGrowthPolicy)
+      : NullsStreamData(memoryPool, descriptor, dataGrowthPolicy),
+        buffer_{&memoryPool},
+        lengths_{&memoryPool},
+        stringViews_{&memoryPool},
+        bufferGrowthPolicy_{bufferGrowthPolicy},
+        lengthGrowthPolicy_{dataGrowthPolicy} {}
+
+  inline std::string_view data() const override {
+    NIMBLE_CHECK(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 elementCount() const {
+    return lengths_.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] * sizeof(char));
+      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