fix Update json_extract to Produce Canonicalized Output

presto-common/src/main/java/com/facebook/presto/common/function/SqlFunctionProperties.java

@@ -37,6 +37,7 @@ public class SqlFunctionProperties
     private final boolean legacyJsonCast;
     private final Map<String, String> extraCredentials;
     private final boolean warnOnCommonNanPatterns;
+    private final boolean canonicalizedJsonExtract;
 
     private SqlFunctionProperties(
             boolean parseDecimalLiteralAsDouble,
@@ -50,7 +51,8 @@ private SqlFunctionProperties(
             boolean fieldNamesInJsonCastEnabled,
             boolean legacyJsonCast,
             Map<String, String> extraCredentials,
-            boolean warnOnCommonNanPatterns)
+            boolean warnOnCommonNanPatterns,
+            boolean canonicalizedJsonExtract)
     {
         this.parseDecimalLiteralAsDouble = parseDecimalLiteralAsDouble;
         this.legacyRowFieldOrdinalAccessEnabled = legacyRowFieldOrdinalAccessEnabled;
@@ -64,6 +66,7 @@ private SqlFunctionProperties(
         this.legacyJsonCast = legacyJsonCast;
         this.extraCredentials = requireNonNull(extraCredentials, "extraCredentials is null");
         this.warnOnCommonNanPatterns = warnOnCommonNanPatterns;
+        this.canonicalizedJsonExtract = canonicalizedJsonExtract;
     }
 
     public boolean isParseDecimalLiteralAsDouble()
@@ -127,6 +130,9 @@ public boolean shouldWarnOnCommonNanPatterns()
         return warnOnCommonNanPatterns;
     }
 
+    public boolean isCanonicalizedJsonExtract()
+    { return canonicalizedJsonExtract; }
+
     @Override
     public boolean equals(Object o)
     {
@@ -146,15 +152,16 @@ public boolean equals(Object o)
                 Objects.equals(sessionLocale, that.sessionLocale) &&
                 Objects.equals(sessionUser, that.sessionUser) &&
                 Objects.equals(extraCredentials, that.extraCredentials) &&
-                Objects.equals(legacyJsonCast, that.legacyJsonCast);
+                Objects.equals(legacyJsonCast, that.legacyJsonCast) &&
+                Objects.equals(canonicalizedJsonExtract, that.legacyJsonCast);
     }
 
     @Override
     public int hashCode()
     {
         return Objects.hash(parseDecimalLiteralAsDouble, legacyRowFieldOrdinalAccessEnabled, timeZoneKey,
                 legacyTimestamp, legacyMapSubscript, sessionStartTime, sessionLocale, sessionUser,
-                extraCredentials, legacyJsonCast);
+                extraCredentials, legacyJsonCast, canonicalizedJsonExtract);
     }
 
     public static Builder builder()
@@ -176,6 +183,7 @@ public static class Builder
         private boolean legacyJsonCast;
         private Map<String, String> extraCredentials = emptyMap();
         private boolean warnOnCommonNanPatterns;
+        private boolean canonicalizedJsonExtract;
 
         private Builder() {}
 
@@ -251,6 +259,12 @@ public Builder setWarnOnCommonNanPatterns(boolean warnOnCommonNanPatterns)
             return this;
         }
 
+        public Builder setCanonicalizedJsonExtract(boolean canonicalizedJsonExtract)
+        {
+            this.canonicalizedJsonExtract = canonicalizedJsonExtract;
+            return this;
+        }
+
         public SqlFunctionProperties build()
         {
             return new SqlFunctionProperties(
@@ -265,7 +279,8 @@ public SqlFunctionProperties build()
                     fieldNamesInJsonCastEnabled,
                     legacyJsonCast,
                     extraCredentials,
-                    warnOnCommonNanPatterns);
+                    warnOnCommonNanPatterns,
+                    canonicalizedJsonExtract);
         }
     }
 }

presto-main/src/main/java/com/facebook/presto/Session.java

@@ -57,6 +57,7 @@
 import java.util.stream.Collectors;
 
 import static com.facebook.presto.SystemSessionProperties.LEGACY_JSON_CAST;
+import static com.facebook.presto.SystemSessionProperties.isCanonicalizedJsonExtract;
 import static com.facebook.presto.SystemSessionProperties.isFieldNameInJsonCastEnabled;
 import static com.facebook.presto.SystemSessionProperties.isLegacyMapSubscript;
 import static com.facebook.presto.SystemSessionProperties.isLegacyRowFieldOrdinalAccessEnabled;
@@ -481,6 +482,7 @@ public SqlFunctionProperties getSqlFunctionProperties()
                 .setLegacyJsonCast(legacyJsonCast)
                 .setExtraCredentials(identity.getExtraCredentials())
                 .setWarnOnCommonNanPatterns(warnOnCommonNanPatterns(this))
+                .setCanonicalizedJsonExtract(isCanonicalizedJsonExtract(this))
                 .build();
     }
 

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -305,6 +305,7 @@ public final class SystemSessionProperties
     public static final String REWRITE_CASE_TO_MAP_ENABLED = "rewrite_case_to_map_enabled";
     public static final String FIELD_NAMES_IN_JSON_CAST_ENABLED = "field_names_in_json_cast_enabled";
     public static final String LEGACY_JSON_CAST = "legacy_json_cast";
+    public static final String CANONICALIZED_JSON_EXTRACT = "canonicalized_json_extract";
     public static final String PULL_EXPRESSION_FROM_LAMBDA_ENABLED = "pull_expression_from_lambda_enabled";
     public static final String REWRITE_CONSTANT_ARRAY_CONTAINS_TO_IN_EXPRESSION = "rewrite_constant_array_contains_to_in_expression";
     public static final String INFER_INEQUALITY_PREDICATES = "infer_inequality_predicates";
@@ -1636,6 +1637,11 @@ public SystemSessionProperties(
                         "Keep the legacy json cast behavior, do not reserve the case for field names when casting to row type",
                         functionsConfig.isLegacyJsonCast(),
                         true),
+                booleanProperty(
+                        CANONICALIZED_JSON_EXTRACT,
+                        "Extracts json data in a canonicalized manner, and raises a PrestoException when encountering invalid json structures within the input json path",
+                        functionsConfig.isCanonicalizedJsonExtract(),
+                        true),
                 booleanProperty(
                         OPTIMIZE_JOIN_PROBE_FOR_EMPTY_BUILD_RUNTIME,
                         "Optimize join probe at runtime if build side is empty",
@@ -3174,4 +3180,9 @@ public static boolean isEnabledAddExchangeBelowGroupId(Session session)
     {
         return session.getSystemProperty(ADD_EXCHANGE_BELOW_PARTIAL_AGGREGATION_OVER_GROUP_ID, Boolean.class);
     }
+
+    public static boolean isCanonicalizedJsonExtract(Session session)
+    {
+        return session.getSystemProperty(CANONICALIZED_JSON_EXTRACT, Boolean.class);
+    }
 }

presto-main/src/main/java/com/facebook/presto/operator/scalar/JsonExtract.java

@@ -13,19 +13,23 @@
  */
 package com.facebook.presto.operator.scalar;
 
+import com.facebook.airlift.json.JsonObjectMapperProvider;
+import com.facebook.presto.common.function.SqlFunctionProperties;
 import com.facebook.presto.spi.PrestoException;
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
 import com.fasterxml.jackson.core.io.SerializedString;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
 import io.airlift.slice.DynamicSliceOutput;
 import io.airlift.slice.Slice;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.UncheckedIOException;
 
 import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
@@ -38,6 +42,7 @@
 import static com.fasterxml.jackson.core.JsonToken.START_ARRAY;
 import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
 import static com.fasterxml.jackson.core.JsonToken.VALUE_NULL;
+import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
 import static io.airlift.slice.Slices.utf8Slice;
 import static java.util.Objects.requireNonNull;
 
@@ -121,13 +126,15 @@ public final class JsonExtract
     private static final JsonFactory JSON_FACTORY = new JsonFactory()
             .disable(CANONICALIZE_FIELD_NAMES);
 
+    private static final ObjectMapper SORTED_MAPPER = new JsonObjectMapperProvider().get().configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
+
     private JsonExtract() {}
 
-    public static <T> T extract(Slice jsonInput, JsonExtractor<T> jsonExtractor)
+    public static <T> T extract(Slice jsonInput, JsonExtractor<T> jsonExtractor, SqlFunctionProperties properties)
     {
         requireNonNull(jsonInput, "jsonInput is null");
         try {
-            return jsonExtractor.extract(jsonInput.getInput());
+            return jsonExtractor.extract(jsonInput.getInput(), properties);
         }
         catch (JsonParseException e) {
             // Return null if we failed to parse something
@@ -156,7 +163,7 @@ public static <T> PrestoJsonExtractor<T> generateExtractor(String path, PrestoJs
 
     public interface JsonExtractor<T>
     {
-        T extract(InputStream inputStream)
+        T extract(InputStream inputStream, SqlFunctionProperties properties)
                 throws IOException;
     }
 
@@ -174,11 +181,11 @@ public abstract static class PrestoJsonExtractor<T>
          *
          * @return the value, or null if not applicable
          */
-        abstract T extract(JsonParser jsonParser)
+        abstract T extract(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException;
 
         @Override
-        public T extract(InputStream inputStream)
+        public T extract(InputStream inputStream, SqlFunctionProperties properties)
                 throws IOException
         {
             try (JsonParser jsonParser = createJsonParser(JSON_FACTORY, inputStream)) {
@@ -187,7 +194,7 @@ public T extract(InputStream inputStream)
                     return null;
                 }
 
-                return extract(jsonParser);
+                return extract(jsonParser, properties);
             }
         }
     }
@@ -214,21 +221,21 @@ public ObjectFieldJsonExtractor(String fieldName, PrestoJsonExtractor<? extends
         }
 
         @Override
-        public T extract(JsonParser jsonParser)
+        public T extract(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException
         {
             if (jsonParser.getCurrentToken() == START_OBJECT) {
-                return processJsonObject(jsonParser);
+                return processJsonObject(jsonParser, properties);
             }
 
             if (jsonParser.getCurrentToken() == START_ARRAY) {
-                return processJsonArray(jsonParser);
+                return processJsonArray(jsonParser, properties);
             }
 
             throw new JsonParseException(jsonParser, "Expected a JSON object or array");
         }
 
-        public T processJsonObject(JsonParser jsonParser)
+        public T processJsonObject(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException
         {
             while (!jsonParser.nextFieldName(fieldName)) {
@@ -244,10 +251,10 @@ public T processJsonObject(JsonParser jsonParser)
 
             jsonParser.nextToken(); // Shift to first token of the value
 
-            return delegate.extract(jsonParser);
+            return delegate.extract(jsonParser, properties);
         }
 
-        public T processJsonArray(JsonParser jsonParser)
+        public T processJsonArray(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException
         {
             int currentIndex = 0;
@@ -270,15 +277,15 @@ public T processJsonArray(JsonParser jsonParser)
                 jsonParser.skipChildren(); // Skip nested structure if currently at the start of one
             }
 
-            return delegate.extract(jsonParser);
+            return delegate.extract(jsonParser, properties);
         }
     }
 
     public static class ScalarValueJsonExtractor
             extends PrestoJsonExtractor<Slice>
     {
         @Override
-        public Slice extract(JsonParser jsonParser)
+        public Slice extract(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException
         {
             JsonToken token = jsonParser.getCurrentToken();
@@ -296,13 +303,31 @@ public static class JsonValueJsonExtractor
             extends PrestoJsonExtractor<Slice>
     {
         @Override
-        public Slice extract(JsonParser jsonParser)
+        public Slice extract(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException
         {
             if (!jsonParser.hasCurrentToken()) {
                 throw new JsonParseException(jsonParser, "Unexpected end of value");
             }
+            if (!properties.isCanonicalizedJsonExtract()) {
+                return legacyExtract(jsonParser);
+            }
+            DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(ESTIMATED_JSON_OUTPUT_SIZE);
+            // Write the JSON to output stream with sorted keys
+            SORTED_MAPPER.writeValue((OutputStream) dynamicSliceOutput, SORTED_MAPPER.readValue(jsonParser, Object.class));
+            // nextToken will throw an exception if there are trailing characters.
+            try {
+                jsonParser.nextToken();
+            }
+            catch (JsonParseException e) {
+                throw new PrestoException(INVALID_FUNCTION_ARGUMENT, e.getMessage());
+            }
+            return dynamicSliceOutput.slice();
+        }
 
+        public Slice legacyExtract(JsonParser jsonParser)
+                throws IOException
+        {
             DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(ESTIMATED_JSON_OUTPUT_SIZE);
             try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, dynamicSliceOutput)) {
                 jsonGenerator.copyCurrentStructure(jsonParser);
@@ -315,7 +340,7 @@ public static class JsonSizeExtractor
             extends PrestoJsonExtractor<Long>
     {
         @Override
-        public Long extract(JsonParser jsonParser)
+        public Long extract(JsonParser jsonParser, SqlFunctionProperties properties)
                 throws IOException
         {
             if (!jsonParser.hasCurrentToken()) {

presto-main/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java

@@ -435,51 +435,51 @@ public static Slice jsonArrayGet(@SqlType(StandardTypes.JSON) Slice json, @SqlTy
     @SqlNullable
     @LiteralParameters("x")
     @SqlType("varchar(x)")
-    public static Slice varcharJsonExtractScalar(@SqlType("varchar(x)") Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
+    public static Slice varcharJsonExtractScalar(SqlFunctionProperties properties, @SqlType("varchar(x)") Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
     {
-        return JsonExtract.extract(json, jsonPath.getScalarExtractor());
+        return JsonExtract.extract(json, jsonPath.getScalarExtractor(), properties);
     }
 
     @ScalarFunction
     @SqlNullable
     @SqlType(StandardTypes.VARCHAR)
-    public static Slice jsonExtractScalar(@SqlType(StandardTypes.JSON) Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
+    public static Slice jsonExtractScalar(SqlFunctionProperties properties, @SqlType(StandardTypes.JSON) Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
     {
-        return JsonExtract.extract(json, jsonPath.getScalarExtractor());
+        return JsonExtract.extract(json, jsonPath.getScalarExtractor(), properties);
     }
 
     @ScalarFunction("json_extract")
     @LiteralParameters("x")
     @SqlNullable
     @SqlType(StandardTypes.JSON)
-    public static Slice varcharJsonExtract(@SqlType("varchar(x)") Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
+    public static Slice varcharJsonExtract(SqlFunctionProperties properties, @SqlType("varchar(x)") Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
     {
-        return JsonExtract.extract(json, jsonPath.getObjectExtractor());
+        return JsonExtract.extract(json, jsonPath.getObjectExtractor(), properties);
     }
 
     @ScalarFunction
     @SqlNullable
     @SqlType(StandardTypes.JSON)
-    public static Slice jsonExtract(@SqlType(StandardTypes.JSON) Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
+    public static Slice jsonExtract(SqlFunctionProperties properties, @SqlType(StandardTypes.JSON) Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
     {
-        return JsonExtract.extract(json, jsonPath.getObjectExtractor());
+        return JsonExtract.extract(json, jsonPath.getObjectExtractor(), properties);
     }
 
     @ScalarFunction("json_size")
     @LiteralParameters("x")
     @SqlNullable
     @SqlType(StandardTypes.BIGINT)
-    public static Long varcharJsonSize(@SqlType("varchar(x)") Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
+    public static Long varcharJsonSize(SqlFunctionProperties properties, @SqlType("varchar(x)") Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
     {
-        return JsonExtract.extract(json, jsonPath.getSizeExtractor());
+        return JsonExtract.extract(json, jsonPath.getSizeExtractor(), properties);
     }
 
     @ScalarFunction
     @SqlNullable
     @SqlType(StandardTypes.BIGINT)
-    public static Long jsonSize(@SqlType(StandardTypes.JSON) Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
+    public static Long jsonSize(SqlFunctionProperties properties, @SqlType(StandardTypes.JSON) Slice json, @SqlType(JsonPathType.NAME) JsonPath jsonPath)
     {
-        return JsonExtract.extract(json, jsonPath.getSizeExtractor());
+        return JsonExtract.extract(json, jsonPath.getSizeExtractor(), properties);
     }
 
     public static Object getJsonObjectValue(Type valueType, SqlFunctionProperties properties, Block block, int position)