Reduce allocations due to Jackson serialization.

codestyle/druid-forbidden-apis.txt

@@ -1,6 +1,8 @@
 com.fasterxml.jackson.databind.ObjectMapper#reader(com.fasterxml.jackson.core.type.TypeReference) @ Use ObjectMapper#readerFor instead
 com.fasterxml.jackson.databind.ObjectMapper#reader(com.fasterxml.jackson.databind.JavaType) @ Use ObjectMapper#readerFor instead
 com.fasterxml.jackson.databind.ObjectMapper#reader(java.lang.Class) @ Use ObjectMapper#readerFor instead
+com.fasterxml.jackson.databind.ObjectMapper#writeValue(com.fasterxml.jackson.core.JsonGenerator, java.lang.Object) @ Use JacksonUtils#writeObjectUsingSerializerProvider to allow SerializerProvider reuse
+com.fasterxml.jackson.core.JsonGenerator#writeObject(java.lang.Object) @ Use JacksonUtils#writeObjectUsingSerializerProvider to allow SerializerProvider reuse
 com.google.common.base.Charsets @ Use java.nio.charset.StandardCharsets instead
 com.google.common.collect.Iterators#emptyIterator() @ Use java.util.Collections#emptyIterator()
 com.google.common.collect.Lists#newArrayList() @ Create java.util.ArrayList directly

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()
   {
   }

core/src/test/java/org/apache/druid/common/jackson/JacksonUtilsTest.java

@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.common.jackson;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+
+public class JacksonUtilsTest
+{
+  @Test
+  public void testWriteObjectUsingSerializerProvider() throws IOException
+  {
+    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+
+    final ObjectMapper objectMapper = new ObjectMapper();
+    final SerializerProvider serializers = objectMapper.getSerializerProviderInstance();
+
+    final JsonGenerator jg = objectMapper.getFactory().createGenerator(baos);
+    jg.writeStartArray();
+    JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, new SerializableClass(2));
+    JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, null);
+    JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, new SerializableClass(3));
+    jg.writeEndArray();
+    jg.close();
+
+    final List<SerializableClass> deserializedValues = objectMapper.readValue(
+        baos.toByteArray(),
+        new TypeReference<List<SerializableClass>>() {}
+    );
+
+    Assert.assertEquals(
+        Arrays.asList(new SerializableClass(2), null, new SerializableClass(3)),
+        deserializedValues
+    );
+  }
+
+  @Test
+  public void testWritePrimitivesUsingSerializerProvider() throws IOException
+  {
+    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+
+    final ObjectMapper objectMapper = new ObjectMapper();
+    final SerializerProvider serializers = objectMapper.getSerializerProviderInstance();
+
+    final JsonGenerator jg = objectMapper.getFactory().createGenerator(baos);
+    jg.writeStartArray();
+    JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, "foo");
+    JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, null);
+    JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, 1.23);
+    jg.writeEndArray();
+    jg.close();
+
+    final List<Object> deserializedValues = objectMapper.readValue(
+        baos.toByteArray(),
+        new TypeReference<List<Object>>() {}
+    );
+
+    Assert.assertEquals(
+        Arrays.asList("foo", null, 1.23),
+        deserializedValues
+    );
+  }
+
+  public static class SerializableClass
+  {
+    private final int value;
+
+    @JsonCreator
+    public SerializableClass(@JsonProperty("value") final int value)
+    {
+      this.value = value;
+    }
+
+    @JsonProperty
+    public int getValue()
+    {
+      return value;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      SerializableClass that = (SerializableClass) o;
+      return value == that.value;
+    }
+
+    @Override
+    public int hashCode()
+    {
+      return Objects.hash(value);
+    }
+
+    @Override
+    public String toString()
+    {
+      return "SerializableClass{" +
+             "value=" + value +
+             '}';
+    }
+  }
+}

core/src/test/java/org/apache/druid/data/input/impl/prefetch/JsonIteratorTest.java

@@ -23,9 +23,11 @@
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
 import com.fasterxml.jackson.dataformat.smile.SmileFactory;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -50,9 +52,10 @@ public void testSerde() throws IOException
     };
     try (FileOutputStream fos = new FileOutputStream(testFile)) {
       final JsonGenerator jg = mapper.getFactory().createGenerator(fos);
+      final SerializerProvider serializers = mapper.getSerializerProviderInstance();
       jg.writeStartArray();
       for (Map<String, Object> mapFromList : expectedList) {
-        jg.writeObject(mapFromList);
+        JacksonUtils.writeObjectUsingSerializerProvider(jg, serializers, mapFromList);
       }
       jg.writeEndArray();
       jg.close();

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));
         }
       }
     };