[WIP] Support for user-defined enums

presto-cli/src/main/java/com/facebook/presto/cli/OutputHandler.java

@@ -13,13 +13,18 @@
  */
 package com.facebook.presto.cli;
 
+import com.facebook.presto.client.Column;
 import com.facebook.presto.client.StatementClient;
+import com.facebook.presto.common.type.TypeMetadata;
+import com.google.common.collect.ImmutableList;
 import io.airlift.units.Duration;
 
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static io.airlift.units.Duration.nanosSince;
@@ -44,19 +49,46 @@ public OutputHandler(OutputPrinter printer)
         this.printer = requireNonNull(printer, "printer is null");
     }
 
-    public void processRow(List<?> row)
+    public void processRow(List<?> row, List<Column> columns)
             throws IOException
     {
         if (rowBuffer.isEmpty()) {
             bufferStart = System.nanoTime();
         }
 
-        rowBuffer.add(row);
+        rowBuffer.add(prettyPrintRow(row, columns));
         if (rowBuffer.size() >= MAX_BUFFERED_ROWS) {
             flush(false);
         }
     }
 
+    private List<?> prettyPrintRow(List<?> row, List<Column> columns)
+    {
+        ImmutableList.Builder<Object> builder = new ImmutableList.Builder<>();
+        for (int i = 0; i < row.size(); i++) {
+            Object value = row.get(i);
+            Column column = columns.get(i);
+            TypeMetadata typeInfo = column.getTypeMetadata();
+            if (typeInfo == null || typeInfo.getEnumValues() == null) {
+                builder.add(value);
+                continue;
+            }
+            builder.add(getDisplayValueForEnum(value, typeInfo.getEnumValues(), column.getType()));
+        }
+        return builder.build();
+    }
+
+    private Object getDisplayValueForEnum(Object value, Map<String, String> enumEntries, String enumName)
+    {
+        Optional<String> key = enumEntries.entrySet().stream()
+                .filter(e -> e.getValue().equals(value.toString()))
+                .map(Map.Entry::getKey).findFirst();
+        if (key.isPresent()) {
+            return String.format("%s'%s'", enumName, key.get());
+        }
+        return value;
+    }
+
     @Override
     public void close()
             throws IOException
@@ -74,7 +106,7 @@ public void processRows(StatementClient client)
             Iterable<List<Object>> data = client.currentData().getData();
             if (data != null) {
                 for (List<Object> row : data) {
-                    processRow(unmodifiableList(row));
+                    processRow(unmodifiableList(row), client.currentData().getColumns());
                 }
             }
 

presto-client/src/main/java/com/facebook/presto/client/Column.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.client;
 
 import com.facebook.presto.common.type.Type;
+import com.facebook.presto.common.type.TypeMetadata;
 import com.facebook.presto.common.type.TypeSignature;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
@@ -28,26 +29,33 @@ public class Column
     private final String name;
     private final String type;
     private final ClientTypeSignature typeSignature;
+    private final TypeMetadata typeMetadata;
 
     public Column(String name, Type type)
     {
-        this(name, type.getTypeSignature());
+        this(
+                name,
+                type.getTypeSignature().toString(),
+                new ClientTypeSignature(type.getTypeSignature()),
+                type.getTypeMetadata());
     }
 
     public Column(String name, TypeSignature signature)
     {
-        this(name, signature.toString(), new ClientTypeSignature(signature));
+        this(name, signature.toString(), new ClientTypeSignature(signature), null);
     }
 
     @JsonCreator
     public Column(
             @JsonProperty("name") String name,
             @JsonProperty("type") String type,
-            @JsonProperty("typeSignature") ClientTypeSignature typeSignature)
+            @JsonProperty("typeSignature") ClientTypeSignature typeSignature,
+            @JsonProperty("typeMetadata") TypeMetadata typeMetadata)
     {
         this.name = requireNonNull(name, "name is null");
         this.type = requireNonNull(type, "type is null");
         this.typeSignature = typeSignature;
+        this.typeMetadata = requireNonNull(typeMetadata, "typeMetadata is null");
     }
 
     @JsonProperty
@@ -67,4 +75,10 @@ public ClientTypeSignature getTypeSignature()
     {
         return typeSignature;
     }
+
+    @JsonProperty
+    public TypeMetadata getTypeMetadata()
+    {
+        return typeMetadata;
+    }
 }

presto-client/src/main/java/com/facebook/presto/client/FixJsonDataUtils.java

@@ -69,7 +69,10 @@ public static Iterable<List<Object>> fixData(List<Column> columns, Iterable<List
         }
         requireNonNull(columns, "columns is null");
         List<TypeSignature> signatures = columns.stream()
-                .map(column -> parseTypeSignature(column.getType()))
+                .map(column -> {
+                    TypeSignature displayType = column.getTypeMetadata().getDisplayType();
+                    return displayType != null ? displayType : parseTypeSignature(column.getType());
+                })
                 .collect(toList());
         ImmutableList.Builder<List<Object>> rows = ImmutableList.builder();
         for (List<Object> row : data) {

presto-client/src/main/java/com/facebook/presto/client/QueryData.java

@@ -18,4 +18,6 @@
 public interface QueryData
 {
     Iterable<List<Object>> getData();
+
+    List<Column> getColumns();
 }

presto-common/src/main/java/com/facebook/presto/common/type/EnumType.java

@@ -0,0 +1,24 @@
+/*
+ * 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.common.type;
+
+import java.util.Map;
+
+public interface EnumType<T>
+        extends Type
+{
+    Map<String, T> getEntries();
+
+    Type getValueType();
+}

presto-common/src/main/java/com/facebook/presto/common/type/IntegerEnumType.java

@@ -0,0 +1,75 @@
+/*
+ * 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.common.type;
+
+import com.facebook.presto.common.block.Block;
+import com.facebook.presto.common.function.SqlFunctionProperties;
+
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class IntegerEnumType
+        extends AbstractLongType
+        implements EnumType<Long>
+{
+    private final Map<String, Long> entries;
+
+    public IntegerEnumType(String name, Map<String, Long> entries)
+    {
+        super(TypeSignature.parseTypeSignature(name));
+        this.entries = entries;
+    }
+
+    @Override
+    public Map<String, Long> getEntries()
+    {
+        return entries;
+    }
+
+    @Override
+    public Object getObjectValue(SqlFunctionProperties properties, Block block, int position)
+    {
+        if (block.isNull(position)) {
+            return null;
+        }
+
+        return block.getLong(position);
+    }
+
+    @Override
+    public boolean isComparable()
+    {
+        return false;
+    }
+
+    @Override
+    public boolean isOrderable()
+    {
+        return false;
+    }
+
+    @Override
+    public Type getValueType()
+    {
+        return BigintType.BIGINT;
+    }
+
+    @Override
+    public TypeMetadata getTypeMetadata()
+    {
+        Map<String, String> stringifiedEntries = this.getEntries().entrySet().stream()
+                .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().toString()));
+        return new TypeMetadata(getValueType().getTypeSignature(), stringifiedEntries);
+    }
+}

presto-common/src/main/java/com/facebook/presto/common/type/StringEnumType.java

@@ -0,0 +1,59 @@
+/*
+ * 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.common.type;
+
+import java.util.Map;
+
+public class StringEnumType
+        extends VarcharType
+        implements EnumType<String>
+{
+    private final Map<String, String> entries;
+
+    public StringEnumType(String name, Map<String, String> entries)
+    {
+        super(VarcharType.MAX_LENGTH, TypeSignature.parseTypeSignature(name));
+        this.entries = entries;
+    }
+
+    @Override
+    public Map<String, String> getEntries()
+    {
+        return entries;
+    }
+
+    @Override
+    public boolean isOrderable()
+    {
+        return false;
+    }
+
+    @Override
+    public boolean isComparable()
+    {
+        return false;
+    }
+
+    @Override
+    public Type getValueType()
+    {
+        return VarcharType.VARCHAR;
+    }
+
+    @Override
+    public TypeMetadata getTypeMetadata()
+    {
+        return new TypeMetadata(getValueType().getTypeSignature(), this.getEntries());
+    }
+}

presto-common/src/main/java/com/facebook/presto/common/type/Type.java

@@ -186,4 +186,12 @@ public interface Type
      * Compare the values in the specified block at the specified positions equal.
      */
     int compareTo(Block leftBlock, int leftPosition, Block rightBlock, int rightPosition);
+
+    /**
+     * Return additional type information relevant to clients
+     */
+    default TypeMetadata getTypeMetadata()
+    {
+        return new TypeMetadata();
+    }
 }

presto-common/src/main/java/com/facebook/presto/common/type/TypeMetadata.java

@@ -0,0 +1,54 @@
+/*
+ * 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.common.type;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import javax.annotation.concurrent.Immutable;
+
+import java.util.Map;
+
+@Immutable
+public class TypeMetadata
+{
+    private final Map<String, String> enumValues;
+    private final TypeSignature displayType;
+
+    TypeMetadata()
+    {
+        this(null, null);
+    }
+
+    @JsonCreator
+    public TypeMetadata(
+            @JsonProperty TypeSignature displayType,
+            @JsonProperty Map<String, String> enumValues)
+    {
+        this.displayType = displayType;
+        this.enumValues = enumValues;
+    }
+
+    @JsonProperty
+    public TypeSignature getDisplayType()
+    {
+        return displayType;
+    }
+
+    @JsonProperty
+    public Map<String, String> getEnumValues()
+    {
+        return enumValues;
+    }
+}