From 254a8eb7e0c04eb8e6fe2c5d5950913faf821be2 Mon Sep 17 00:00:00 2001
From: Adarsh Sanjeev <adarshsanjeev@gmail.com>
Date: Fri, 8 Dec 2023 11:43:04 +0530
Subject: [PATCH] Add null checks for HllSketchHolder (#15502)

Fixes a potential NPE which could occur while folding the HllSketchAggregator. If the sketch is null, druid could return a null HllSketchHolder object. Adding a null check here could help here

Resolves a null pointer exception in HllSketchAggregatorFactory
---
 .../hll/HllSketchAggregatorFactory.java        |  8 ++++++--
 .../hll/HllSketchHolderObjectStrategy.java     |  1 +
 .../hll/HllSketchAggregatorFactoryTest.java    | 18 ++++++++++++++++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactory.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactory.java
index 1c908be244a..eb874dbea27 100644
--- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactory.java
+++ b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactory.java
@@ -172,8 +172,12 @@ public abstract class HllSketchAggregatorFactory extends AggregatorFactory
       @Override
       public void fold(final ColumnValueSelector selector)
       {
-        final HllSketchHolder sketch = (HllSketchHolder) selector.getObject();
-        union.update(sketch.getSketch());
+        final HllSketchHolder sketchHolder = (HllSketchHolder) selector.getObject();
+        // sketchHolder can be null here, if the sketch is empty. This is an optimisation done by
+        // HllSketchHolderObjectStrategy. If the holder is null, this should be a no-op.
+        if (sketchHolder != null) {
+          union.update(sketchHolder.getSketch());
+        }
       }
 
       @Nullable
diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchHolderObjectStrategy.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchHolderObjectStrategy.java
index 18b7f03319e..93597a2621e 100644
--- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchHolderObjectStrategy.java
+++ b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchHolderObjectStrategy.java
@@ -45,6 +45,7 @@ public class HllSketchHolderObjectStrategy implements ObjectStrategy<HllSketchHo
     return HllSketchAggregatorFactory.COMPARATOR.compare(sketch1, sketch2);
   }
 
+  @Nullable
   @Override
   public HllSketchHolder fromByteBuffer(final ByteBuffer buf, final int size)
   {
diff --git a/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactoryTest.java b/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactoryTest.java
index 09a20ab8213..dcc12f5c384 100644
--- a/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactoryTest.java
+++ b/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactoryTest.java
@@ -24,10 +24,12 @@ import org.apache.datasketches.hll.TgtHllType;
 import org.apache.druid.java.util.common.StringEncoding;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.query.Druids;
+import org.apache.druid.query.aggregation.AggregateCombiner;
 import org.apache.druid.query.aggregation.Aggregator;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.BufferAggregator;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.query.aggregation.TestObjectColumnSelector;
 import org.apache.druid.query.aggregation.post.FieldAccessPostAggregator;
 import org.apache.druid.query.aggregation.post.FinalizingFieldAccessPostAggregator;
 import org.apache.druid.query.timeseries.TimeseriesQuery;
@@ -348,6 +350,22 @@ public class HllSketchAggregatorFactoryTest
     );
   }
 
+  @Test
+  public void testFoldWithNullObject()
+  {
+    TestHllSketchAggregatorFactory factory = new TestHllSketchAggregatorFactory(
+        NAME,
+        FIELD_NAME,
+        LG_K,
+        TGT_HLL_TYPE,
+        STRING_ENCODING,
+        !ROUND
+    );
+    AggregateCombiner aggregateCombiner = factory.makeAggregateCombiner();
+    TestObjectColumnSelector objectColumnSelector = new TestObjectColumnSelector(new Object[] {null});
+    aggregateCombiner.fold(objectColumnSelector);
+  }
+
   private static boolean isToStringField(Field field)
   {
     int modfiers = field.getModifiers();