From 8a5cc976a933bb0313be74750fbb6754eed9ce19 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Wed, 24 Apr 2024 21:59:59 -0700
Subject: [PATCH] ArrayOfDoublesSketchBuildAggregator: Fix NPE in get() for
 empty sketch. (#16330)

Fixes a bug introduced in #16296, where the sketch might not be
initialized if get() is called without calling aggregate(). Also adds
a test for this case.
---
 .../ArrayOfDoublesSketchBuildAggregator.java  | 18 +++++++---
 ...ArrayOfDoublesSketchSqlAggregatorTest.java | 34 +++++++++++++++++++
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchBuildAggregator.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchBuildAggregator.java
index b093e730f0b..11ad6b218e1 100644
--- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchBuildAggregator.java
+++ b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchBuildAggregator.java
@@ -84,11 +84,6 @@ public class ArrayOfDoublesSketchBuildAggregator implements Aggregator
       values = new double[valueSelectors.length];
     }
 
-    if (sketch == null) {
-      sketch = new ArrayOfDoublesUpdatableSketchBuilder().setNominalEntries(nominalEntries)
-                                                         .setNumberOfValues(valueSelectors.length).build();
-    }
-
     final IndexedInts keys = keySelector.getRow();
     for (int i = 0; i < valueSelectors.length; i++) {
       if (valueSelectors[i].isNull()) {
@@ -113,6 +108,7 @@ public class ArrayOfDoublesSketchBuildAggregator implements Aggregator
             key.get(bytes);
             key.reset();
 
+            initializeSketchIfNeeded();
             sketch.update(bytes, values);
           }
         }
@@ -125,6 +121,7 @@ public class ArrayOfDoublesSketchBuildAggregator implements Aggregator
             key = keySelector.lookupName(keys.get(i));
           }
 
+          initializeSketchIfNeeded();
           sketch.update(key, values);
         }
       }
@@ -142,6 +139,7 @@ public class ArrayOfDoublesSketchBuildAggregator implements Aggregator
   @Override
   public synchronized Object get()
   {
+    initializeSketchIfNeeded();
     return sketch.compact();
   }
 
@@ -164,4 +162,14 @@ public class ArrayOfDoublesSketchBuildAggregator implements Aggregator
     values = null;
   }
 
+  /**
+   * Initialize {@link #sketch} if it is null.
+   */
+  private void initializeSketchIfNeeded()
+  {
+    if (sketch == null) {
+      sketch = new ArrayOfDoublesUpdatableSketchBuilder().setNominalEntries(nominalEntries)
+                                                         .setNumberOfValues(valueSelectors.length).build();
+    }
+  }
 }
diff --git a/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/tuple/sql/ArrayOfDoublesSketchSqlAggregatorTest.java b/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/tuple/sql/ArrayOfDoublesSketchSqlAggregatorTest.java
index 0ed01f3f411..7d6e322ae61 100644
--- a/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/tuple/sql/ArrayOfDoublesSketchSqlAggregatorTest.java
+++ b/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/tuple/sql/ArrayOfDoublesSketchSqlAggregatorTest.java
@@ -427,6 +427,40 @@ public class ArrayOfDoublesSketchSqlAggregatorTest extends BaseCalciteQueryTest
     );
   }
 
+  @Test
+  public void testNoInputGroupByAll()
+  {
+    cannotVectorize();
+
+    final String sql = "SELECT\n"
+                       + "  DS_TUPLE_DOUBLES(tuplesketch_dim2),\n"
+                       + "  DS_TUPLE_DOUBLES(dim2, m1)\n"
+                       + "FROM druid.foo\n"
+                       + "WHERE dim2 = 'nonexistent'\n"
+                       + "GROUP BY ()";
+
+    testQuery(
+        sql,
+        ImmutableList.of(
+            Druids.newTimeseriesQueryBuilder()
+                  .dataSource(CalciteTests.DATASOURCE1)
+                  .intervals(querySegmentSpec(Filtration.eternity()))
+                  .virtualColumns(expressionVirtualColumn("v0", "'nonexistent'", ColumnType.STRING))
+                  .filters(equality("dim2", "nonexistent", ColumnType.STRING))
+                  .granularity(Granularities.ALL)
+                  .aggregators(
+                      new ArrayOfDoublesSketchAggregatorFactory("a0", "tuplesketch_dim2", null, null, 1),
+                      new ArrayOfDoublesSketchAggregatorFactory("a1", "v0", null, ImmutableList.of("m1"), 1)
+                  )
+                  .context(QUERY_CONTEXT_DEFAULT)
+                  .build()
+        ),
+        ImmutableList.of(
+            new Object[]{"0.0", "0.0"}
+        )
+    );
+  }
+
   @Test
   public void testArrayOfDoublesSketchIntersectOnScalarExpression()
   {