Fix NPE in ArrayAgg functions (#13358)

apache · Jun 11, 2024 · 5c75b15 · 5c75b15
1 parent ae4e7de
commit 5c75b15
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 0 deletions.
diff --git a/...va/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggDoubleFunction.java b/...va/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggDoubleFunction.java
@@ -93,6 +93,9 @@ public I merge(I intermediateResult1, I intermediateResult2) {
 
   @Override
   public DoubleArrayList extractFinalResult(I doubleArrayList) {
+    if (doubleArrayList == null) {
+      return new DoubleArrayList();
+    }
     return new DoubleArrayList(doubleArrayList);
   }
 }
diff --git a/...ava/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggFloatFunction.java b/...ava/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggFloatFunction.java
@@ -93,6 +93,9 @@ public I merge(I intermediateResult1, I intermediateResult2) {
 
   @Override
   public FloatArrayList extractFinalResult(I floatArrayList) {
+    if (floatArrayList == null) {
+      return new FloatArrayList();
+    }
     return new FloatArrayList(floatArrayList);
   }
 }
diff --git a/.../java/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggIntFunction.java b/.../java/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggIntFunction.java
@@ -98,6 +98,9 @@ public I merge(I intermediateResult1, I intermediateResult2) {
 
   @Override
   public IntArrayList extractFinalResult(I intArrayList) {
+    if (intArrayList == null) {
+      return new IntArrayList();
+    }
     return new IntArrayList(intArrayList);
   }
 }
diff --git a/...java/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggLongFunction.java b/...java/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggLongFunction.java
@@ -96,6 +96,9 @@ public I merge(I intermediateResult1, I intermediateResult2) {
 
   @Override
   public LongArrayList extractFinalResult(I arrayList) {
+    if (arrayList == null) {
+      return new LongArrayList();
+    }
     return new LongArrayList(arrayList);
   }
 }
diff --git a/...va/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggStringFunction.java b/...va/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggStringFunction.java
@@ -94,6 +94,9 @@ public I merge(I intermediateResult1, I intermediateResult2) {
 
   @Override
   public ObjectArrayList<String> extractFinalResult(I stringArrayList) {
+    if (stringArrayList == null) {
+      return new ObjectArrayList<>();
+    }
     // NOTE: Wrap a String[] to work around the bug of ObjectArrayList constructor creating Object[] internally.
     String[] stringArray = new String[stringArrayList.size()];
     ObjectIterators.unwrap(stringArrayList.iterator(), stringArray);

diff --git a/...src/main/java/org/apache/pinot/core/query/aggregation/function/array/ListAggFunction.java b/...src/main/java/org/apache/pinot/core/query/aggregation/function/array/ListAggFunction.java
@@ -152,6 +152,9 @@ public DataSchema.ColumnDataType getFinalResultColumnType() {
 
   @Override
   public String extractFinalResult(AbstractObjectCollection<String> strings) {
+    if (strings == null) {
+      return null;
+    }
     return StringUtils.join(strings, _separator);
   }
 }
diff --git a/...-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ArrayTest.java b/...-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ArrayTest.java
@@ -54,6 +54,30 @@ protected long getCountStarResult() {
     return 1000;
   }
 
+  @Test(dataProvider = "useBothQueryEngines")
+  public void testArrayAggWithEmptyPredicate(boolean useMultiStageQueryEngine) throws Exception {
+    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+    String query =
+        String.format("SELECT "
+            + "arrayAgg(boolCol, 'BOOLEAN'), "
+            + "arrayAgg(intCol, 'INT'), "
+            + "arrayAgg(longCol, 'LONG'), "
+            // NOTE: FLOAT array is auto converted to DOUBLE array
+            + (useMultiStageQueryEngine ? "arrayAgg(floatCol, 'DOUBLE'), " : "arrayAgg(floatCol, 'FLOAT'), ")
+            + "arrayAgg(doubleCol, 'DOUBLE'), "
+            + "arrayAgg(stringCol, 'STRING'), "
+            + "arrayAgg(timestampCol, 'TIMESTAMP') "
+            + "FROM %s WHERE intCol < 0 LIMIT %d", getTableName(), getCountStarResult());
+    JsonNode jsonNode = postQuery(query);
+    JsonNode rows = jsonNode.get("resultTable").get("rows");
+    assertEquals(rows.size(), 1);
+    JsonNode row = rows.get(0);
+    assertEquals(row.size(), 7);
+    for (int i = 0; i < 7; i++) {
+      assertEquals(row.get(i).size(), 0);
+    }
+  }
+
   @Test(dataProvider = "useBothQueryEngines")
   public void testArrayAggQueries(boolean useMultiStageQueryEngine)
       throws Exception {