Expand support of types handled in EXPLAIN (TYPE IO)

presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveIntegrationSmokeTest.java

@@ -33,8 +33,19 @@
 import io.prestosql.spi.connector.Constraint;
 import io.prestosql.spi.security.Identity;
 import io.prestosql.spi.security.SelectedRole;
+import io.prestosql.spi.type.BigintType;
+import io.prestosql.spi.type.BooleanType;
+import io.prestosql.spi.type.CharType;
+import io.prestosql.spi.type.DateType;
+import io.prestosql.spi.type.DecimalType;
+import io.prestosql.spi.type.DoubleType;
+import io.prestosql.spi.type.IntegerType;
+import io.prestosql.spi.type.SmallintType;
+import io.prestosql.spi.type.TimestampType;
+import io.prestosql.spi.type.TinyintType;
 import io.prestosql.spi.type.Type;
 import io.prestosql.spi.type.TypeSignature;
+import io.prestosql.spi.type.VarcharType;
 import io.prestosql.sql.planner.Plan;
 import io.prestosql.sql.planner.plan.ExchangeNode;
 import io.prestosql.sql.planner.planPrinter.IoPlanPrinter.ColumnConstraint;
@@ -232,6 +243,51 @@ public void testIOExplain()
         assertUpdate("DROP TABLE test_orders");
     }
 
+    @Test
+    public void testIOExplainWithPrimitiveTypes()
+    {
+        Map<Object, Type> sampleDataMap = new HashMap<>();
+        sampleDataMap.put("foo", VarcharType.createUnboundedVarcharType());
+        sampleDataMap.put(Byte.toString((byte) (Byte.MAX_VALUE / 2)), TinyintType.TINYINT);
+        sampleDataMap.put(Short.toString((short) (Short.MAX_VALUE / 2)), SmallintType.SMALLINT);
+        sampleDataMap.put(Integer.toString(Integer.MAX_VALUE / 2), IntegerType.INTEGER);
+        sampleDataMap.put(Long.toString(Long.MAX_VALUE / 2), BigintType.BIGINT);
+        sampleDataMap.put(Boolean.TRUE.toString(), BooleanType.BOOLEAN);
+        sampleDataMap.put("bar", CharType.createCharType(3));
+        sampleDataMap.put("1.2345678901234578E14", DoubleType.DOUBLE);
+        sampleDataMap.put("123456789012345678901234.567", DecimalType.createDecimalType(30, 3));
+        sampleDataMap.put("2019-01-01", DateType.DATE);
+        sampleDataMap.put("2019-01-01 23:22:21.123", TimestampType.TIMESTAMP);
+        for (Map.Entry<Object, Type> entry : sampleDataMap.entrySet()) {
+            @Language("SQL") String query = "" +
+                    "CREATE TABLE test_types_table  WITH (partitioned_by = ARRAY['my_col']) AS " +
+                    "SELECT " +
+                    "'foo' my_non_partition_col" +
+                    ",CAST('" + entry.getKey() + "' as " + entry.getValue().getDisplayName() + ") my_col";
+
+            assertUpdate(query, 1);
+            MaterializedResult result = computeActual("EXPLAIN (TYPE IO, FORMAT JSON) SELECT * FROM test_types_table");
+            assertEquals(
+                    jsonCodec(IoPlan.class).fromJson((String) getOnlyElement(result.getOnlyColumnAsSet())),
+                    new IoPlan(
+                            ImmutableSet.of(new TableColumnInfo(
+                                    new CatalogSchemaTableName(catalog, "tpch", "test_types_table"),
+                                    ImmutableSet.of(
+                                            new ColumnConstraint(
+                                                    "my_col",
+                                                    entry.getValue().getTypeSignature(),
+                                                    new FormattedDomain(
+                                                            false,
+                                                            ImmutableSet.of(
+                                                                    new FormattedRange(
+                                                                            new FormattedMarker(Optional.of(entry.getKey().toString()), EXACTLY),
+                                                                            new FormattedMarker(Optional.of(entry.getKey().toString()), EXACTLY)))))))),
+                            Optional.empty()));
+
+            assertUpdate("DROP TABLE test_types_table");
+        }
+    }
+
     @Test
     public void testReadNoColumns()
     {

presto-main/src/main/java/io/prestosql/sql/planner/planPrinter/IoPlanPrinter.java

@@ -16,7 +16,6 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableSet;
-import io.airlift.slice.Slice;
 import io.prestosql.Session;
 import io.prestosql.metadata.Metadata;
 import io.prestosql.metadata.TableHandle;
@@ -31,8 +30,14 @@
 import io.prestosql.spi.predicate.TupleDomain;
 import io.prestosql.spi.type.BigintType;
 import io.prestosql.spi.type.BooleanType;
+import io.prestosql.spi.type.CharType;
+import io.prestosql.spi.type.DateType;
+import io.prestosql.spi.type.DecimalType;
+import io.prestosql.spi.type.DoubleType;
 import io.prestosql.spi.type.IntegerType;
 import io.prestosql.spi.type.SmallintType;
+import io.prestosql.spi.type.TimeType;
+import io.prestosql.spi.type.TimestampType;
 import io.prestosql.spi.type.TinyintType;
 import io.prestosql.spi.type.Type;
 import io.prestosql.spi.type.TypeSignature;
@@ -61,11 +66,21 @@
 import static io.airlift.json.JsonCodec.jsonCodec;
 import static io.prestosql.spi.StandardErrorCode.NOT_SUPPORTED;
 import static io.prestosql.spi.predicate.Marker.Bound.EXACTLY;
+import static io.prestosql.sql.planner.planPrinter.PlanPrinterUtil.castToVarchar;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public class IoPlanPrinter
 {
+    /**
+     * List of supported data types for column constraints appearing in
+     * EXPLAIN (TYPE IO). Other types would cause a NOT_SUPPORTED
+     * Presto exception to be thrown.
+     */
+    private static final Set<Class> WHITELISTED_TYPES_FOR_OUTPUT = ImmutableSet.of(
+            VarcharType.class, TinyintType.class, SmallintType.class, IntegerType.class, BigintType.class,
+            BooleanType.class, CharType.class, DecimalType.class, DoubleType.class, DateType.class,
+            TimeType.class, TimestampType.class);
     private final Metadata metadata;
     private final Session session;
 
@@ -564,14 +579,10 @@ private FormattedMarker formatMarker(Marker marker)
 
         private String getVarcharValue(Type type, Object value)
         {
-            if (type instanceof VarcharType) {
-                return ((Slice) value).toStringUtf8();
-            }
-            if (type instanceof TinyintType || type instanceof SmallintType || type instanceof IntegerType || type instanceof BigintType) {
-                return ((Long) value).toString();
-            }
-            if (type instanceof BooleanType) {
-                return ((Boolean) value).toString();
+            for (Class<Type> whitelistedType : WHITELISTED_TYPES_FOR_OUTPUT) {
+                if (whitelistedType.getClass().isInstance(type.getClass())) {
+                    return castToVarchar(type, value, metadata.getFunctionRegistry(), session);
+                }
             }
             throw new PrestoException(NOT_SUPPORTED, format("Unsupported data type in EXPLAIN (TYPE IO): %s", type.getDisplayName()));
         }

presto-main/src/main/java/io/prestosql/sql/planner/planPrinter/PlanPrinter.java

@@ -21,7 +21,6 @@
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Streams;
-import io.airlift.slice.Slice;
 import io.airlift.units.Duration;
 import io.prestosql.Session;
 import io.prestosql.SystemSessionProperties;
@@ -31,8 +30,6 @@
 import io.prestosql.execution.StageInfo;
 import io.prestosql.execution.StageStats;
 import io.prestosql.metadata.FunctionRegistry;
-import io.prestosql.metadata.OperatorNotFoundException;
-import io.prestosql.metadata.Signature;
 import io.prestosql.metadata.TableHandle;
 import io.prestosql.operator.StageExecutionDescriptor;
 import io.prestosql.spi.connector.ColumnHandle;
@@ -44,7 +41,6 @@
 import io.prestosql.spi.statistics.ColumnStatisticMetadata;
 import io.prestosql.spi.statistics.TableStatisticType;
 import io.prestosql.spi.type.Type;
-import io.prestosql.sql.InterpretedFunctionInvoker;
 import io.prestosql.sql.planner.OrderingScheme;
 import io.prestosql.sql.planner.Partitioning;
 import io.prestosql.sql.planner.PartitioningScheme;
@@ -122,9 +118,9 @@
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static io.prestosql.execution.StageInfo.getAllStages;
 import static io.prestosql.operator.StageExecutionDescriptor.ungroupedExecution;
-import static io.prestosql.spi.type.VarcharType.VARCHAR;
 import static io.prestosql.sql.planner.SystemPartitioningHandle.SINGLE_DISTRIBUTION;
 import static io.prestosql.sql.planner.planPrinter.PlanNodeStatsSummarizer.aggregateStageStats;
+import static io.prestosql.sql.planner.planPrinter.PlanPrinterUtil.castToVarchar;
 import static io.prestosql.sql.planner.planPrinter.TextRenderer.formatDouble;
 import static io.prestosql.sql.planner.planPrinter.TextRenderer.formatPositions;
 import static io.prestosql.sql.planner.planPrinter.TextRenderer.indentString;
@@ -1182,22 +1178,6 @@ public NodeRepresentation addNode(PlanNode rootNode, String name, String identif
         }
     }
 
-    private static String castToVarchar(Type type, Object value, FunctionRegistry functionRegistry, Session session)
-    {
-        if (value == null) {
-            return "NULL";
-        }
-
-        try {
-            Signature coercion = functionRegistry.getCoercion(type, VARCHAR);
-            Slice coerced = (Slice) new InterpretedFunctionInvoker(functionRegistry).invoke(coercion, session.toConnectorSession(), value);
-            return coerced.toStringUtf8();
-        }
-        catch (OperatorNotFoundException e) {
-            return "<UNREPRESENTABLE VALUE>";
-        }
-    }
-
     private static String formatFrame(WindowNode.Frame frame)
     {
         StringBuilder builder = new StringBuilder(frame.getType().toString());

presto-main/src/main/java/io/prestosql/sql/planner/planPrinter/PlanPrinterUtil.java

@@ -0,0 +1,45 @@
+/*
+ * 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 io.prestosql.sql.planner.planPrinter;
+
+import io.airlift.slice.Slice;
+import io.prestosql.Session;
+import io.prestosql.metadata.FunctionRegistry;
+import io.prestosql.metadata.OperatorNotFoundException;
+import io.prestosql.metadata.Signature;
+import io.prestosql.spi.type.Type;
+import io.prestosql.sql.InterpretedFunctionInvoker;
+
+import static io.prestosql.spi.type.VarcharType.VARCHAR;
+
+public final class PlanPrinterUtil
+{
+    private PlanPrinterUtil() {}
+
+    static String castToVarchar(Type type, Object value, FunctionRegistry functionRegistry, Session session)
+    {
+        if (value == null) {
+            return "NULL";
+        }
+
+        try {
+            Signature coercion = functionRegistry.getCoercion(type, VARCHAR);
+            Slice coerced = (Slice) new InterpretedFunctionInvoker(functionRegistry).invoke(coercion, session.toConnectorSession(), value);
+            return coerced.toStringUtf8();
+        }
+        catch (OperatorNotFoundException e) {
+            return "<UNREPRESENTABLE VALUE>";
+        }
+    }
+}