[Backport] add protobuf flattener, direct to plain java conversion for faster flattening

core/src/main/java/org/apache/druid/java/util/common/parsers/FlattenerJsonProvider.java

@@ -0,0 +1,138 @@
+/*
+ * 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.java.util.common.parsers;
+
+import com.jayway.jsonpath.InvalidJsonException;
+import com.jayway.jsonpath.spi.json.JsonProvider;
+
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+public abstract class FlattenerJsonProvider implements JsonProvider
+{
+  @Override
+  public Object createArray()
+  {
+    return new ArrayList<>();
+  }
+
+  @Override
+  public Object createMap()
+  {
+    return new LinkedHashMap<>();
+  }
+
+  @Override
+  public int length(final Object o)
+  {
+    if (o instanceof List) {
+      return ((List) o).size();
+    } else {
+      return 0;
+    }
+  }
+
+  @Override
+  public Iterable<?> toIterable(final Object o)
+  {
+    if (o instanceof List) {
+      return (List) o;
+    }
+    throw new UnsupportedOperationException(o.getClass().getName());
+  }
+
+  @Override
+  public Object getArrayIndex(final Object o, final int i)
+  {
+    if (o instanceof List) {
+      return ((List) o).get(i);
+    }
+    throw new UnsupportedOperationException(o.getClass().getName());
+  }
+
+  @Override
+  public void setArrayIndex(final Object o, final int i, final Object o1)
+  {
+    if (o instanceof List) {
+      final List list = (List) o;
+      if (list.size() == i) {
+        list.add(o1);
+      } else {
+        list.set(i, o1);
+      }
+    } else {
+      throw new UnsupportedOperationException(o.getClass().getName());
+    }
+  }
+
+  @Override
+  public void setProperty(final Object o, final Object o1, final Object o2)
+  {
+    if (o instanceof Map) {
+      ((Map) o).put(o1, o2);
+    } else {
+      throw new UnsupportedOperationException(o.getClass().getName());
+    }
+  }
+
+  @Override
+  public void removeProperty(final Object o, final Object o1)
+  {
+    if (o instanceof Map) {
+      ((Map) o).remove(o1);
+    } else {
+      throw new UnsupportedOperationException(o.getClass().getName());
+    }
+  }
+
+  @Override
+  @Deprecated
+  public Object getArrayIndex(final Object o, final int i, final boolean b)
+  {
+    throw new UnsupportedOperationException("Deprecated");
+  }
+
+  @Override
+  public Object parse(final String s) throws InvalidJsonException
+  {
+    throw new UnsupportedOperationException("Unused");
+  }
+
+  @Override
+  public Object parse(final InputStream inputStream, final String s) throws InvalidJsonException
+  {
+    throw new UnsupportedOperationException("Unused");
+  }
+
+  @Override
+  public String toJson(final Object o)
+  {
+    throw new UnsupportedOperationException("Unused");
+  }
+
+  @Override
+  public Object unwrap(final Object o)
+  {
+    return o;
+  }
+}

core/src/main/java/org/apache/druid/java/util/common/parsers/JSONFlattenerMaker.java

@@ -54,9 +54,9 @@ public class JSONFlattenerMaker implements ObjectFlatteners.FlattenerMaker<JsonN
                    .options(EnumSet.of(Option.SUPPRESS_EXCEPTIONS))
                    .build();
 
+  private final CharsetEncoder enc = StandardCharsets.UTF_8.newEncoder();
   private final boolean keepNullValues;
 
-  private final CharsetEncoder enc = StandardCharsets.UTF_8.newEncoder();
 
   public JSONFlattenerMaker(boolean keepNullValues)
   {
@@ -66,7 +66,7 @@ public JSONFlattenerMaker(boolean keepNullValues)
   @Override
   public Iterable<String> discoverRootFields(final JsonNode obj)
   {
-    return FluentIterable.from(() -> obj.fields())
+    return FluentIterable.from(obj::fields)
                          .filter(
                              entry -> {
                                final JsonNode val = entry.getValue();
@@ -137,13 +137,13 @@ public JsonProvider getJsonProvider()
   public Object finalizeConversionForMap(Object o)
   {
     if (o instanceof JsonNode) {
-      return convertJsonNode((JsonNode) o);
+      return convertJsonNode((JsonNode) o, enc);
     }
     return o;
   }
 
   @Nullable
-  private Object convertJsonNode(JsonNode val)
+  public static Object convertJsonNode(JsonNode val, CharsetEncoder enc)
   {
     if (val == null || val.isNull()) {
       return null;
@@ -158,7 +158,7 @@ private Object convertJsonNode(JsonNode val)
     }
 
     if (val.isTextual()) {
-      return charsetFix(val.asText());
+      return charsetFix(val.asText(), enc);
     }
 
     if (val.isBoolean()) {
@@ -175,7 +175,7 @@ private Object convertJsonNode(JsonNode val)
       List<Object> newList = new ArrayList<>();
       for (JsonNode entry : val) {
         if (!entry.isNull()) {
-          newList.add(finalizeConversionForMap(entry));
+          newList.add(convertJsonNode(entry, enc));
         }
       }
       return newList;
@@ -185,7 +185,7 @@ private Object convertJsonNode(JsonNode val)
       Map<String, Object> newMap = new LinkedHashMap<>();
       for (Iterator<Map.Entry<String, JsonNode>> it = val.fields(); it.hasNext(); ) {
         Map.Entry<String, JsonNode> entry = it.next();
-        newMap.put(entry.getKey(), finalizeConversionForMap(entry.getValue()));
+        newMap.put(entry.getKey(), convertJsonNode(entry.getValue(), enc));
       }
       return newMap;
     }
@@ -197,7 +197,7 @@ private Object convertJsonNode(JsonNode val)
   }
 
   @Nullable
-  private String charsetFix(String s)
+  private static String charsetFix(String s, CharsetEncoder enc)
   {
     if (s != null && !enc.canEncode(s)) {
       // Some whacky characters are in this string (e.g. \uD900). These are problematic because they are decodeable
@@ -209,7 +209,7 @@ private String charsetFix(String s)
     }
   }
 
-  private boolean isFlatList(JsonNode list)
+  private static boolean isFlatList(JsonNode list)
   {
     for (JsonNode obj : list) {
       if (obj.isObject() || obj.isArray()) {

core/src/test/java/org/apache/druid/java/util/common/parsers/FlattenerJsonProviderTest.java

@@ -0,0 +1,146 @@
+/*
+ * 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.java.util.common.parsers;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.apache.druid.java.util.common.StringUtils;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+
+public class FlattenerJsonProviderTest
+{
+  FlattenerJsonProvider jsonProvider = new FlattenerJsonProvider()
+  {
+    @Override
+    public boolean isArray(final Object o)
+    {
+      throw new RuntimeException("not tested");
+    }
+
+    @Override
+    public boolean isMap(final Object o)
+    {
+      throw new RuntimeException("not tested");
+    }
+
+    @Override
+    public Collection<String> getPropertyKeys(final Object o)
+    {
+      throw new RuntimeException("not tested");
+    }
+
+    @Override
+    public Object getMapValue(final Object o, final String s)
+    {
+      throw new RuntimeException("not tested");
+    }
+  };
+
+  @Test
+  public void testMapStuff()
+  {
+    Object aMap = jsonProvider.createMap();
+    jsonProvider.setProperty(aMap, "key", "value");
+    Assert.assertEquals(ImmutableMap.of("key", "value"), aMap);
+    jsonProvider.removeProperty(aMap, "key");
+    Assert.assertEquals(ImmutableMap.of(), aMap);
+    Assert.assertEquals(aMap, jsonProvider.unwrap(aMap));
+
+    Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.setProperty(jsonProvider.createArray(), "key", "value")
+    );
+    Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.removeProperty(jsonProvider.createArray(), "key")
+    );
+  }
+
+  @Test
+  public void testArrayStuff()
+  {
+    Object aList = jsonProvider.createArray();
+    jsonProvider.setArrayIndex(aList, 0, "a");
+    jsonProvider.setArrayIndex(aList, 1, "b");
+    jsonProvider.setArrayIndex(aList, 2, "c");
+    Assert.assertEquals(3, jsonProvider.length(aList));
+    Assert.assertEquals("a", jsonProvider.getArrayIndex(aList, 0));
+    Assert.assertEquals("b", jsonProvider.getArrayIndex(aList, 1));
+    Assert.assertEquals("c", jsonProvider.getArrayIndex(aList, 2));
+    List<String> expected = ImmutableList.of("a", "b", "c");
+    Assert.assertEquals(expected, aList);
+    Iterator<?> iter = jsonProvider.toIterable(aList).iterator();
+    Iterator<String> expectedIter = expected.iterator();
+    while (iter.hasNext()) {
+      Assert.assertEquals(expectedIter.next(), iter.next());
+    }
+    Assert.assertFalse(expectedIter.hasNext());
+    Assert.assertEquals(aList, jsonProvider.unwrap(aList));
+
+    Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.getArrayIndex(jsonProvider.createMap(), 0)
+    );
+    Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.setArrayIndex(jsonProvider.createMap(), 0, "a")
+    );
+    Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.toIterable(jsonProvider.createMap())
+    );
+  }
+
+  @Test
+  public void testNotImplementedOnPurpose()
+  {
+    Object aList = jsonProvider.createArray();
+    Throwable t = Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.toJson(aList)
+    );
+    Assert.assertEquals("Unused", t.getMessage());
+
+    t = Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.parse("{}")
+    );
+    Assert.assertEquals("Unused", t.getMessage());
+
+
+    t = Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.parse(new ByteArrayInputStream(StringUtils.toUtf8("{}")), "UTF-8")
+    );
+    Assert.assertEquals("Unused", t.getMessage());
+
+    t = Assert.assertThrows(
+        UnsupportedOperationException.class,
+        () -> jsonProvider.getArrayIndex(aList, 0, false)
+    );
+    Assert.assertEquals("Deprecated", t.getMessage());
+  }
+}