New Flush Policy Implementation With Chunking (#240)

Summary: Pull Request resolved: #240

Differential Revision: D81516697

Author: macvincent

dwio/nimble/velox/FlushPolicy.cpp

@@ -23,4 +23,56 @@ 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 ||
+      stripeProgress.stripeEncodedSize == 0) {
+    return lastFlushDecision_;
+  }
+
+  // When no writer memory pressure, optimize for storage stripe size
+  const auto optimizeStorageSize = [&]() {
+    double compressionRatio = config_.compressionRatioFactor *
+        (stripeProgress.stripeEncodedRawSize /
+         stripeProgress.stripeEncodedSize);
+    double expectedEncodedStripeSize = stripeProgress.stripeEncodedSize +
+        compressionRatio * stripeProgress.stripeRawSize;
+    return expectedEncodedStripeSize >= config_.targetStripeSizeBytes
+        ? FlushDecision::Stripe
+        : FlushDecision::None;
+  };
+  lastFlushDecision_ = optimizeStorageSize();
+
+  return lastFlushDecision_;
+}
+
 } // namespace facebook::nimble

dwio/nimble/velox/FlushPolicy.h

@@ -20,11 +20,15 @@
 
 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;
+struct FlushPolicyConfig {
+  // 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,6 +37,17 @@ 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;
@@ -64,4 +79,16 @@ class LambdaFlushPolicy : public FlushPolicy {
   std::function<FlushDecision(const StripeProgress&)> lambda_;
 };
 
+class ChunkFlushPolicy : public FlushPolicy {
+ public:
+  explicit ChunkFlushPolicy(const FlushPolicyConfig& config)
+      : config_{config} {}
+
+  FlushDecision shouldFlush(const StripeProgress& stripeProgress) override;
+
+ private:
+  FlushPolicyConfig config_;
+  FlushDecision lastFlushDecision_{FlushDecision::None};
+};
+
 } // namespace facebook::nimble