Reduce allocations due to Jackson serialization.

core/src/main/java/org/apache/druid/java/util/common/jackson/JacksonUtils.java

@@ -19,9 +19,14 @@
 
 package org.apache.druid.java.util.common.jackson;
 
+import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.JsonMappingException;
+import com.fasterxml.jackson.databind.JsonSerializer;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
 
+import javax.annotation.Nullable;
 import java.io.IOException;
 import java.util.Map;
 
@@ -40,7 +45,9 @@ public final class JacksonUtils
       {
       };
 
-  /** Silences Jackson's {@link IOException}. */
+  /**
+   * Silences Jackson's {@link IOException}.
+   */
   public static <T> T readValue(ObjectMapper mapper, byte[] bytes, Class<T> valueClass)
   {
     try {
@@ -51,6 +58,37 @@ public static <T> T readValue(ObjectMapper mapper, byte[] bytes, Class<T> valueC
     }
   }
 
+  /**
+   * Returns a serializer for a particular class. If you have a {@link SerializerProvider}, this is better than calling
+   * {@link JsonGenerator#writeObject(Object)} or {@link ObjectMapper#writeValue(JsonGenerator, Object)}, because it
+   * avoids re-creating the {@link SerializerProvider} for each serialized object.
+   */
+  public static JsonSerializer<Object> getSerializer(final SerializerProvider serializerProvider, final Class<?> clazz)
+      throws JsonMappingException
+  {
+    // cache = true, property = null because this is what DefaultSerializerProvider.serializeValue would do.
+    return serializerProvider.findTypedValueSerializer(clazz, true, null);
+  }
+
+  /**
+   * Serializes an object using a {@link JsonGenerator}. If you have a {@link SerializerProvider}, this is better than
+   * calling {@link JsonGenerator#writeObject(Object)}, because it avoids re-creating the {@link SerializerProvider}
+   * for each serialized object.
+   */
+  public static void writeObjectUsingSerializerProvider(
+      final JsonGenerator jsonGenerator,
+      final SerializerProvider serializers,
+      @Nullable final Object o
+  ) throws IOException
+  {
+    if (o == null) {
+      jsonGenerator.writeNull();
+    } else {
+      final JsonSerializer<Object> serializer = getSerializer(serializers, o.getClass());
+      serializer.serialize(o, jsonGenerator, serializers);
+    }
+  }
+
   private JacksonUtils()
   {
   }

docs/querying/groupbyquery.md

@@ -409,6 +409,7 @@ Supported runtime properties:
 |--------|-----------|-------|
 |`druid.query.groupBy.defaultStrategy`|Default groupBy query strategy.|v2|
 |`druid.query.groupBy.singleThreaded`|Merge results using a single thread.|false|
+|`druid.query.groupBy.intermediateResultAsMapCompat`|Whether Brokers are able to understand map-based result rows. Setting this to `true` adds some overhead to all groupBy queries. It is required for compatibility with data servers running versions older than 0.16.0, which introduced [array-based result rows](#array-based-result-rows).|false|
 
 Supported query contexts:
 

processing/src/main/java/org/apache/druid/jackson/DruidDefaultSerializersModule.java

@@ -31,6 +31,7 @@
 import org.apache.druid.java.util.common.guava.Accumulator;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Yielder;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.query.context.ResponseContext;
 import org.apache.druid.query.context.ResponseContextDeserializer;
 import org.joda.time.DateTimeZone;
@@ -93,11 +94,26 @@ public void serialize(Sequence value, final JsonGenerator jgen, SerializerProvid
                 null,
                 new Accumulator<Object, Object>()
                 {
+                  // Save allocations in jgen.writeObject by caching serializer.
+                  JsonSerializer<Object> serializer = null;
+                  Class<?> serializerClass = null;
+
                   @Override
-                  public Object accumulate(Object o, Object o1)
+                  public Object accumulate(Object ignored, Object object)
                   {
                     try {
-                      jgen.writeObject(o1);
+                      if (object == null) {
+                        jgen.writeNull();
+                      } else {
+                        final Class<?> clazz = object.getClass();
+
+                        if (serializerClass != clazz) {
+                          serializer = JacksonUtils.getSerializer(provider, clazz);
+                          serializerClass = clazz;
+                        }
+
+                        serializer.serialize(object, jgen, provider);
+                      }
                     }
                     catch (IOException e) {
                       throw new RuntimeException(e);
@@ -119,11 +135,27 @@ public Object accumulate(Object o, Object o1)
           public void serialize(Yielder yielder, final JsonGenerator jgen, SerializerProvider provider)
               throws IOException
           {
+            // Save allocations in jgen.writeObject by caching serializer.
+            JsonSerializer<Object> serializer = null;
+            Class<?> serializerClass = null;
+
             try {
               jgen.writeStartArray();
               while (!yielder.isDone()) {
                 final Object o = yielder.get();
-                jgen.writeObject(o);
+                if (o == null) {
+                  jgen.writeNull();
+                } else {
+                  final Class<?> clazz = o.getClass();
+
+                  if (serializerClass != clazz) {
+                    serializer = JacksonUtils.getSerializer(provider, clazz);
+                    serializerClass = clazz;
+                  }
+
+                  serializer.serialize(o, jgen, provider);
+                }
+
                 yielder = yielder.next(null);
               }
               jgen.writeEndArray();

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

@@ -105,6 +105,9 @@ public class GroupByQueryConfig
   @JsonProperty
   private boolean vectorize = true;
 
+  @JsonProperty
+  private boolean intermediateResultAsMapCompat = false;
+
   @JsonProperty
   private boolean enableMultiValueUnnesting = true;
 
@@ -203,6 +206,11 @@ public boolean isVectorize()
     return vectorize;
   }
 
+  public boolean isIntermediateResultAsMapCompat()
+  {
+    return intermediateResultAsMapCompat;
+  }
+
   public boolean isForcePushDownNestedQuery()
   {
     return forcePushDownNestedQuery;

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

@@ -31,6 +31,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Functions;
+import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.inject.Inject;
@@ -40,6 +41,7 @@
 import org.apache.druid.java.util.common.guava.MappedSequence;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.query.CacheStrategy;
 import org.apache.druid.query.DataSource;
 import org.apache.druid.query.Query;
@@ -91,21 +93,24 @@ public class GroupByQueryQueryToolChest extends QueryToolChest<ResultRow, GroupB
   public static final String GROUP_BY_MERGE_KEY = "groupByMerge";
 
   private final GroupByStrategySelector strategySelector;
+  private final GroupByQueryConfig queryConfig;
   private final GroupByQueryMetricsFactory queryMetricsFactory;
 
   @VisibleForTesting
   public GroupByQueryQueryToolChest(GroupByStrategySelector strategySelector)
   {
-    this(strategySelector, DefaultGroupByQueryMetricsFactory.instance());
+    this(strategySelector, GroupByQueryConfig::new, DefaultGroupByQueryMetricsFactory.instance());
   }
 
   @Inject
   public GroupByQueryQueryToolChest(
       GroupByStrategySelector strategySelector,
+      Supplier<GroupByQueryConfig> queryConfigSupplier,
       GroupByQueryMetricsFactory queryMetricsFactory
   )
   {
     this.strategySelector = strategySelector;
+    this.queryConfig = queryConfigSupplier.get();
     this.queryMetricsFactory = queryMetricsFactory;
   }
 
@@ -415,6 +420,12 @@ public ObjectMapper decorateObjectMapper(final ObjectMapper objectMapper, final
   {
     final boolean resultAsArray = query.getContextBoolean(GroupByQueryConfig.CTX_KEY_ARRAY_RESULT_ROWS, false);
 
+    if (resultAsArray && !queryConfig.isIntermediateResultAsMapCompat()) {
+      // We can assume ResultRow are serialized and deserialized as arrays. No need for special decoration,
+      // and we can save the overhead of making a copy of the ObjectMapper.
+      return objectMapper;
+    }
+
     // Serializer that writes array- or map-based rows as appropriate, based on the "resultAsArray" setting.
     final JsonSerializer<ResultRow> serializer = new JsonSerializer<ResultRow>()
     {
@@ -426,9 +437,9 @@ public void serialize(
       ) throws IOException
       {
         if (resultAsArray) {
-          jg.writeObject(resultRow.getArray());
+          JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, resultRow.getArray());
         } else {
-          jg.writeObject(resultRow.toMapBasedRow(query));
+          JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, resultRow.toMapBasedRow(query));
         }
       }
     };

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/SpillingGrouper.java

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.databind.MappingIterator;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
@@ -30,6 +31,7 @@
 import net.jpountz.lz4.LZ4BlockOutputStream;
 import org.apache.druid.java.util.common.CloseableIterators;
 import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.java.util.common.parsers.CloseableIterator;
 import org.apache.druid.query.BaseQuery;
@@ -309,10 +311,11 @@ private <T> File spill(Iterator<T> iterator) throws IOException
         final LZ4BlockOutputStream compressedOut = new LZ4BlockOutputStream(out);
         final JsonGenerator jsonGenerator = spillMapper.getFactory().createGenerator(compressedOut)
     ) {
+      final SerializerProvider serializers = spillMapper.getSerializerProviderInstance();
+
       while (iterator.hasNext()) {
         BaseQuery.checkInterrupted();
-
-        jsonGenerator.writeObject(iterator.next());
+        JacksonUtils.writeObjectUsingSerializerProvider(jsonGenerator, serializers, iterator.next());
       }
 
       return out.getFile();

processing/src/test/java/org/apache/druid/jackson/DefaultObjectMapperTest.java

@@ -22,11 +22,18 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.guava.Sequence;
+import org.apache.druid.java.util.common.guava.Sequences;
+import org.apache.druid.java.util.common.guava.Yielders;
 import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.Arrays;
+
 /**
+ *
  */
 public class DefaultObjectMapperTest
 {
@@ -39,4 +46,25 @@ public void testDateTime() throws Exception
 
     Assert.assertEquals(StringUtils.format("\"%s\"", time), mapper.writeValueAsString(time));
   }
+
+  @Test
+  public void testYielder() throws Exception
+  {
+    final Sequence<Object> sequence = Sequences.simple(
+        Arrays.asList(
+            "a",
+            "b",
+            null,
+            DateTimes.utc(2L),
+            5,
+            DateTimeZone.UTC,
+            "c"
+        )
+    );
+
+    Assert.assertEquals(
+        "[\"a\",\"b\",null,\"1970-01-01T00:00:00.002Z\",5,\"UTC\",\"c\"]",
+        mapper.writeValueAsString(Yielders.each(sequence))
+    );
+  }
 }