Add Elasticsearch array support using definitions in the _meta field

Cherry-pick of trinodb/trino@e5bd9eb

Co-authored-by: Brian Olsen brianolsen87@gmail.com

presto-docs/src/main/sphinx/connector/elasticsearch.rst

@@ -190,6 +190,53 @@ Elasticsearch Presto
 (others)      (unsupported)
 ============= =============
 
+
+Array Types
+^^^^^^^^^^^
+
+Fields in Elasticsearch can contain `zero or more values <https://www.elastic.co/guide/en/elasticsearch/reference/current/array.html>`_
+, but there is no dedicated array type. To indicate a field contains an array, it can be annotated in a Presto-specific structure in
+the `_meta <https://www.elastic.co/guide/en/elasticsearch/reference/current/mapping-meta-field.html>`_ section of the index mapping.
+
+For example, you can have an Elasticsearch index that contains documents with the following structure:
+
+.. code-block:: json
+
+    {
+        "array_string_field": ["presto","is","the","besto"],
+        "long_field": 314159265359,
+        "id_field": "564e6982-88ee-4498-aa98-df9e3f6b6109",
+        "timestamp_field": "1987-09-17T06:22:48.000Z",
+        "object_field": {
+            "array_int_field": [86,75,309],
+            "int_field": 2
+        }
+    }
+
+The array fields of this structure can be defined by using the following command to add the field
+property definition to the ``_meta.presto`` property of the target index mapping.
+
+.. code-block:: shell
+
+    curl --request PUT \
+        --url localhost:9200/doc/_mapping \
+        --header 'content-type: application/json' \
+        --data '
+    {
+        "_meta": {
+            "presto":{
+                "array_string_field":{
+                    "isArray":true
+                },
+                "object_field":{
+                    "array_int_field":{
+                        "isArray":true
+                    }
+                },
+            }
+        }
+    }'
+
 Special Columns
 ---------------
 

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java

@@ -27,6 +27,7 @@
 import com.facebook.presto.spi.PrestoException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.NullNode;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import io.airlift.units.Duration;
@@ -445,15 +446,17 @@ public IndexMetadata getIndexMetadata(String index)
                     mappings = mappings.elements().next();
                 }
 
-                return new IndexMetadata(parseType(mappings.get("properties")));
+                JsonNode metaNode = nullSafeNode(mappings, "_meta");
+
+                return new IndexMetadata(parseType(mappings.get("properties"), nullSafeNode(metaNode, "presto")));
             }
             catch (IOException e) {
                 throw new PrestoException(ELASTICSEARCH_INVALID_RESPONSE, e);
             }
         });
     }
 
-    private IndexMetadata.ObjectType parseType(JsonNode properties)
+    private IndexMetadata.ObjectType parseType(JsonNode properties, JsonNode metaProperties)
     {
         Iterator<Map.Entry<String, JsonNode>> entries = properties.fields();
 
@@ -463,28 +466,38 @@ private IndexMetadata.ObjectType parseType(JsonNode properties)
 
             String name = field.getKey();
             JsonNode value = field.getValue();
+
+            JsonNode metaNode = nullSafeNode(metaProperties, name);
+            boolean isArray = !metaNode.isNull() && metaNode.has("isArray") && metaNode.get("isArray").asBoolean();
             if (value.has("type")) {
                 String type = value.get("type").asText();
-
                 if (type.equals("date")) {
                     List<String> formats = ImmutableList.of();
                     if (value.has("format")) {
                         formats = Arrays.asList(value.get("format").asText().split("\\|\\|"));
                     }
-                    result.add(new IndexMetadata.Field(name, new IndexMetadata.DateTimeType(formats)));
+                    result.add(new IndexMetadata.Field(isArray, name, new IndexMetadata.DateTimeType(formats)));
                 }
                 else {
-                    result.add(new IndexMetadata.Field(name, new IndexMetadata.PrimitiveType(type)));
+                    result.add(new IndexMetadata.Field(isArray, name, new IndexMetadata.PrimitiveType(type)));
                 }
             }
             else if (value.has("properties")) {
-                result.add(new IndexMetadata.Field(name, parseType(value.get("properties"))));
+                result.add(new IndexMetadata.Field(isArray, name, parseType(value.get("properties"), metaNode)));
             }
         }
 
         return new IndexMetadata.ObjectType(result.build());
     }
 
+    private JsonNode nullSafeNode(JsonNode jsonNode, String name)
+    {
+        if (jsonNode == null || jsonNode.isNull() || jsonNode.get(name) == null) {
+            return NullNode.getInstance();
+        }
+        return jsonNode.get(name);
+    }
+
     public SearchResponse beginSearch(String index, int shard, QueryBuilder query, Optional<List<String>> fields, List<String> documentFields, Optional<String> sort)
     {
         SearchSourceBuilder sourceBuilder = SearchSourceBuilder.searchSource()

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java

@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.elasticsearch;
 
+import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.common.type.RowType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.elasticsearch.client.ElasticsearchClient;
@@ -140,7 +141,7 @@ private List<ColumnMetadata> toColumnMetadata(IndexMetadata metadata)
         result.add(BuiltinColumns.SCORE.getMetadata());
 
         for (IndexMetadata.Field field : metadata.getSchema().getFields()) {
-            Type type = toPrestoType(field.getType());
+            Type type = toPrestoType(field);
             if (type == null) {
                 continue;
             }
@@ -174,9 +175,19 @@ private static boolean supportsPredicates(IndexMetadata.Type type)
         return false;
     }
 
-    private Type toPrestoType(IndexMetadata.Type type)
+    private Type toPrestoType(IndexMetadata.Field metaDataField)
     {
-        if (type instanceof PrimitiveType) {
+        return toPrestoType(metaDataField, metaDataField.isArray());
+    }
+
+    private Type toPrestoType(IndexMetadata.Field metaDataField, boolean isArray)
+    {
+        IndexMetadata.Type type = metaDataField.getType();
+        if (isArray) {
+            Type elementType = toPrestoType(metaDataField, false);
+            return new ArrayType(elementType);
+        }
+        else if (type instanceof PrimitiveType) {
             switch (((PrimitiveType) type).getName()) {
                 case "float":
                     return REAL;
@@ -210,7 +221,7 @@ else if (type instanceof ObjectType) {
             ObjectType objectType = (ObjectType) type;
 
             List<Field> fields = objectType.getFields().stream()
-                    .map(field -> RowType.field(field.getName(), toPrestoType(field.getType())))
+                    .map(field -> RowType.field(field.getName(), toPrestoType(field)))
                     .collect(toImmutableList());
 
             return RowType.from(fields);

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchPageSource.java

@@ -18,9 +18,11 @@
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.BlockBuilder;
 import com.facebook.presto.common.block.PageBuilderStatus;
+import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.common.type.RowType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.elasticsearch.client.ElasticsearchClient;
+import com.facebook.presto.elasticsearch.decoders.ArrayDecoder;
 import com.facebook.presto.elasticsearch.decoders.BigintDecoder;
 import com.facebook.presto.elasticsearch.decoders.BooleanDecoder;
 import com.facebook.presto.elasticsearch.decoders.Decoder;
@@ -334,6 +336,11 @@ else if (type instanceof RowType) {
 
             return new RowDecoder(fieldNames, decoders);
         }
+        else if (type instanceof ArrayType) {
+            Type elementType = ((ArrayType) type).getElementType();
+
+            return new ArrayDecoder(createDecoder(session, path, elementType));
+        }
 
         throw new UnsupportedOperationException("Type not supported: " + type);
     }

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/client/IndexMetadata.java

@@ -35,15 +35,22 @@ public ObjectType getSchema()
 
     public static class Field
     {
+        private final boolean isArray;
         private final String name;
         private final Type type;
 
-        public Field(String name, Type type)
+        public Field(boolean isArray, String name, Type type)
         {
+            this.isArray = isArray;
             this.name = requireNonNull(name, "name is null");
             this.type = requireNonNull(type, "type is null");
         }
 
+        public boolean isArray()
+        {
+            return isArray;
+        }
+
         public String getName()
         {
             return name;

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/decoders/ArrayDecoder.java

@@ -0,0 +1,52 @@
+/*
+ * 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.elasticsearch.decoders;
+
+import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.spi.PrestoException;
+import org.elasticsearch.search.SearchHit;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+import static com.facebook.presto.elasticsearch.ElasticsearchErrorCode.ELASTICSEARCH_TYPE_MISMATCH;
+
+public class ArrayDecoder
+        implements Decoder
+{
+    private final Decoder elementDecoder;
+
+    public ArrayDecoder(Decoder elementDecoder)
+    {
+        this.elementDecoder = elementDecoder;
+    }
+
+    @Override
+    public void decode(SearchHit hit, Supplier<Object> getter, BlockBuilder output)
+    {
+        Object data = getter.get();
+
+        if (data == null) {
+            output.appendNull();
+        }
+        else if (data instanceof List) {
+            BlockBuilder array = output.beginBlockEntry();
+            ((List) data).forEach(element -> elementDecoder.decode(hit, () -> element, array));
+            output.closeEntry();
+        }
+        else {
+            throw new PrestoException(ELASTICSEARCH_TYPE_MISMATCH, "Expected list of elements for ARRAY field");
+        }
+    }
+}