feat(Nimble): New Flush Policy Implementation With Chunking (facebookincubator#240)

Summary:
X-link: facebookincubator/velox#14846

Pull Request resolved: facebookincubator#240

This is an implementation of the new chunking policy described in this [doc](https://fburl.com/gdoc/gkdwwju1). It has two phases:

**Phase 1 - Memory Pressure Management (shouldChunk)**
The policy monitors total in-memory data size:
*  When memory usage exceeds the maximum threshold, initiates chunking to reduce memory footprint while continuing data ingestion
*  When previous chunking attempts succeeded and memory remains above the minimum threshold, continues chunking to further reduce memory usage

**Phase 2 - Storage Size Optimization (shouldFlush)**
 Implements compression-aware stripe size prediction:
*   When chunking fails to reduce memory usage effectively and memory stays above the minimum threshold, forces a full stripe flush to guarantee memory relief
*   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

`shouldChunk` is also now a separate method required by all flush policies. We updated all previous tests and code references

NOTE: The Velox repo change here is just test copied into an experimental directory that references the flush policy.

Differential Revision: D81516697

Author: macvincent

dwio/nimble/velox/FlushPolicy.cpp

@@ -23,4 +23,54 @@ FlushDecision StripeRawSizeFlushPolicy::shouldFlush(
                                                         : FlushDecision::None;
 }
 
+// Relieve memory pressure with chunking. Tracks state between calls.
+ChunkDecision ChunkFlushPolicy::shouldChunk(
+    const StripeProgress& stripeProgress) {
+  const auto relieveMemoryPressure = [&]() {
+    uint64_t inMemoryBytes =
+        stripeProgress.stripeRawSize + stripeProgress.stripeEncodedSize;
+    if (lastChunkDecision_ == ChunkDecision::None) {
+      return inMemoryBytes > config_.writerMemoryHighThresholdBytes
+          ? ChunkDecision::Chunk
+          : ChunkDecision::None;
+    }
+
+    // Continue chunking while memory pressure remains high but stripe raw
+    // size is decreasing, indicating progress in memory reduction.
+    if (lastStripeRawSize_ > stripeProgress.stripeRawSize &&
+        inMemoryBytes > config_.writerMemoryLowThresholdBytes) {
+      return ChunkDecision::Chunk;
+    }
+    return ChunkDecision::None;
+  };
+
+  lastChunkDecision_ = relieveMemoryPressure();
+  lastStripeRawSize_ = stripeProgress.stripeRawSize;
+  return lastChunkDecision_;
+}
+
+// Optimize for expected storage stripe size.
+// Does not track state between calls.
+FlushDecision ChunkFlushPolicy::shouldFlush(
+    const StripeProgress& stripeProgress) {
+  // When chunking is unable to relieve memory pressure, we flush stripe.
+  if (stripeProgress.stripeRawSize + stripeProgress.stripeEncodedSize >
+      config_.writerMemoryHighThresholdBytes) {
+    return FlushDecision::Stripe;
+  }
+
+  double compressionFactor = config_.estimatedCompressionFactor;
+  // Use historical compression ratio as a heuristic when available.
+  if (stripeProgress.stripeEncodedSize > 0) {
+    compressionFactor =
+        static_cast<double>(stripeProgress.stripeEncodedLogicalSize) /
+        stripeProgress.stripeEncodedSize;
+  }
+  double expectedEncodedStripeSize = stripeProgress.stripeEncodedSize +
+      stripeProgress.stripeRawSize / std::max(compressionFactor, 1.0);
+  return (expectedEncodedStripeSize >= config_.targetStripeSizeBytes)
+      ? FlushDecision::Stripe
+      : FlushDecision::None;
+}
+
 } // namespace facebook::nimble

dwio/nimble/velox/FlushPolicy.h

@@ -20,23 +20,45 @@
 
 namespace facebook::nimble {
 
+// TODO: Set default values for these parameters based on DISCO experiments.
+// Use abitrary values for now.
+struct ChunkFlushPolicyConfig {
+  // Threshold to trigger chunking to relieve memory pressure
+  const uint64_t writerMemoryHighThresholdBytes{200 * 1024L * 1024L};
+  // Threshold below which chunking stops and stripe size optimization resumes
+  const uint64_t writerMemoryLowThresholdBytes{100 * 1024L * 1024L};
+  // Target size for encoded stripes
+  const uint64_t targetStripeSizeBytes{100 * 1024L * 1024L};
+  // Expected ratio of raw to encoded data
+  const double estimatedCompressionFactor{1.3};
+};
+
+enum class FlushDecision : bool {
+  None = false,
+  Stripe = true,
+};
+
+enum class ChunkDecision : bool {
+  None = false,
+  Chunk = true,
+};
+
 struct StripeProgress {
   // Size of the stripe data when it's fully decompressed and decoded
   const uint64_t stripeRawSize;
   // Size of the stripe after buffered data is encoded and optionally compressed
   const uint64_t stripeEncodedSize;
-};
-
-enum class FlushDecision : uint8_t {
-  None = 0,
-  Stripe = 1,
-  Chunk = 2,
+  // Logical size of the now encoded stripe data
+  const uint64_t stripeEncodedLogicalSize;
 };
 
 class FlushPolicy {
  public:
   virtual ~FlushPolicy() = default;
   virtual FlushDecision shouldFlush(const StripeProgress& stripeProgress) = 0;
+  virtual ChunkDecision shouldChunk(const StripeProgress&) {
+    return ChunkDecision::None;
+  }
 };
 
 class StripeRawSizeFlushPolicy final : public FlushPolicy {
@@ -53,15 +75,44 @@ class StripeRawSizeFlushPolicy final : public FlushPolicy {
 class LambdaFlushPolicy : public FlushPolicy {
  public:
   explicit LambdaFlushPolicy(
-      std::function<FlushDecision(const StripeProgress&)> lambda)
-      : lambda_{std::move(lambda)} {}
+      std::function<FlushDecision(const StripeProgress&)> flushLambda =
+          [](const StripeProgress&) { return FlushDecision::None; },
+      std::function<ChunkDecision(const StripeProgress&)> chunkLambda =
+          [](const StripeProgress&) { return ChunkDecision::None; })
+      : flushLambda_{std::move(flushLambda)},
+        chunkLambda_{std::move(chunkLambda)} {}
 
   FlushDecision shouldFlush(const StripeProgress& stripeProgress) override {
-    return lambda_(stripeProgress);
+    return flushLambda_(stripeProgress);
   }
 
+  ChunkDecision shouldChunk(const StripeProgress& stripeProgress) override {
+    return chunkLambda_(stripeProgress);
+  }
+
+ private:
+  std::function<FlushDecision(const StripeProgress&)> flushLambda_;
+  std::function<ChunkDecision(const StripeProgress&)> chunkLambda_;
+};
+
+class ChunkFlushPolicy : public FlushPolicy {
+ public:
+  explicit ChunkFlushPolicy(ChunkFlushPolicyConfig config)
+      : config_{std::move(config)},
+        lastChunkDecision_{ChunkDecision::None},
+        lastStripeRawSize_{0} {}
+
+  // Optimize for expected storage stripe size.
+  // Does not track state between calls.
+  FlushDecision shouldFlush(const StripeProgress& stripeProgress) override;
+
+  // Relieve memory pressure with chunking. Tracks state between calls.
+  ChunkDecision shouldChunk(const StripeProgress& stripeProgress) override;
+
  private:
-  std::function<FlushDecision(const StripeProgress&)> lambda_;
+  const ChunkFlushPolicyConfig config_;
+  ChunkDecision lastChunkDecision_;
+  uint64_t lastStripeRawSize_;
 };
 
 } // namespace facebook::nimble

dwio/nimble/velox/VeloxWriter.cpp

@@ -846,18 +846,20 @@ bool VeloxWriter::tryWriteStripe(bool force) {
         .stripeEncodedSize = context_->stripeSize});
   };
 
-  auto decision = force ? FlushDecision::Stripe : shouldFlush();
-  if (decision == FlushDecision::None) {
-    return false;
-  }
+  auto shouldChunk = [&]() {
+    return context_->flushPolicy->shouldChunk(StripeProgress{
+        .stripeRawSize = context_->memoryUsed,
+        .stripeEncodedSize = context_->stripeSize});
+  };
 
   try {
     // TODO: we can improve merge the last chunk write with stripe
-    if (decision == FlushDecision::Chunk && context_->options.enableChunking) {
+    if (context_->options.enableChunking &&
+        shouldChunk() == ChunkDecision::Chunk) {
       writeChunk(false);
-      decision = shouldFlush();
     }
 
+    auto decision = force ? FlushDecision::Stripe : shouldFlush();
     if (decision != FlushDecision::Stripe) {
       return false;
     }

dwio/nimble/velox/selective/tests/E2EFilterTest.cpp

@@ -128,13 +128,14 @@ class E2EFilterTest : public dwio::common::E2EFilterTestBase {
     writeSchema_ = rowType_;
     VeloxWriterOptions options;
     options.enableChunking = true;
-    options.flushPolicyFactory = [] {
+    auto i = 0;
+    options.flushPolicyFactory = [&] {
       return std::make_unique<LambdaFlushPolicy>(
-          [i = 0](const StripeProgress&) mutable {
-            if (i++ % 3 == 2) {
-              return FlushDecision::Stripe;
-            }
-            return FlushDecision::Chunk;
+          [&](const StripeProgress&) {
+            return (i++ % 3 == 2) ? FlushDecision::Stripe : FlushDecision::None;
+          },
+          [&](const StripeProgress&) {
+            return (i++ % 3 == 2) ? ChunkDecision::Chunk : ChunkDecision::None;
           });
     };
     if (!flatMapColumns_.empty()) {

dwio/nimble/velox/selective/tests/SelectiveNimbleReaderTest.cpp

@@ -566,7 +566,8 @@ TEST_F(SelectiveNimbleReaderTest, multiChunkNulls) {
   options.minStreamChunkRawSize = 0;
   options.flushPolicyFactory = [] {
     return std::make_unique<LambdaFlushPolicy>(
-        [](const StripeProgress&) { return FlushDecision::Chunk; });
+        [](const StripeProgress&) { return FlushDecision::None; },
+        [](const StripeProgress&) { return ChunkDecision::Chunk; });
   };
   auto file = test::createNimbleFile(
       *rootPool(), {chunk1, chunk2, chunk3}, options, false);
@@ -612,7 +613,8 @@ TEST_F(SelectiveNimbleReaderTest, multiChunkInt16RowSetOverBoundary) {
   options.minStreamChunkRawSize = 0;
   options.flushPolicyFactory = [] {
     return std::make_unique<LambdaFlushPolicy>(
-        [](const StripeProgress&) { return FlushDecision::Chunk; });
+        [](const StripeProgress&) { return FlushDecision::None; },
+        [](const StripeProgress&) { return ChunkDecision::Chunk; });
   };
   auto file =
       test::createNimbleFile(*rootPool(), {chunk1, chunk2}, options, false);