Support Per Stream String Buffers in Velox Writer (facebookincubator#392)

Summary:

 Integrates the per-stream string buffer infrastructure (from D89933055) into the Velox Writer, enabling string fields to use per-field buffers instead of a shared global buffer. This architectural change enables incremental memory reclamation during chunking operations, improving memory efficiency for string-heavy workloads. The feature is controlled by a global [JustKnob](https://www.internalfb.com/intern/justknobs/?name=dwio%2Fnimble%3Adisable_shared_string_buffers). When the knob returns true, per-stream string buffers are enabled. When set to false (default), the existing shared buffer behavior is preserved.

Reviewed By: helfman

Differential Revision: D89938211

Author: macvincent

dwio/nimble/stats/ColumnStatsUtils.cpp

@@ -18,17 +18,30 @@
 #include "velox/dwio/common/Range.h"
 
 namespace facebook::nimble {
-uint64_t getRawSizeFromVector(
-    const velox::VectorPtr& vector,
-    const OrderedRanges& ranges,
-    RawSizeContext& context) {
+
+namespace {
+velox::common::Ranges toVeloxRanges(const OrderedRanges& ranges) {
   velox::common::Ranges veloxRanges;
   for (auto& range : ranges.ranges()) {
     veloxRanges.add(
         std::get<0>(range), std::get<0>(range) + std::get<1>(range));
   }
-  return getRawSizeFromVector(vector, veloxRanges, context);
-};
+  return veloxRanges;
+}
+} // namespace
+
+uint64_t getRawSizeFromVector(
+    const velox::VectorPtr& vector,
+    const OrderedRanges& ranges,
+    RawSizeContext& context) {
+  return getRawSizeFromVector(vector, toVeloxRanges(ranges), context);
+}
+
+uint64_t getRawSizeFromVector(
+    const velox::VectorPtr& vector,
+    const OrderedRanges& ranges) {
+  return getRawSizeFromVector(vector, toVeloxRanges(ranges));
+}
 
 void aggregateStats(
     const TypeBuilder& builder,

dwio/nimble/stats/ColumnStatsUtils.h

@@ -39,6 +39,10 @@ uint64_t getRawSizeFromVector(
     const OrderedRanges& ranges,
     RawSizeContext& context);
 
+uint64_t getRawSizeFromVector(
+    const velox::VectorPtr& vector,
+    const OrderedRanges& ranges);
+
 void aggregateStats(
     const TypeBuilder& builder,
     std::unordered_map<offset_size, ColumnStats>& columnStats,

dwio/nimble/velox/BufferGrowthPolicy.h

@@ -56,6 +56,15 @@ class DefaultInputBufferGrowthPolicy : public InputBufferGrowthPolicy {
             {32UL, 4.0}, {512UL, 1.414}, {4096UL, 1.189}});
   }
 
+  // Growth policy for per-stream string buffers when disableSharedStringBuffers
+  // is enabled. Uses byte-based thresholds since string data is stored as raw
+  // bytes.
+  static std::unique_ptr<InputBufferGrowthPolicy> withStringBufferRanges() {
+    return std::make_unique<DefaultInputBufferGrowthPolicy>(
+        std::map<uint64_t, float>{
+            {4096UL, 2.0}, {4194304UL, 1.414}, {12582912UL, 1.189}});
+  }
+
  private:
   // Map of range lowerbounds and the growth factor for the range.
   // The first lowerbound is the smallest unit of allocation and the last range

dwio/nimble/velox/FieldWriter.cpp

@@ -338,6 +338,7 @@ class SimpleFieldWriter : public FieldWriter {
       nullCount = size - nonNullCount;
     }
 
+    // NOTE: This logic is wrong. Will be removed with new stats changes.
     columnStats_.logicalSize += nullCount +
         ((K == velox::TypeKind::VARCHAR || K == velox::TypeKind::VARBINARY)
              ? valuesStream_.extraMemory()
@@ -355,6 +356,76 @@ class SimpleFieldWriter : public FieldWriter {
   ColumnStats& columnStats_;
 };
 
+template <velox::TypeKind K>
+class StringFieldWriter : public FieldWriter {
+  using SourceType = typename velox::TypeTraits<K>::NativeType;
+
+ public:
+  explicit StringFieldWriter(FieldWriterContext& context)
+      : FieldWriter(
+            context,
+            context.schemaBuilder().createScalarTypeBuilder(
+                NimbleTypeTraits<K>::scalarKind)),
+        valuesStream_{context.createNullableContentStringStreamData(
+            typeBuilder_->asScalar().scalarDescriptor())},
+        columnStats_{context.columnStats(valuesStream_.descriptor().offset())} {
+  }
+
+  void write(
+      const velox::VectorPtr& vector,
+      const OrderedRanges& ranges,
+      folly::Executor*) override {
+    // Ensure string buffer capacity.
+    auto size = ranges.size();
+    uint64_t totalBytes = getRawSizeFromVector(vector, ranges);
+    valuesStream_.ensureStringBufferCapacity(size, totalBytes);
+
+    // Append to string buffer.
+    uint64_t memoryUsed = 0;
+    auto appendToStringBuffer = [&](SourceType sv) {
+      memoryUsed += sv.size();
+      auto stringBuffer = valuesStream_.mutableData();
+      auto& buffer = stringBuffer.buffer;
+      buffer.insert(buffer.end(), sv.begin(), sv.end());
+      auto& mutableLengths = stringBuffer.lengths;
+      mutableLengths.push_back(sv.size());
+    };
+
+    uint64_t nonNullCount = 0;
+    if (auto flat = vector->asFlatVector<SourceType>()) {
+      valuesStream_.ensureAdditionalNullsCapacity(flat->mayHaveNulls(), size);
+      nonNullCount = iterateNonNullValues(
+          ranges,
+          valuesStream_.mutableNonNulls(),
+          Flat<SourceType>{vector},
+          appendToStringBuffer);
+    } else {
+      auto decodingContext = context_.decodingContext();
+      auto& decoded = decodingContext.decode(vector, ranges);
+      valuesStream_.ensureAdditionalNullsCapacity(decoded.mayHaveNulls(), size);
+      nonNullCount = iterateNonNullValues(
+          ranges,
+          valuesStream_.mutableNonNulls(),
+          Decoded<SourceType>{decoded},
+          appendToStringBuffer);
+    }
+    uint64_t nullCount = size - nonNullCount;
+
+    // TODO: Validate that this logic is correct.
+    columnStats_.logicalSize += nullCount + memoryUsed;
+    columnStats_.nullCount += nullCount;
+    columnStats_.valueCount += size;
+  }
+
+  void reset() override {
+    valuesStream_.reset();
+  }
+
+ private:
+  NullableContentStringStreamData& valuesStream_;
+  ColumnStats& columnStats_;
+};
+
 class TimestampFieldWriter : public FieldWriter {
  public:
   explicit TimestampFieldWriter(FieldWriterContext& context)
@@ -1964,15 +2035,25 @@ std::unique_ptr<FieldWriter> FieldWriter::create(
       break;
     }
     case velox::TypeKind::VARCHAR: {
-      field = std::make_unique<
-          SimpleFieldWriter<velox::TypeKind::VARCHAR, StringConverter>>(
-          context);
+      if (context.disableSharedStringBuffers()) {
+        field = std::make_unique<StringFieldWriter<velox::TypeKind::VARCHAR>>(
+            context);
+      } else {
+        field = std::make_unique<
+            SimpleFieldWriter<velox::TypeKind::VARCHAR, StringConverter>>(
+            context);
+      }
       break;
     }
     case velox::TypeKind::VARBINARY: {
-      field = std::make_unique<
-          SimpleFieldWriter<velox::TypeKind::VARBINARY, StringConverter>>(
-          context);
+      if (context.disableSharedStringBuffers()) {
+        field = std::make_unique<StringFieldWriter<velox::TypeKind::VARBINARY>>(
+            context);
+      } else {
+        field = std::make_unique<
+            SimpleFieldWriter<velox::TypeKind::VARBINARY, StringConverter>>(
+            context);
+      }
       break;
     }
     case velox::TypeKind::TIMESTAMP: {

dwio/nimble/velox/FieldWriter.h

@@ -296,6 +296,22 @@ class FieldWriterContext {
             *bufferMemoryPool_, descriptor, *inputBufferGrowthPolicy_)));
   }
 
+  nimble::NullableContentStringStreamData&
+  createNullableContentStringStreamData(
+      const nimble::StreamDescriptorBuilder& descriptor) {
+    return static_cast<nimble::NullableContentStringStreamData&>(
+        *streams_.emplace_back(
+            std::make_unique<nimble::NullableContentStringStreamData>(
+                *bufferMemoryPool_,
+                descriptor,
+                *inputBufferGrowthPolicy_,
+                *stringBufferGrowthPolicy_)));
+  }
+
+  inline bool disableSharedStringBuffers() const {
+    return disableSharedStringBuffers_;
+  }
+
  protected:
   MemoryPoolHolder bufferMemoryPool_;
   std::mutex flatMapSchemaMutex_;
@@ -305,8 +321,10 @@ class FieldWriterContext {
   folly::F14FastSet<uint32_t> dictionaryArrayNodeIds_;
   folly::F14FastSet<uint32_t> deduplicatedMapNodeIds_;
   bool ignoreTopLevelNulls_{false};
+  bool disableSharedStringBuffers_{false};
 
   std::unique_ptr<InputBufferGrowthPolicy> inputBufferGrowthPolicy_;
+  std::unique_ptr<InputBufferGrowthPolicy> stringBufferGrowthPolicy_;
   InputBufferGrowthStats inputBufferGrowthStats_;
 
   std::unordered_map<offset_size, ColumnStats> columnStats_;

dwio/nimble/velox/VeloxWriter.cpp

@@ -57,7 +57,11 @@ class WriterContext : public FieldWriterContext {
     inputBufferGrowthPolicy_ = this->options_.lowMemoryMode
         ? std::make_unique<ExactGrowthPolicy>()
         : this->options_.inputGrowthPolicyFactory();
+    stringBufferGrowthPolicy_ = this->options_.lowMemoryMode
+        ? std::make_unique<ExactGrowthPolicy>()
+        : this->options_.stringBufferGrowthPolicyFactory();
     ignoreTopLevelNulls_ = options_.ignoreTopLevelNulls;
+    disableSharedStringBuffers_ = options_.disableSharedStringBuffers;
   }
 
   const VeloxWriterOptions& options() const {
@@ -991,6 +995,7 @@ void VeloxWriter::processStream(
       encodeStream(nullsStreamData, streamSize, chunkSize);
     }
   } else {
+    streamData.materialize();
     if (!streamData.data().empty()) {
       encodeStream(streamData, streamSize, chunkSize);
     }

dwio/nimble/velox/VeloxWriterOptions.h

@@ -142,6 +142,14 @@ struct VeloxWriterOptions {
     return DefaultInputBufferGrowthPolicy::withDefaultRanges();
   };
 
+  // When per-stream string buffers are enabled (disableSharedStringBuffers),
+  // this policy controls how the string buffer vectors grow.
+  std::function<std::unique_ptr<InputBufferGrowthPolicy>()>
+      stringBufferGrowthPolicyFactory =
+          []() -> std::unique_ptr<InputBufferGrowthPolicy> {
+    return DefaultInputBufferGrowthPolicy::withStringBufferRanges();
+  };
+
   std::function<std::unique_ptr<velox::memory::MemoryReclaimer>()>
       reclaimerFactory = []() { return nullptr; };
 
@@ -178,6 +186,10 @@ struct VeloxWriterOptions {
   bool ignoreTopLevelNulls{false};
 
   bool enableStreamDeduplication{true};
+
+  // When true, string fields use per-field buffers instead of a shared buffer.
+  // This enables incremental memory reclamation during chunking.
+  bool disableSharedStringBuffers{false};
 };
 
 } // namespace facebook::nimble

dwio/nimble/velox/tests/FieldWriterStatsTests.cpp

@@ -66,6 +66,7 @@ class FieldWriterStatsTests : public ::testing::Test {
       rowType = input->type();
     }
     options.enableChunking = true;
+    options.disableSharedStringBuffers = true;
     nimble::VeloxWriter writer(
         rowType, std::move(writeFile), *rootPool_, options);
     writer.write(input);

dwio/nimble/velox/tests/VeloxWriterTest.cpp

@@ -2184,7 +2184,7 @@ TEST_P(ChunkFlushPolicyTest, ChunkFlushPolicyIntegration) {
 
 TEST_F(VeloxWriterTest, fuzzComplex) {
   auto type = velox::ROW(
-      {{"array", velox::ARRAY(velox::REAL())},
+      {{"array", velox::ARRAY(velox::VARCHAR())},
        {"dict_array", velox::ARRAY(velox::REAL())},
        {"map", velox::MAP(velox::INTEGER(), velox::DOUBLE())},
        {"row",
@@ -2204,13 +2204,14 @@ TEST_F(VeloxWriterTest, fuzzComplex) {
         velox::ARRAY(
             velox::ROW({
                 {"a", velox::INTEGER()},
-                {"b", velox::MAP(velox::REAL(), velox::REAL())},
+                {"b", velox::MAP(velox::REAL(), velox::VARBINARY())},
             }))},
        {"nested_map_array1",
         velox::MAP(velox::INTEGER(), velox::ARRAY(velox::REAL()))},
        {"nested_map_array2",
         velox::MAP(velox::INTEGER(), velox::ARRAY(velox::INTEGER()))},
        {"dict_map", velox::MAP(velox::INTEGER(), velox::INTEGER())}});
+
   auto rowType = std::dynamic_pointer_cast<const velox::RowType>(type);
   const uint32_t seed = FLAGS_writer_tests_seed > 0 ? FLAGS_writer_tests_seed
                                                     : folly::Random::rand32();
@@ -2220,6 +2221,7 @@ TEST_F(VeloxWriterTest, fuzzComplex) {
     std::shared_ptr<folly::CPUThreadPoolExecutor> executor;
     nimble::VeloxWriterOptions writerOptions;
     writerOptions.enableChunking = true;
+    writerOptions.disableSharedStringBuffers = true;
     writerOptions.flushPolicyFactory =
         []() -> std::unique_ptr<nimble::FlushPolicy> {
       return std::make_unique<nimble::ChunkFlushPolicy>(
@@ -2334,6 +2336,7 @@ TEST_F(VeloxWriterTest, batchedChunkingRelievesMemoryPressure) {
   nimble::VeloxWriterOptions writerOptions;
   writerOptions.chunkedStreamBatchSize = kBatchSize;
   writerOptions.enableChunking = true;
+  writerOptions.disableSharedStringBuffers = true;
   writerOptions.minStreamChunkRawSize = minChunkSize;
   writerOptions.maxStreamChunkRawSize = maxChunkSize;
   writerOptions.flushPolicyFactory =