refactor: polish record size estimation in compaction

1. polish record size estimation in compaction Signed-off-by: TheR1sing3un <[email protected]>
apache · Sep 27, 2024 · 25b4c74 · 25b4c74
1 parent 083b608
commit 25b4c74
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 13 deletions.
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/model/CompactionContext.java b/hudi-common/src/main/java/org/apache/hudi/common/model/CompactionContext.java
@@ -171,6 +171,7 @@ public String toString() {
         + ", estimatedBaseFileSize=" + estimatedBaseFileSize
         + ", baseFileSorted=" + baseFileSorted
         + ", estimatedLogFileSize=" + estimatedLogFileSize
+        + ", logFileNum=" + logFilePaths.size()
         + ", logFilePaths=" + logFilePaths
         + ", logFilesRealSize=" + logFilesRealSize
         + '}';

diff --git a/hudi-common/src/main/java/org/apache/hudi/common/util/SampleEstimator.java b/hudi-common/src/main/java/org/apache/hudi/common/util/SampleEstimator.java
@@ -38,15 +38,15 @@ public class SampleEstimator<R> implements SizeEstimator<R> {
   private int sampleCount;
 
   public SampleEstimator(SizeEstimator<R> underlyingEstimator) {
-    this(DEFAULT_TRIGGER_SAMPLE_THRESHOLD, DEFAULT_SAMPLE_WRIGHT, underlyingEstimator);
+    this(underlyingEstimator, DEFAULT_TRIGGER_SAMPLE_THRESHOLD);
   }
 
   public SampleEstimator(SizeEstimator<R> underlyingEstimator, int triggerSampleThreshold) {
-    this(triggerSampleThreshold, DEFAULT_SAMPLE_WRIGHT, underlyingEstimator);
+    this(underlyingEstimator, triggerSampleThreshold, DEFAULT_SAMPLE_WRIGHT);
   }
 
   // TODO: configure the triggerSampleThreshold and sampleWeight in the write config
-  public SampleEstimator(int triggerSampleThreshold, double sampleWeight, SizeEstimator<R> underlyingEstimator) {
+  public SampleEstimator(SizeEstimator<R> underlyingEstimator, int triggerSampleThreshold, double sampleWeight) {
     this.triggerSampleThreshold = triggerSampleThreshold;
     this.sampleWeight = sampleWeight;
     this.underlyingEstimator = underlyingEstimator;
@@ -56,7 +56,7 @@ public SampleEstimator(int triggerSampleThreshold, double sampleWeight, SizeEsti
   }
 
   public SampleEstimator<R> newInstance() {
-    return new SampleEstimator<>(triggerSampleThreshold, sampleWeight, underlyingEstimator);
+    return new SampleEstimator<>(underlyingEstimator, triggerSampleThreshold, sampleWeight);
   }
 
   public long getPerEstimatedSize() {

diff --git a/...ain/java/org/apache/hudi/common/util/collection/SortedAppendOnlyExternalSpillableMap.java b/...ain/java/org/apache/hudi/common/util/collection/SortedAppendOnlyExternalSpillableMap.java
@@ -411,10 +411,10 @@ public void close() throws IOException {
     }
   }
 
-  class MemoryCombinedMap {
+  private class MemoryCombinedMap {
     Map<K, V> map;
     CombineFunc<K, V, V> combineFunc;
-    long currentMemorySize;
+    long estimatedAverageKeyValueSize;
 
     public MemoryCombinedMap(Option<CombineFunc<K, V, V>> func) {
       this.map = new HashMap<K, V>();
@@ -429,20 +429,21 @@ public V insert(K key, V value) {
       map.compute(key, (k, oldValue) -> {
         if (oldValue == null) {
           V initValue = combineFunc.initCombine(key, value);
-          long estimatedSize = keySizeEstimator.sizeEstimate(key) + valueSizeEstimator.sizeEstimate(initValue);
-          currentMemorySize += estimatedSize;
+          sampleSize(key, initValue);
           return initValue;
         } else {
           V combined = combineFunc.combine(key, value, oldValue);
-          // NOTE: call size-estimate function to sample, but not add to currentMemorySize
-          keySizeEstimator.sizeEstimate(key);
-          valueSizeEstimator.sizeEstimate(combined);
+          sampleSize(key, combined);
           return combined;
         }
       });
       return map.get(key);
     }
 
+    private void sampleSize(K k, V v) {
+      estimatedAverageKeyValueSize = keySizeEstimator.sizeEstimate(k) + valueSizeEstimator.sizeEstimate(v);
+    }
+
     public V get(K key) {
       return map.get(key);
     }
@@ -456,7 +457,7 @@ public int size() {
     }
 
     public long getCurrentMemorySize() {
-      return currentMemorySize;
+      return map.size() * estimatedAverageKeyValueSize;
     }
 
     public Iterator<V> getSortedIterator() {

diff --git a/hudi-common/src/test/java/org/apache/hudi/common/util/TestSampleEstimator.java b/hudi-common/src/test/java/org/apache/hudi/common/util/TestSampleEstimator.java
@@ -31,7 +31,7 @@ public class TestSampleEstimator {
   @Test
   public void testSampleEstimator() {
     final AtomicLong returnValue = new AtomicLong(100);
-    SampleEstimator estimator = new SampleEstimator(100, 0.1, (r) -> returnValue.get());
+    SampleEstimator estimator = new SampleEstimator((r) -> returnValue.get(), 100, 0.1);
     long estimatedSize = 0;
     for (int i = 0; i < 1000; i++) {
       estimatedSize = estimator.sizeEstimate(i);