New Flush Policy Implementation With Chunking (#240)

Summary:

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

**Phase 1 - Memory Pressure Management:**

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
*   When chunking fails to reduce memory usage effectively and memory stays above the minimum threshold, forces a full stripe flush to guarantee memory relief

**Phase 2 - Storage Size Optimization:**

Only executed when no memory pressure exists. Implements compression-aware stripe size prediction:

*  Uses a configurable compression factor as baseline, enhanced with historical compression ratios from previously encoded data when available
*   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: D81516697

Author: macvincent

dwio/nimble/velox/FlushPolicy.cpp

@@ -23,4 +23,62 @@ FlushDecision RawStripeSizeFlushPolicy::shouldFlush(
                                                         : FlushDecision::None;
 }
 
+FlushDecision ChunkFlushPolicy::shouldFlush(
+    const StripeProgress& stripeProgress) {
+  const auto relieveMemoryPressure = [&]() {
+    const uint64_t inMemoryByteSize =
+        stripeProgress.stripeRawSize + stripeProgress.stripeEncodedSize;
+    // Determine when we need to start relieving memory pressure with chunking
+    if (lastFlushDecision_ == FlushDecision::None &&
+        inMemoryByteSize > config_.writerMaxMemoryBytes) {
+      return FlushDecision::Chunk;
+    }
+
+    // Try chunking when possible to relieve memory pressure
+    else if (
+        lastFlushDecision_ == FlushDecision::Chunk &&
+        stripeProgress.successfullyChunked &&
+        inMemoryByteSize > config_.writerMinMemoryBytes) {
+      return FlushDecision::Chunk;
+    }
+
+    // When chunking is unable to relieve memory pressure, we flush
+    else if (
+        lastFlushDecision_ == FlushDecision::Chunk &&
+        !stripeProgress.successfullyChunked &&
+        inMemoryByteSize > config_.writerMinMemoryBytes) {
+      return FlushDecision::Stripe;
+    }
+
+    return FlushDecision::None;
+  };
+
+  lastFlushDecision_ = relieveMemoryPressure();
+  if (lastFlushDecision_ != FlushDecision::None) {
+    return lastFlushDecision_;
+  }
+
+  // When no writer memory pressure, optimize for storage stripe size
+  const auto optimizeStorageSize = [&]() {
+    double compressionRatio = config_.compressionRatioFactor;
+    // Use historical compression ratio as a heuristic when available.
+    if (stripeProgress.stripeEncodedSize > 0) {
+      compressionRatio *= stripeProgress.stripeEncodedRawSize /
+          stripeProgress.stripeEncodedSize;
+    }
+    double expectedEncodedStripeSize = stripeProgress.stripeEncodedSize +
+        compressionRatio * stripeProgress.stripeRawSize;
+    return expectedEncodedStripeSize >= config_.targetStripeSizeBytes
+        ? FlushDecision::Stripe
+        : FlushDecision::None;
+  };
+  lastFlushDecision_ = optimizeStorageSize();
+
+  return lastFlushDecision_;
+}
+
+void ChunkFlushPolicy::reset() {
+  lastFlushDecision_ = FlushDecision::None;
+}
+
 } // namespace facebook::nimble

dwio/nimble/velox/FlushPolicy.h

@@ -20,11 +20,16 @@
 
 namespace facebook::nimble {
 
-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;
+// TODO: Set default values for these parameters.
+struct ChunkFlushPolicyConfig {
+  // Threshold to trigger chunking to relieve memory pressure
+  const uint64_t writerMaxMemoryBytes;
+  // Threshold below which chunking stops and stripe size optimization resumes
+  const uint64_t writerMinMemoryBytes;
+  // Target size for encoded stripes
+  const uint64_t targetStripeSizeBytes;
+  // Controls effect of current compression ratio on final stripe size
+  const double compressionRatioFactor{0.0};
 };
 
 enum class FlushDecision : uint8_t {
@@ -33,10 +38,22 @@ enum class FlushDecision : uint8_t {
   Chunk = 2,
 };
 
+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;
+  // Previous raw size of the now encoded stripe data
+  const uint64_t stripeEncodedRawSize;
+  // True if the last suggested chunk operation was successful
+  const bool successfullyChunked;
+};
+
 class FlushPolicy {
  public:
   virtual ~FlushPolicy() = default;
   virtual FlushDecision shouldFlush(const StripeProgress& stripeProgress) = 0;
+  virtual void reset() { /* no-op */ }
 };
 
 class RawStripeSizeFlushPolicy final : public FlushPolicy {
@@ -64,4 +81,18 @@ class LambdaFlushPolicy : public FlushPolicy {
   std::function<FlushDecision(const StripeProgress&)> lambda_;
 };
 
+class ChunkFlushPolicy : public FlushPolicy {
+ public:
+  explicit ChunkFlushPolicy(const ChunkFlushPolicyConfig& config)
+      : config_{config} {}
+
+  FlushDecision shouldFlush(const StripeProgress& stripeProgress) override;
+
+  void reset() override;
+
+ private:
+  ChunkFlushPolicyConfig config_;
+  FlushDecision lastFlushDecision_{FlushDecision::None};
+};
+
 } // namespace facebook::nimble