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

Summary:
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
*  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

@@ -17,10 +17,41 @@
 
 namespace facebook::nimble {
 
-FlushDecision StripeRawSizeFlushPolicy::shouldFlush(
+bool StripeRawSizeFlushPolicy::shouldFlush(
     const StripeProgress& stripeProgress) {
-  return stripeProgress.stripeRawSize >= stripeRawSize_ ? FlushDecision::Stripe
-                                                        : FlushDecision::None;
+  return stripeProgress.stripeRawSize >= stripeRawSize_;
+}
+
+// Relieve memory pressure with chunking. Tracks state between calls.
+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.
+// Does not track state between calls.
+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;
+  }
+
+  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,43 @@
 
 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&) {
+    return false;
+  }
 };
 
 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;
 
  private:
   const uint64_t stripeRawSize_;
@@ -53,15 +65,41 @@ 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.
+  // Does not track state between calls.
+  bool shouldFlush(const StripeProgress& stripeProgress) override;
+
+  // Relieve memory pressure with chunking. Tracks state between calls.
+  bool shouldChunk(const StripeProgress& stripeProgress) override;
+
+ private:
+  const ChunkFlushPolicyConfig config_;
+  bool lastChunkDecision_;
 };
 
 } // namespace facebook::nimble

dwio/nimble/velox/VeloxWriter.cpp

@@ -846,19 +846,19 @@ 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,11 @@ 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) ? true : false; },
+          [&](const StripeProgress&) { return (i++ % 3 == 2) ? true : false; });
     };
     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; });
+        [](const StripeProgress&) { return false; },
+        [](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; });
+        [](const StripeProgress&) { return false; },
+        [](const StripeProgress&) { return true; });
   };
   auto file =
       test::createNimbleFile(*rootPool(), {chunk1, chunk2}, options, false);