If responseContext keys are null.....

Author: GWphua

processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryMetrics.java

@@ -42,7 +42,7 @@ public interface GroupByQueryMetrics extends QueryMetrics<GroupByQuery>
   void numMetrics(GroupByQuery query);
 
   /**
-   * Sets the number of "complex" metrics of the given groupBy query as dimension. By default it is assumed that
+   * Sets the number of "complex" metrics of the given groupBy query as dimension. By default, it is assumed that
    * "complex" metric is a metric of not long or double type, but it could be redefined in the implementation of this
    * method.
    */

processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryQueryToolChest.java

@@ -189,9 +189,7 @@ private Sequence<ResultRow> initAndMergeGroupByResults(
 
       if (reportMetricsForEmission) {
         GroupByQueryMetrics queryMetrics = (GroupByQueryMetrics) queryPlus.getQueryMetrics();
-        queryMetrics.bytesSpilledToStorage((Long) context.get(GroupByResponseContextKeys.GROUPBY_BYTES_SPILLED_TO_STORAGE_KEY));
-        queryMetrics.mergeDictionarySize((Long) context.get(GroupByResponseContextKeys.GROUPBY_MERGE_DICTIONARY_SIZE_KEY));
-        queryMetrics.mergeBufferAcquisitionTime((Long) context.get(GroupByResponseContextKeys.GROUPBY_MERGE_BUFFER_ACQUISITION_TIME_KEY));
+        populateQueryMetrics(queryMetrics, context);
         groupByStatsProvider.aggregateStats(queryMetrics);
       }
 
@@ -207,6 +205,27 @@ private Sequence<ResultRow> initAndMergeGroupByResults(
     }
   }
 
+  private void populateQueryMetrics(GroupByQueryMetrics queryMetrics, ResponseContext context)
+  {
+    Object bytesSpilledToStorage = context.get(GroupByResponseContextKeys.GROUPBY_BYTES_SPILLED_TO_STORAGE_KEY);
+
+    if (bytesSpilledToStorage != null) {
+      queryMetrics.bytesSpilledToStorage((Long) bytesSpilledToStorage);
+    }
+
+    Object mergeDictionarySize = context.get(GroupByResponseContextKeys.GROUPBY_MERGE_DICTIONARY_SIZE_KEY);
+
+    if (mergeDictionarySize != null) {
+      queryMetrics.mergeDictionarySize((Long) mergeDictionarySize);
+    }
+
+    Object mergeBufferAcquisitionTime = context.get(GroupByResponseContextKeys.GROUPBY_MERGE_BUFFER_ACQUISITION_TIME_KEY);
+
+    if (mergeBufferAcquisitionTime != null) {
+      queryMetrics.mergeBufferAcquisitionTime((Long) mergeBufferAcquisitionTime);
+    }
+  }
+
   private Sequence<ResultRow> mergeGroupByResults(
       final GroupByQuery query,
       GroupByQueryResources resource,

processing/src/main/java/org/apache/druid/query/groupby/GroupByResponseContextKeys.java

@@ -21,6 +21,8 @@
 
 import org.apache.druid.query.context.ResponseContext;
 
+import javax.annotation.Nullable;
+
 /**
  * Response context keys for GroupBy query metrics.
  * These keys are used to aggregate metrics from parallel query execution threads
@@ -32,6 +34,20 @@ public class GroupByResponseContextKeys
   public static final String GROUPBY_BYTES_SPILLED_TO_STORAGE_NAME = "groupByBytesSpilledToStorage";
   public static final String GROUPBY_MERGE_DICTIONARY_SIZE_NAME = "groupByMergeDictionarySize";
 
+  private static Object mergeMax(@Nullable Object oldValue, @Nullable Object newValue)
+  {
+    if (oldValue == null && newValue == null) {
+      return 0L;
+    }
+
+    if (oldValue == null) {
+      return newValue;
+    } else if (newValue == null) {
+      return oldValue;
+    }
+
+    return Math.max((Long) oldValue, (Long) newValue);
+  }
   /**
    * Maximum bytes spilled to storage across all parallel threads processing segments.
    * This represents the peak disk usage during query execution.
@@ -42,7 +58,7 @@ public class GroupByResponseContextKeys
     @Override
     public Object mergeValues(Object oldValue, Object newValue)
     {
-      return Math.max((Long) oldValue, (Long) newValue);
+      return mergeMax(oldValue, newValue);
     }
   };
 
@@ -56,7 +72,7 @@ public Object mergeValues(Object oldValue, Object newValue)
     @Override
     public Object mergeValues(Object oldValue, Object newValue)
     {
-      return Math.max((Long) oldValue, (Long) newValue);
+      return mergeMax(oldValue, newValue);
     }
   };
 
@@ -70,7 +86,7 @@ public Object mergeValues(Object oldValue, Object newValue)
     @Override
     public Object mergeValues(Object oldValue, Object newValue)
     {
-      return Math.max((Long) oldValue, (Long) newValue);
+      return mergeMax(oldValue, newValue);
     }
   };