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

Summary:
X-link: facebookexternal/presto-facebook#3412

X-link: facebookincubator/velox#14846


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
*  While 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 maximum 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.

Reviewed By: helfman

Differential Revision: D81516697

Author: macvincent

dwio/nimble/velox/FlushPolicy.cpp

@@ -16,11 +16,35 @@
 #include "dwio/nimble/velox/FlushPolicy.h"
 
 namespace facebook::nimble {
+// Relieve memory pressure with chunking.
+bool ChunkFlushPolicy::shouldChunk(const StripeProgress& stripeProgress) {
+  const uint64_t inMemoryBytes =
+      stripeProgress.stripeRawSize + stripeProgress.stripeEncodedSize;
+  const auto writerMemoryThreshold = (lastChunkDecision_ == false)
+      ? config_.writerMemoryHighThresholdBytes
+      : config_.writerMemoryLowThresholdBytes;
+  lastChunkDecision_ = inMemoryBytes > writerMemoryThreshold;
+  return lastChunkDecision_;
+}
+
+// Optimize for expected storage stripe size.
+bool ChunkFlushPolicy::shouldFlush(const StripeProgress& stripeProgress) {
+  // When chunking is unable to relieve memory pressure, we flush stripe.
+  if (stripeProgress.stripeRawSize + stripeProgress.stripeEncodedSize >
+      config_.writerMemoryHighThresholdBytes) {
+    return true;
+  }
 
-FlushDecision StripeRawSizeFlushPolicy::shouldFlush(
-    const StripeProgress& stripeProgress) {
-  return stripeProgress.stripeRawSize >= stripeRawSize_ ? FlushDecision::Stripe
-                                                        : FlushDecision::None;
+  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);
 }
 
 } // namespace facebook::nimble

dwio/nimble/velox/FlushPolicy.h

@@ -20,31 +20,47 @@
 
 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};
+};
+
 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 bool shouldFlush(const StripeProgress& stripeProgress) = 0;
+  virtual bool shouldChunk(const StripeProgress& stripeProgress) = 0;
 };
 
 class StripeRawSizeFlushPolicy final : public FlushPolicy {
  public:
   explicit StripeRawSizeFlushPolicy(uint64_t stripeRawSize)
       : stripeRawSize_{stripeRawSize} {}
 
-  FlushDecision shouldFlush(const StripeProgress& stripeProgress) override;
+  bool shouldFlush(const StripeProgress& stripeProgress) override {
+    return stripeProgress.stripeRawSize >= stripeRawSize_;
+  }
+
+  bool shouldChunk(const StripeProgress&) override {
+    return false;
+  }
 
  private:
   const uint64_t stripeRawSize_;
@@ -53,15 +69,40 @@ class StripeRawSizeFlushPolicy final : public FlushPolicy {
 class LambdaFlushPolicy : public FlushPolicy {
  public:
   explicit LambdaFlushPolicy(
-      std::function<FlushDecision(const StripeProgress&)> lambda)
-      : lambda_{std::move(lambda)} {}
+      std::function<bool(const StripeProgress&)> flushLambda =
+          [](const StripeProgress&) { return false; },
+      std::function<bool(const StripeProgress&)> chunkLambda =
+          [](const StripeProgress&) { return false; })
+      : flushLambda_{std::move(flushLambda)},
+        chunkLambda_{std::move(chunkLambda)} {}
+
+  bool shouldFlush(const StripeProgress& stripeProgress) override {
+    return flushLambda_(stripeProgress);
+  }
 
-  FlushDecision shouldFlush(const StripeProgress& stripeProgress) override {
-    return lambda_(stripeProgress);
+  bool shouldChunk(const StripeProgress& stripeProgress) override {
+    return chunkLambda_(stripeProgress);
   }
 
  private:
-  std::function<FlushDecision(const StripeProgress&)> lambda_;
+  std::function<bool(const StripeProgress&)> flushLambda_;
+  std::function<bool(const StripeProgress&)> chunkLambda_;
+};
+
+class ChunkFlushPolicy : public FlushPolicy {
+ public:
+  explicit ChunkFlushPolicy(ChunkFlushPolicyConfig config)
+      : config_{std::move(config)}, lastChunkDecision_{false} {}
+
+  // Optimize for expected storage stripe size.
+  bool shouldFlush(const StripeProgress& stripeProgress) override;
+
+  // Relieve memory pressure with chunking.
+  bool shouldChunk(const StripeProgress& stripeProgress) override;
+
+ private:
+  const ChunkFlushPolicyConfig config_;
+  bool lastChunkDecision_;
 };
 
 } // namespace facebook::nimble

dwio/nimble/velox/VeloxWriter.cpp

@@ -848,19 +848,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()) {
       writeChunk(false);
-      decision = shouldFlush();
     }
 
-    if (decision != FlushDecision::Stripe) {
+    if (!(force || shouldFlush())) {
       return false;
     }
 

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

@@ -128,14 +128,12 @@ 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;
-          });
+          /*flushLambda=*/[&](const StripeProgress&) { return (i++ % 3 == 2); },
+          /*chunkLambda=*/
+          [&](const StripeProgress&) { return (i++ % 3 == 2); });
     };
     if (!flatMapColumns_.empty()) {
       setUpFlatMapColumns();

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; });
+        /*flushLambda=*/[](const StripeProgress&) { return false; },
+        /*chunkLambda=*/[](const StripeProgress&) { return true; });
   };
   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; });
+        /*flushLambda=*/[](const StripeProgress&) { return false; },
+        /*chunkLambda=*/[](const StripeProgress&) { return true; });
   };
   auto file =
       test::createNimbleFile(*rootPool(), {chunk1, chunk2}, options, false);