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

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

Author: macvincent

dwio/nimble/velox/VeloxWriter.cpp

@@ -33,6 +33,7 @@
 #include "dwio/nimble/velox/SchemaSerialization.h"
 #include "dwio/nimble/velox/SchemaTypes.h"
 #include "dwio/nimble/velox/StatsGenerated.h"
+#include "dwio/nimble/velox/StreamChunker.h"
 #include "velox/common/time/CpuWallTimer.h"
 #include "velox/dwio/common/ExecutorBarrier.h"
 #include "velox/type/Type.h"
@@ -813,45 +814,34 @@ bool VeloxWriter::writeChunks(
     streams_.resize(context_->schemaBuilder.nodeCount());
 
     auto processStream = [&](StreamData& streamData) {
-      // TODO: Breakdown large streams above a threshold into smaller chunks.
-      const auto minStreamSize =
-          lastChunk ? 0 : context_->options.minStreamChunkRawSize;
       const auto* context =
           streamData.descriptor().context<WriterStreamContext>();
       bool isNullStream = context && context->isNullStream;
-      bool shouldChunkStream = false;
-      if (isNullStream) {
-        // We apply the same null logic, where if all values
-        // are non-nulls, we omit the entire stream.
-        shouldChunkStream = streamData.hasNulls() &&
-            streamData.nonNulls().size() > minStreamSize;
-      } else {
-        shouldChunkStream = streamData.data().size() > minStreamSize;
-      }
-
-      // If we have previous written chunks for this stream, during final
-      // chunk, always write any remaining data.
-      const auto offset = streamData.descriptor().offset();
+      const auto& offset = streamData.descriptor().offset();
+      auto& streamSize = context_->columnStats[offset].physicalSize;
+      logicalSizeBeforeEncoding += streamData.memoryUsed();
       NIMBLE_DASSERT(offset < streams_.size(), "Stream offset out of range.");
       auto& stream = streams_[offset];
-      if (lastChunk && !shouldChunkStream && !stream.content.empty()) {
-        shouldChunkStream =
-            !streamData.empty() || !streamData.nonNulls().empty();
-      }
-
-      if (shouldChunkStream) {
+      auto chunker = getStreamChunker(
+          streamData,
+          context_->options.maxStreamChunkRawSize,
+          lastChunk ? 0 : context_->options.minStreamChunkRawSize,
+          stream.content.empty(),
+          isNullStream);
+      while (chunker->hasNext()) {
+        auto streamDataView = chunker->next();
         std::string_view encoded;
         if (isNullStream) {
-          // For null streams we promote the null values to be written as
-          // boolean data.
+          // Null stream values are converted to boolean data for encoding.
           encoded = encodeStream(
-              *context_, *encodingBuffer_, NullsAsDataStreamData{streamData});
+              *context_,
+              *encodingBuffer_,
+              NullsAsDataStreamData{streamDataView});
         } else {
-          encoded = encodeStream(*context_, *encodingBuffer_, streamData);
+          encoded = encodeStream(*context_, *encodingBuffer_, streamDataView);
         }
 
         if (!encoded.empty()) {
-          auto& streamSize = context_->columnStats[offset].physicalSize;
           ChunkedStreamWriter chunkWriter{*encodingBuffer_};
           for (auto& buffer : chunkWriter.encode(encoded)) {
             streamSize += buffer.size();
@@ -860,9 +850,10 @@ bool VeloxWriter::writeChunks(
           }
         }
         wroteChunk = true;
-        logicalSizeBeforeEncoding += streamData.memoryUsed();
-        streamData.reset();
       }
+      // Remove processed stream data to reclaim memory.
+      chunker->compact();
+      logicalSizeBeforeEncoding -= streamData.memoryUsed();
     };
 
     const auto& streams = context_->streams();
@@ -985,39 +976,56 @@ bool VeloxWriter::tryWriteStripe(bool force) {
     // TODO: we can improve merge the last chunk write with stripe
     if (context_->options.enableChunking &&
         shouldChunk() == ChunkDecision::Chunk) {
-      const auto& streams = context_->streams();
-      const size_t streamCount = streams.size();
-      // Sort streams for chunking based on raw memory usage.
-      // TODO(T240072104): Improve performance by bucketing the streams by size
-      // (most significant bit) instead of sorting.
-      std::vector<uint32_t> streamIndices(streamCount);
-      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();
-          });
+      bool continueChunking = true;
+      auto batchChunkStreams = [&](const std::vector<uint32_t>& indices) {
+        size_t currentIndex = 0;
+        size_t indicesCount = indices.size();
+        NIMBLE_DASSERT(
+            context_->options.chunkedStreamBatchSize > 0,
+            "streamEncodingBatchSize must be greater than 0");
+        while (currentIndex < indicesCount && continueChunking) {
+          size_t endStreamIndex = std::min(
+              indicesCount,
+              currentIndex + context_->options.chunkedStreamBatchSize);
+          std::span<const uint32_t> batchIndices(
+              indices.begin() + currentIndex, endStreamIndex - currentIndex);
+          // Stop attempting chunking once streams are too small to chunk.
+          if (!writeChunks(
+                  /*lastChunk=*/false, batchIndices)) {
+            continueChunking = false;
+            break;
+          }
+          currentIndex = endStreamIndex;
+          continueChunking = (shouldChunk() == ChunkDecision::Chunk);
+        }
+      };
 
-      // Chunk streams in batches.
-      size_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) {
-        size_t endStreamIndex = std::min(
-            streamCount,
-            currentIndex + context_->options.chunkedStreamBatchSize);
-        std::span<const uint32_t> batchIndices(
-            streamIndices.data() + currentIndex, endStreamIndex - currentIndex);
-        // Stop attempting chunking once streams are too small to chunk.
-        if (!writeChunks(false, batchIndices)) {
-          break;
+      // Relieve memory pressure by chunking streams above max size.
+      const auto& streams = context_->streams();
+      std::vector<uint32_t> streamsAboveMaxChunkSize;
+      streamsAboveMaxChunkSize.reserve(streams.size());
+      for (auto streamIndex = 0; streamIndex < streams.size(); ++streamIndex) {
+        if (streams[streamIndex]->memoryUsed() >=
+            context_->options.maxStreamChunkRawSize) {
+          streamsAboveMaxChunkSize.push_back(streamIndex);
         }
-        currentIndex = endStreamIndex;
-        decision = shouldChunk();
+      }
+      batchChunkStreams(streamsAboveMaxChunkSize);
+
+      if (continueChunking) {
+        // Relieve memory pressure by chunking small streams.
+        // Sort streams for chunking based on raw memory usage.
+        // TODO(T240072104): Improve performance by bucketing the streams
+        // by size (by most significant bit) instead of sorting them.
+        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);
       }
     }
 

dwio/nimble/velox/VeloxWriterOptions.h

@@ -100,6 +100,11 @@ struct VeloxWriterOptions {
   // Note: this is ignored when it is time to flush a stripe.
   size_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 = 4 << 20;
+
   // 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

dwio/nimble/velox/tests/VeloxWriterTests.cpp

@@ -1962,6 +1962,7 @@ struct ChunkFlushPolicyTestCase {
   const uint64_t writerMemoryLowThresholdBytes{75 << 10};
   const double estimatedCompressionFactor{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> {
@@ -2096,6 +2098,7 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 75 << 10,
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 4,
             .expectedMaxChunkCount = 1,
             .expectedMinChunkCount = 1,
@@ -2110,13 +2113,29 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 75 << 10,
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 7,
             .expectedMaxChunkCount = 2,
             .expectedMinChunkCount = 1,
             .chunkedStreamBatchSize = 10,
         },
+        // Reducing maxStreamChunkRawSize produces more chunks
+        ChunkFlushPolicyTestCase{
+            .batchCount = 20,
+            .enableChunking = true,
+            .targetStripeSizeBytes = 250 << 10, // 250KB
+            .writerMemoryHighThresholdBytes = 80 << 10,
+            .writerMemoryLowThresholdBytes = 75 << 10,
+            .estimatedCompressionFactor = 1.0,
+            .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 12 << 10, // -126KB
+            .expectedStripeCount = 7,
+            .expectedMaxChunkCount = 9, // +7
+            .expectedMinChunkCount = 2, // +1
+            .chunkedStreamBatchSize = 10,
+        },
         // High memory regression threshold
-        // Produces file identical to RawStripeSizeFlushPolicy
+        // Stripe count identical to RawStripeSizeFlushPolicy
         ChunkFlushPolicyTestCase{
             .batchCount = 20,
             .enableChunking = true,
@@ -2125,8 +2144,9 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 495 << 10, // +420KB
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 4,
-            .expectedMaxChunkCount = 1,
+            .expectedMaxChunkCount = 2,
             .expectedMinChunkCount = 1,
             .chunkedStreamBatchSize = 10,
         },
@@ -2140,13 +2160,14 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 35 << 10, // -40KB
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 10,
             .expectedMaxChunkCount = 2,
             .expectedMinChunkCount = 2,
             .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,
@@ -2155,9 +2176,10 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 1 << 20, // +1MB
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 1, // -2 stripes
-            .expectedMaxChunkCount = 1,
-            .expectedMinChunkCount = 1,
+            .expectedMaxChunkCount = 5,
+            .expectedMinChunkCount = 2,
             .chunkedStreamBatchSize = 10,
 
         },
@@ -2171,6 +2193,7 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 1 << 20, // +1MB
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 7, // +6 stripes
             .expectedMaxChunkCount = 1,
             .expectedMinChunkCount = 1,
@@ -2186,6 +2209,7 @@ INSTANTIATE_TEST_CASE_P(
             .writerMemoryLowThresholdBytes = 75 << 10,
             .estimatedCompressionFactor = 1.0,
             .minStreamChunkRawSize = 100,
+            .maxStreamChunkRawSize = 128 << 10,
             .expectedStripeCount = 7,
             .expectedMaxChunkCount = 2,
             .expectedMinChunkCount = 1,