From 14e51e87b8334b7fbcc5b3dfe67a7061710c5279 Mon Sep 17 00:00:00 2001
From: Yash Mayya <yash.mayya@gmail.com>
Date: Wed, 14 Aug 2024 13:31:51 +0530
Subject: [PATCH] Optimize NullableSingleInputAggregationFunction when entire
 block is null (#13758)

---
 ...ullableSingleInputAggregationFunction.java | 10 +++++++++-
 .../pinot/perf/BenchmarkModeAggregation.java  | 20 +++++++++++--------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NullableSingleInputAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NullableSingleInputAggregationFunction.java
index 907f0139d2a..af2a41610e0 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NullableSingleInputAggregationFunction.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NullableSingleInputAggregationFunction.java
@@ -81,7 +81,10 @@ public void forEachNotNull(int length, BlockValSet blockValSet, BatchConsumer co
       return;
     }
 
-    forEachNotNull(length, roaringBitmap.getIntIterator(), consumer);
+    // Skip if entire block is null
+    if (!roaringBitmap.contains(0, length)) {
+      forEachNotNull(length, roaringBitmap.getIntIterator(), consumer);
+    }
   }
 
   /**
@@ -118,6 +121,11 @@ public <A> A foldNotNull(int length, BlockValSet blockValSet, A initialAcum, Red
    * @param <A> The type of the accumulator
    */
   public <A> A foldNotNull(int length, @Nullable RoaringBitmap roaringBitmap, A initialAcum, Reducer<A> reducer) {
+    // Exit early if entire block is null
+    if (_nullHandlingEnabled && roaringBitmap != null && roaringBitmap.contains(0, length)) {
+      return initialAcum;
+    }
+
     IntIterator intIterator = roaringBitmap == null ? null : roaringBitmap.getIntIterator();
     return foldNotNull(length, intIterator, initialAcum, reducer);
   }
diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkModeAggregation.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkModeAggregation.java
index b1051a1476a..dcfff95c417 100644
--- a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkModeAggregation.java
+++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkModeAggregation.java
@@ -58,11 +58,11 @@ public class BenchmarkModeAggregation extends AbstractAggregationFunctionBenchma
   @Param({"100", "50", "0"})
   public int _nullHandlingEnabledPerCent;
   private boolean _nullHandlingEnabled;
-  @Param({"2", "4", "8", "16", "32", "64", "128"})
+  @Param({"1", "2", "4", "8", "16", "32", "64", "128"})
   protected int _nullPeriod;
   private final Random _segmentNullRandomGenerator = new Random(42);
-  private double _modeIgnoringNull;
-  private double _modeNullAware;
+  private Double _modeIgnoringNull;
+  private Double _modeNullAware;
   private final int _numDocs = DocIdSetPlanNode.MAX_DOC_PER_CALL;
 
   public static void main(String[] args)
@@ -122,11 +122,15 @@ public void setupTrial() {
         .get()
         .getKey()
         .doubleValue();
-    _modeNullAware = nullAwareDistribution.entrySet().stream()
-        .max(Comparator.comparingInt(Map.Entry::getValue))
-        .get()
-        .getKey()
-        .doubleValue();
+    if (nullAwareDistribution.isEmpty()) {
+      _modeNullAware = null;
+    } else {
+      _modeNullAware = nullAwareDistribution.entrySet().stream()
+              .max(Comparator.comparingInt(Map.Entry::getValue))
+              .get()
+              .getKey()
+              .doubleValue();
+    }
   }
 
   @Override