Add support for json type in MongoDB

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

@@ -13,13 +13,9 @@
  */
 package com.facebook.presto.operator.scalar;
 
-import com.facebook.airlift.json.JsonObjectMapperProvider;
-import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.function.OperatorType;
-import com.facebook.presto.common.function.SqlFunctionProperties;
-import com.facebook.presto.common.type.SqlDecimal;
 import com.facebook.presto.common.type.StandardTypes;
-import com.facebook.presto.common.type.Type;
+import com.facebook.presto.plugin.base.JsonTypeUtil;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
@@ -32,14 +28,10 @@
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
 import com.fasterxml.jackson.databind.MappingJsonFactory;
-import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.primitives.Doubles;
-import io.airlift.slice.DynamicSliceOutput;
 import io.airlift.slice.Slice;
-import io.airlift.slice.SliceOutput;
 
 import java.io.IOException;
-import java.io.OutputStream;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -57,9 +49,7 @@
 import static com.fasterxml.jackson.core.JsonToken.VALUE_NUMBER_INT;
 import static com.fasterxml.jackson.core.JsonToken.VALUE_STRING;
 import static com.fasterxml.jackson.core.JsonToken.VALUE_TRUE;
-import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
 import static io.airlift.slice.Slices.utf8Slice;
-import static java.lang.String.format;
 
 public final class JsonFunctions
 {
@@ -69,8 +59,6 @@ public final class JsonFunctions
     private static final JsonFactory MAPPING_JSON_FACTORY = new MappingJsonFactory()
             .disable(CANONICALIZE_FIELD_NAMES);
 
-    private static final ObjectMapper SORTED_MAPPER = new JsonObjectMapperProvider().get().configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
-
     private JsonFunctions() {}
 
     @ScalarOperator(OperatorType.CAST)
@@ -139,20 +127,7 @@ public static Slice jsonFormat(@SqlType(StandardTypes.JSON) Slice slice)
     @SqlType(StandardTypes.JSON)
     public static Slice jsonParse(@SqlType("varchar(x)") Slice slice)
     {
-        // cast(json_parse(x) AS t)` will be optimized into `$internal$json_string_to_array/map/row_cast` in ExpressionOptimizer
-        // If you make changes to this function (e.g. use parse JSON string into some internal representation),
-        // make sure `$internal$json_string_to_array/map/row_cast` is changed accordingly.
-        try (JsonParser parser = createJsonParser(JSON_FACTORY, slice)) {
-            SliceOutput dynamicSliceOutput = new DynamicSliceOutput(slice.length());
-            SORTED_MAPPER.writeValue((OutputStream) dynamicSliceOutput, SORTED_MAPPER.readValue(parser, Object.class));
-            // nextToken() returns null if the input is parsed correctly,
-            // but will throw an exception if there are trailing characters.
-            parser.nextToken();
-            return dynamicSliceOutput.slice();
-        }
-        catch (Exception e) {
-            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Cannot convert '%s' to JSON", slice.toStringUtf8()));
-        }
+        return JsonTypeUtil.jsonParse(slice);
     }
 
     @SqlNullable
@@ -481,13 +456,4 @@ public static Long jsonSize(@SqlType(StandardTypes.JSON) Slice json, @SqlType(Js
     {
         return JsonExtract.extract(json, jsonPath.getSizeExtractor());
     }
-
-    public static Object getJsonObjectValue(Type valueType, SqlFunctionProperties properties, Block block, int position)
-    {
-        Object objectValue = valueType.getObjectValue(properties, block, position);
-        if (objectValue instanceof SqlDecimal) {
-            objectValue = ((SqlDecimal) objectValue).toBigDecimal();
-        }
-        return objectValue;
-    }
 }

presto-mongodb/pom.xml

@@ -18,6 +18,11 @@
     </properties>
 
     <dependencies>
+        <dependency>
+            <groupId>com.facebook.presto</groupId>
+            <artifactId>presto-plugin-toolkit</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>org.mongodb</groupId>
             <artifactId>mongo-java-driver</artifactId>

presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoPageSink.java

@@ -40,6 +40,7 @@
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.model.InsertManyOptions;
 import io.airlift.slice.Slice;
+import org.bson.BsonInvalidOperationException;
 import org.bson.Document;
 import org.bson.types.Binary;
 import org.bson.types.ObjectId;
@@ -60,6 +61,7 @@
 import static com.facebook.presto.common.type.Varchars.isVarcharType;
 import static com.facebook.presto.mongodb.ObjectIdType.OBJECT_ID;
 import static com.facebook.presto.mongodb.TypeUtils.isArrayType;
+import static com.facebook.presto.mongodb.TypeUtils.isJsonType;
 import static com.facebook.presto.mongodb.TypeUtils.isMapType;
 import static com.facebook.presto.mongodb.TypeUtils.isRowType;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
@@ -175,6 +177,15 @@ private Object getObjectValue(Type type, Block block, int position)
             }
             return new BigDecimal(unscaledValue);
         }
+        if (isJsonType(type)) {
+            String json = type.getSlice(block, position).toStringUtf8();
+            try {
+                return Document.parse(json);
+            }
+            catch (BsonInvalidOperationException e) {
+                throw new PrestoException(NOT_SUPPORTED, "Can't convert json to MongoDB Document: " + json, e);
+            }
+        }
         if (isArrayType(type)) {
             Type elementType = type.getTypeParameters().get(0);
 

presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoPageSource.java

@@ -46,8 +46,10 @@
 import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
 import static com.facebook.presto.mongodb.ObjectIdType.OBJECT_ID;
 import static com.facebook.presto.mongodb.TypeUtils.isArrayType;
+import static com.facebook.presto.mongodb.TypeUtils.isJsonType;
 import static com.facebook.presto.mongodb.TypeUtils.isMapType;
 import static com.facebook.presto.mongodb.TypeUtils.isRowType;
+import static com.facebook.presto.plugin.base.JsonTypeUtil.jsonParse;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
@@ -222,6 +224,9 @@ else if (type.equals(VARBINARY)) {
                 output.appendNull();
             }
         }
+        else if (isJsonType(type)) {
+            type.writeSlice(output, jsonParse(utf8Slice(toVarcharValue(value))));
+        }
         else {
             throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Slice: " + type.getTypeSignature());
         }

presto-mongodb/src/main/java/com/facebook/presto/mongodb/TypeUtils.java

@@ -19,6 +19,7 @@
 import java.util.function.Predicate;
 
 import static com.facebook.presto.common.type.DateType.DATE;
+import static com.facebook.presto.common.type.StandardTypes.JSON;
 import static com.facebook.presto.common.type.TimeType.TIME;
 import static com.facebook.presto.common.type.TimestampType.TIMESTAMP;
 import static com.facebook.presto.common.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE;
@@ -27,6 +28,11 @@ public final class TypeUtils
 {
     private TypeUtils() {}
 
+    public static boolean isJsonType(Type type)
+    {
+        return type.getTypeSignature().getBase().equals(JSON);
+    }
+
     public static boolean isArrayType(Type type)
     {
         return type.getTypeSignature().getBase().equals(StandardTypes.ARRAY);

presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java

@@ -78,7 +78,8 @@ public void createTableWithEveryType()
                 ", DATE '1980-05-07' _date" +
                 ", TIMESTAMP '1980-05-07 11:22:33.456' _timestamp" +
                 ", ObjectId('ffffffffffffffffffffffff') _objectid" +
-                ", cast(ObjectId('ffffffffffffffffffffffff') as varchar) _objectid_string";
+                ", cast(ObjectId('ffffffffffffffffffffffff') as varchar) _objectid_string" +
+                ", JSON '{\"name\":\"alice\"}' _json";
 
         assertUpdate(query, 1);
 
@@ -93,6 +94,7 @@ public void createTableWithEveryType()
         assertEquals(row.getField(5), LocalDate.of(1980, 5, 7));
         assertEquals(row.getField(6), LocalDateTime.of(1980, 5, 7, 11, 22, 33, 456_000_000));
         assertEquals(row.getField(8), "ffffffffffffffffffffffff");
+        assertEquals(row.getField(9), "{\"name\":\"alice\"}");
         assertUpdate("DROP TABLE test_types_table");
 
         assertFalse(getQueryRunner().tableExists(getSession(), "test_types_table"));
@@ -114,6 +116,7 @@ public void testInsertWithEveryType()
                 ", ts  timestamp" +
                 ", tm time" +
                 ", objid objectid" +
+                ", _json json" +
                 ")";
         getQueryRunner().execute(getSession(), createSql);
 
@@ -128,7 +131,8 @@ public void testInsertWithEveryType()
                 ", DATE '1980-05-07' _date" +
                 ", TIMESTAMP '1980-05-07 11:22:33.456' _timestamp" +
                 ", TIME '11:22:33.456' _time" +
-                ", ObjectId('ffffffffffffffffffffffff') _objectid";
+                ", ObjectId('ffffffffffffffffffffffff') _objectid" +
+                ", JSON '{\"name\":\"alice\"}' _json";
         getQueryRunner().execute(getSession(), insertSql);
 
         MaterializedResult results = getQueryRunner().execute(getSession(), "SELECT * FROM test_insert_types_table").toTestTypes();
@@ -142,10 +146,36 @@ public void testInsertWithEveryType()
         assertEquals(row.getField(5), LocalDate.of(1980, 5, 7));
         assertEquals(row.getField(6), LocalDateTime.of(1980, 5, 7, 11, 22, 33, 456_000_000));
         assertEquals(row.getField(7), LocalTime.of(11, 22, 33, 456_000_000));
+        assertEquals(row.getField(9), "{\"name\":\"alice\"}");
         assertUpdate("DROP TABLE test_insert_types_table");
         assertFalse(getQueryRunner().tableExists(getSession(), "test_insert_types_table"));
     }
 
+    @Test
+    public void testJson()
+    {
+        assertUpdate("CREATE TABLE test_json (id INT, col JSON)");
+
+        assertUpdate("INSERT INTO test_json VALUES (1, JSON '{\"name\":\"alice\"}')", 1);
+        assertQuery("SELECT json_extract_scalar(col, '$.name') FROM test_json WHERE id = 1", "SELECT 'alice'");
+
+        assertUpdate("INSERT INTO test_json VALUES (2, JSON '{\"numbers\":[1, 2, 3]}')", 1);
+        assertQuery("SELECT json_extract(col, '$.numbers[0]') FROM test_json WHERE id = 2", "SELECT 1");
+
+        assertUpdate("INSERT INTO test_json VALUES (3, NULL)", 1);
+        assertQuery("SELECT col FROM test_json WHERE id = 3", "SELECT NULL");
+
+        assertQueryFails(
+                "CREATE TABLE test_json_scalar AS SELECT JSON '1' AS col",
+                "Can't convert json to MongoDB Document.*");
+
+        assertQueryFails(
+                "CREATE TABLE test_json_array AS SELECT JSON '[\"a\", \"b\", \"c\"]' AS col",
+                "Can't convert json to MongoDB Document.*");
+
+        assertUpdate("DROP TABLE test_json");
+    }
+
     @Test
     public void testArrays()
     {

presto-plugin-toolkit/pom.xml

@@ -78,6 +78,16 @@
             <artifactId>jackson-annotations</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-core</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>slice</artifactId>
+        </dependency>
+
         <!-- for testing -->
         <dependency>
             <groupId>com.facebook.presto</groupId>

presto-plugin-toolkit/src/main/java/com/facebook/presto/plugin/base/JsonTypeUtil.java

@@ -0,0 +1,76 @@
+/*
+ * Licensed 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 com.facebook.presto.plugin.base;
+
+import com.facebook.airlift.json.JsonObjectMapperProvider;
+import com.facebook.presto.spi.PrestoException;
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonFactoryBuilder;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.airlift.slice.DynamicSliceOutput;
+import io.airlift.slice.Slice;
+import io.airlift.slice.SliceOutput;
+import io.airlift.slice.Slices;
+
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+
+import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
+import static com.fasterxml.jackson.core.JsonFactory.Feature.CANONICALIZE_FIELD_NAMES;
+import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
+import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+public class JsonTypeUtil
+{
+    private static final JsonFactory JSON_FACTORY = new JsonFactoryBuilder().disable(CANONICALIZE_FIELD_NAMES).build();
+    private static final ObjectMapper SORTED_MAPPER = new JsonObjectMapperProvider().get().configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
+
+    private JsonTypeUtil() {}
+
+    // TODO: this should be available from the engine
+    public static Slice jsonParse(Slice slice)
+    {
+        // cast(json_parse(x) AS t)` will be optimized into `$internal$json_string_to_array/map/row_cast` in ExpressionOptimizer
+        // If you make changes to this function (e.g. use parse JSON string into some internal representation),
+        // make sure `$internal$json_string_to_array/map/row_cast` is changed accordingly.
+        try (JsonParser parser = createJsonParser(JSON_FACTORY, slice)) {
+            SliceOutput dynamicSliceOutput = new DynamicSliceOutput(slice.length());
+            SORTED_MAPPER.writeValue((OutputStream) dynamicSliceOutput, SORTED_MAPPER.readValue(parser, Object.class));
+            // nextToken() returns null if the input is parsed correctly,
+            // but will throw an exception if there are trailing characters.
+            parser.nextToken();
+            return dynamicSliceOutput.slice();
+        }
+        catch (IOException | RuntimeException e) {
+            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Cannot convert value to JSON: '%s'", slice.toStringUtf8()), e);
+        }
+    }
+
+    public static Slice toJsonValue(Object value)
+            throws IOException
+    {
+        return Slices.wrappedBuffer(SORTED_MAPPER.writeValueAsBytes(value));
+    }
+
+    private static JsonParser createJsonParser(JsonFactory factory, Slice json)
+            throws IOException
+    {
+        // Jackson tries to detect the character encoding automatically when
+        // using InputStream, so we pass an InputStreamReader instead.
+        return factory.createParser(new InputStreamReader(json.getInput(), UTF_8));
+    }
+}