Integrate Max Stream Size Chunking in Velox Writer (#249)

macvincent · facebook-github-bot · commit 1c1fd00c2a25 · 2025-09-29T18:58:49.000-07:00
Summary: This is the last feature of the new chunking policy described in this [doc](https://fburl.com/gdoc/gkdwwju1). Here, we break down large streams into multiple chunks of size up to `maxStreamChunkRawSize`. This protects the reader from attempting to materialize huge chunks. We included StreamData support for this in the previous diff. In this diff, we integrate with the VeloxWriter. With this change, while memory pressure is detected, we: 1. Chunk large streams above `maxStreamChunkRawSize`, retaining stream data below the limit. 2. If there is still memory pressure after the first step, chunk streams with size above `minStreamChunkRawSize`. During stripe flush, we chunk all remaining data, breaking down streams above `maxStreamChunkRawSize` into smaller chunks. --- The general chunking policy has two phases: ## **Phase 1 - Memory Pressure Management (shouldChunk)** The policy monitors total in-memory data size: - When memory usage exceeds the maximum threshold, it initiates chunking to reduce memory footprint while continuing data ingestion. - When previous chunking attempts succeeded and memory remains above the minimum threshold, it continues chunking to further reduce memory usage. - When chunking fails to reduce memory usage effectively and memory stays above the minimum threshold, it forces a full stripe flush to guarantee memory relief. ## **Phase 2 - Storage Size Optimization (shouldFlush)** Implements compression-aware stripe size prediction: - Calculates the anticipated final compressed stripe size by applying the estimated compression ratio to unencoded data. - Triggers stripe flush when the predicted compressed size reaches the target stripe size threshold. Differential Revision: D82175496
diff --git a/dwio/nimble/velox/VeloxWriter.cpp b/dwio/nimble/velox/VeloxWriter.cpp
@@ -156,6 +156,69 @@ std::string_view encode(
   }
 }
 
+template <typename T>
+std::string_view encode(
+    std::optional<EncodingLayout> encodingLayout,
+    detail::WriterContext& context,
+    Buffer& buffer,
+    const StreamDataView& streamData) {
+  NIMBLE_DASSERT(
+      streamData.data().size() % sizeof(T) == 0,
+      fmt::format("Unexpected size {}", streamData.data().size()));
+  std::span<const T> data{
+      reinterpret_cast<const T*>(streamData.data().data()),
+      streamData.data().size() / sizeof(T)};
+
+  std::unique_ptr<EncodingSelectionPolicy<T>> policy;
+  if (encodingLayout.has_value()) {
+    policy = std::make_unique<ReplayedEncodingSelectionPolicy<T>>(
+        encodingLayout.value(),
+        context.options.compressionOptions,
+        context.options.encodingSelectionPolicyFactory);
+
+  } else {
+    policy = std::unique_ptr<EncodingSelectionPolicy<T>>(
+        static_cast<EncodingSelectionPolicy<T>*>(
+            context.options
+                .encodingSelectionPolicyFactory(TypeTraits<T>::dataType)
+                .release()));
+  }
+
+  if (streamData.hasNulls()) {
+    std::span<const bool> notNulls = streamData.nonNulls();
+    return EncodingFactory::encodeNullable(
+        std::move(policy), data, notNulls, buffer);
+  } else {
+    return EncodingFactory::encode(std::move(policy), data, buffer);
+  }
+}
+
+template <typename T>
+std::string_view encodeStreamTyped(
+    detail::WriterContext& context,
+    Buffer& buffer,
+    const StreamDataView& streamData) {
+  const auto* streamContext =
+      streamData.descriptor().context<WriterStreamContext>();
+
+  std::optional<EncodingLayout> encodingLayout;
+  if (streamContext && streamContext->encoding) {
+    encodingLayout.emplace(*streamContext->encoding);
+  }
+
+  try {
+    return encode<T>(encodingLayout, context, buffer, streamData);
+  } catch (const NimbleUserError& e) {
+    if (e.errorCode() != error_code::IncompatibleEncoding ||
+        !encodingLayout.has_value()) {
+      throw;
+    }
+
+    // Incompatible captured encoding.Try again without a captured encoding.
+    return encode<T>(std::nullopt, context, buffer, streamData);
+  }
+}
+
 template <typename T>
 std::string_view encodeStreamTyped(
     detail::WriterContext& context,
@@ -213,6 +276,37 @@ std::string_view encodeStream(
   }
 }
 
+std::string_view encodeStream(
+    detail::WriterContext& context,
+    Buffer& buffer,
+    const StreamDataView& streamData) {
+  auto scalarKind = streamData.descriptor().scalarKind();
+  switch (scalarKind) {
+    case ScalarKind::Bool:
+      return encodeStreamTyped<bool>(context, buffer, streamData);
+    case ScalarKind::Int8:
+      return encodeStreamTyped<int8_t>(context, buffer, streamData);
+    case ScalarKind::Int16:
+      return encodeStreamTyped<int16_t>(context, buffer, streamData);
+    case ScalarKind::Int32:
+      return encodeStreamTyped<int32_t>(context, buffer, streamData);
+    case ScalarKind::UInt32:
+      return encodeStreamTyped<uint32_t>(context, buffer, streamData);
+    case ScalarKind::Int64:
+      return encodeStreamTyped<int64_t>(context, buffer, streamData);
+    case ScalarKind::Float:
+      return encodeStreamTyped<float>(context, buffer, streamData);
+    case ScalarKind::Double:
+      return encodeStreamTyped<double>(context, buffer, streamData);
+    case ScalarKind::String:
+    case ScalarKind::Binary:
+      return encodeStreamTyped<std::string_view>(context, buffer, streamData);
+    default:
+      NIMBLE_UNREACHABLE(
+          fmt::format("Unsupported scalar kind {}", toString(scalarKind)));
+  }
+}
+
 template <typename Set>
 void findNodeIds(
     const velox::dwio::common::TypeWithId& typeWithId,
@@ -643,6 +737,7 @@ void VeloxWriter::flush() {
 
 bool VeloxWriter::writeChunk(
     bool lastChunk,
+    bool encodeBelowMax,
     const std::unordered_set<uint32_t>& streamIndices) {
   uint64_t previousFlushWallTime = context_->stripeFlushTiming.wallNanos;
   std::atomic<uint64_t> chunkSize = 0;
@@ -713,20 +808,48 @@ bool VeloxWriter::writeChunk(
 
     auto encode = [&](StreamData& streamData, uint64_t& streamSize) {
       const auto offset = streamData.descriptor().offset();
-      auto encoded = encodeStream(*context_, *encodingBuffer_, streamData);
-      if (!encoded.empty()) {
-        ChunkedStreamWriter chunkWriter{*encodingBuffer_};
-        NIMBLE_DASSERT(offset < streams_.size(), "Stream offset out of range.");
-        auto& stream = streams_[offset];
-        for (auto& buffer : chunkWriter.encode(encoded)) {
-          streamSize += buffer.size();
-          chunkSize += buffer.size();
-          stream.content.push_back(std::move(buffer));
+      uint64_t streamSizeBeforeEncoding = streamData.memoryUsed();
+      auto writeEncoded = [&](std::string_view encoded) {
+        if (!encoded.empty()) {
+          ChunkedStreamWriter chunkWriter{*encodingBuffer_};
+          NIMBLE_DASSERT(
+              offset < streams_.size(), "Stream offset out of range.");
+          auto& stream = streams_[offset];
+          for (auto& buffer : chunkWriter.encode(encoded)) {
+            streamSize += buffer.size();
+            chunkSize += buffer.size();
+            stream.content.push_back(std::move(buffer));
+          }
+        }
+      };
+
+      // Encoded large streams as smaller chunks.
+      if (context_->options.enableChunking) {
+        while (auto chunkedStream = streamData.nextChunk(
+                   context_->options.maxStreamChunkRawSize)) {
+          auto encoded =
+              encodeStream(*context_, *encodingBuffer_, *chunkedStream);
+          writeEncoded(encoded);
+        }
+
+        // Encode small streams.
+        if ((lastChunk || encodeBelowMax) && streamData.memoryUsed()) {
+          if (auto lastStreamDataChunk = streamData.lastChunk()) {
+            auto encoded =
+                encodeStream(*context_, *encodingBuffer_, *lastStreamDataChunk);
+            writeEncoded(encoded);
+          }
+          streamData.reset();
         }
+      } else {
+        auto encoded = encodeStream(*context_, *encodingBuffer_, streamData);
+        writeEncoded(encoded);
+        streamData.reset();
       }
+
+      logicalSizeBeforeEncoding +=
+          streamSizeBeforeEncoding - streamData.memoryUsed();
       wroteChunk = true;
-      logicalSizeBeforeEncoding += streamData.memoryUsed();
-      streamData.reset();
     };
 
     auto processStream = [&](StreamData& streamData,
@@ -894,38 +1017,66 @@ bool VeloxWriter::tryWriteStripe(bool force) {
     // TODO: we can improve merge the last chunk write with stripe
     if (context_->options.enableChunking &&
         shouldChunk() == ChunkDecision::Chunk) {
+      ChunkDecision decision = ChunkDecision::Chunk;
       const auto& streams = context_->streams();
-      // Sort streams for chunking based on raw memory usage.
-      std::vector<uint32_t> streamIndices(streams.size());
-      std::iota(streamIndices.begin(), streamIndices.end(), 0);
-      std::sort(
-          streamIndices.begin(),
-          streamIndices.end(),
-          [&](const uint32_t& a, const uint32_t& b) {
-            return streams[a]->memoryUsed() > streams[b]->memoryUsed();
-          });
+      auto batchChunkStreams = [&](const std::vector<uint32_t>& indices,
+                                   bool encodeBelowMax) {
+        uint32_t currentIndex = 0;
+        NIMBLE_DASSERT(
+            context_->options.chunkedStreamBatchSize > 0,
+            "streamEncodingBatchSize must be greater than 0");
+        while (currentIndex < indices.size() &&
+               decision == ChunkDecision::Chunk) {
+          uint32_t endStreamIndex = std::min(
+              static_cast<uint32_t>(indices.size()),
+              currentIndex + context_->options.chunkedStreamBatchSize);
+          std::unordered_set<uint32_t> streamIndicesToChunk(
+              indices.begin() + currentIndex, indices.begin() + endStreamIndex);
+          currentIndex = endStreamIndex;
+          // Stop attempting chunking once streams are too small to chunk.
+          if (!writeChunk(
+                  /*lastChunk=*/false,
+                  /*encodeBelowMax=*/encodeBelowMax,
+                  streamIndicesToChunk)) {
+            decision = ChunkDecision::None;
+            break;
+          }
+          decision = shouldChunk();
+        }
+      };
+
+      // Chunk streams above maxStreamChunkRawSize to relieve memory pressure
+      auto chunkStreamsAboveMaxSize = [&]() {
+        std::vector<uint32_t> streamsAboveMaxChunkSize;
+        for (auto streamIndex = 0; streamIndex < streams.size();
+             streamIndex++) {
+          if (streams[streamIndex]->memoryUsed() >=
+              context_->options.maxStreamChunkRawSize) {
+            streamsAboveMaxChunkSize.push_back(streamIndex);
+          }
+        }
+        batchChunkStreams(streamsAboveMaxChunkSize, /*encodeBelowMax=*/false);
+      };
 
-      // Chunk streams in batches.
-      uint32_t currentIndex = 0;
-      ChunkDecision decision = ChunkDecision::Chunk;
-      NIMBLE_DASSERT(
-          context_->options.chunkedStreamBatchSize > 0,
-          "streamEncodingBatchSize must be greater than 0");
-      while (currentIndex < streams.size() &&
-             decision == ChunkDecision::Chunk) {
-        uint32_t endStreamIndex = std::min(
-            static_cast<uint32_t>(streams.size()),
-            currentIndex + context_->options.chunkedStreamBatchSize);
-        std::unordered_set<uint32_t> streamIndicesToChunk(
-            streamIndices.begin() + currentIndex,
-            streamIndices.begin() + endStreamIndex);
-        currentIndex = endStreamIndex;
-        // Stop attempting chunking once streams are too small to chunk.
-        if (!writeChunk(false, streamIndicesToChunk)) {
-          break;
+      // Relieve memory pressure by chunking small streams.
+      auto chunkSmallStreams = [&]() {
+        if (decision != ChunkDecision::Chunk) {
+          return;
         }
-        decision = shouldChunk();
-      }
+        // Sort streams for chunking based on raw memory usage.
+        std::vector<uint32_t> streamIndices(streams.size());
+        std::iota(streamIndices.begin(), streamIndices.end(), 0);
+        std::sort(
+            streamIndices.begin(),
+            streamIndices.end(),
+            [&](const uint32_t& a, const uint32_t& b) {
+              return streams[a]->memoryUsed() > streams[b]->memoryUsed();
+            });
+        batchChunkStreams(streamIndices, /*encodeBelowMax=*/true);
+      };
+
+      chunkStreamsAboveMaxSize();
+      chunkSmallStreams();
     }
 
     auto decision = force ? FlushDecision::Stripe : shouldFlush();
diff --git a/dwio/nimble/velox/VeloxWriter.h b/dwio/nimble/velox/VeloxWriter.h
@@ -89,6 +89,7 @@ class VeloxWriter {
   // Returns 'true' if chunk was written.
   bool writeChunk(
       bool lastChunk = true,
+      bool encodeBelowMax = true,
       const std::unordered_set<uint32_t>& streamIndices = {});
   uint32_t writeStripe();
 };
diff --git a/dwio/nimble/velox/VeloxWriterOptions.h b/dwio/nimble/velox/VeloxWriterOptions.h
@@ -100,6 +100,11 @@ struct VeloxWriterOptions {
   // Note: this is ignored when it is time to flush a stripe.
   uint32_t chunkedStreamBatchSize = 1024;
 
+  // When flushing data streams into chunks, streams with raw data size larger
+  // than this threshold will be broken down into multiple smaller chunks. Each
+  // chunk will be at most this size.
+  uint64_t maxStreamChunkRawSize = std::numeric_limits<int32_t>::max() / 2;
+
   // The factory function that produces the root encoding selection policy.
   // Encoding selection policy is the way to balance the tradeoffs of
   // different performance factors (at both read and write times). Heuristics
diff --git a/dwio/nimble/velox/tests/VeloxWriterTests.cpp b/dwio/nimble/velox/tests/VeloxWriterTests.cpp
@@ -1962,6 +1962,7 @@ struct ChunkFlushPolicyTestCase {
   const uint64_t writerMemoryLowThreshold{75 << 10};
   const double compressionRatioFactor{1.0};
   const uint32_t minStreamChunkRawSize{100};
+  const uint32_t maxStreamChunkRawSize{128 << 10};
   const uint32_t expectedStripeCount{0};
   const uint32_t expectedMaxChunkCount{0};
   const uint32_t expectedMinChunkCount{0};
@@ -1977,6 +1978,7 @@ TEST_P(ChunkFlushPolicyTest, ChunkFlushPolicyIntegration) {
       {{"BIGINT", velox::BIGINT()}, {"SMALLINT", velox::SMALLINT()}});
   nimble::VeloxWriterOptions writerOptions{
       .minStreamChunkRawSize = GetParam().minStreamChunkRawSize,
+            .maxStreamChunkRawSize = GetParam().maxStreamChunkRawSize,
       .chunkedStreamBatchSize = GetParam().chunkedStreamBatchSize,
       .flushPolicyFactory = GetParam().enableChunking
           ? []() -> std::unique_ptr<nimble::FlushPolicy> {
@@ -2097,6 +2099,7 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThreshold = 75 << 10,
             .compressionRatioFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 4,
             .expectedMaxChunkCount = 1,
             .expectedMinChunkCount = 1,
@@ -2111,13 +2114,29 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThreshold = 75 << 10,
             .compressionRatioFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 3,
             .expectedMaxChunkCount = 7,
             .expectedMinChunkCount = 3,
             .chunkedStreamBatchSize = 10,
         },
+        // Reducing maxStreamChunkRawSize produces more chunks
+        ChunkFlushPolicyTestCase{
+            .batchCount = 20,
+            .enableChunking = true,
+            .targetStripeSizeBytes = 250 << 10, // 250KB
+            .writerMemoryHighThreshold = 80 << 10,
+            .writerMemoryLowThreshold = 75 << 10,
+            .compressionRatioFactor = 1.0,
+            .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 12 << 10, // -126KB
+            .expectedStripeCount = 3,
+            .expectedMaxChunkCount = 19, // +12
+            .expectedMinChunkCount = 1, // -2
+            .chunkedStreamBatchSize = 10,
+        },
         // High memory regression threshold
-        // Produces file identical to RawStripeSizeFlushPolicy
+        // Stripe count identical to RawStripeSizeFlushPolicy
         ChunkFlushPolicyTestCase{
             .batchCount = 20,
             .enableChunking = true,
@@ -2126,8 +2145,9 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThreshold = 495 << 10, // +420KB
             .compressionRatioFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 4,
-            .expectedMaxChunkCount = 1,
+            .expectedMaxChunkCount = 2,
             .expectedMinChunkCount = 1,
             .chunkedStreamBatchSize = 10,
         },
@@ -2141,13 +2161,14 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThreshold = 35 << 10, // -40KB
             .compressionRatioFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 3,
             .expectedMaxChunkCount = 8,
             .expectedMinChunkCount = 4,
             .chunkedStreamBatchSize = 10,
         },
         // High target stripe size bytes (with disabled memory pressure
-        // optimization) produces fewer stripes. Single chunks.
+        // optimization) produces fewer stripes.
         ChunkFlushPolicyTestCase{
             .batchCount = 20,
             .enableChunking = true,
@@ -2156,9 +2177,10 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThreshold = 1 << 20, // +1MB
             .compressionRatioFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 1, // -2 stripes
-            .expectedMaxChunkCount = 1,
-            .expectedMinChunkCount = 1,
+            .expectedMaxChunkCount = 5,
+            .expectedMinChunkCount = 2,
             .chunkedStreamBatchSize = 10,
 
         },
@@ -2172,6 +2194,7 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThreshold = 1 << 20, // +1MB
             .compressionRatioFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 7, // +6 stripes
             .expectedMaxChunkCount = 1,
             .expectedMinChunkCount = 1,
