Variable precision timestamp support for Hive write operations

presto-hive/src/main/java/io/prestosql/plugin/hive/HiveMetadata.java

@@ -390,7 +390,7 @@ public HiveTableHandle getTableHandle(ConnectorSession session, SchemaTableName
                 tableName.getTableName(),
                 table.get().getParameters(),
                 getPartitionKeyColumnHandles(table.get(), typeManager),
-                getRegularColumnHandles(table.get(), typeManager, getTimestampPrecision(session).getPrecision()),
+                getRegularColumnHandles(table.get(), typeManager, getTimestampPrecision(session)),
                 getHiveBucketHandle(session, table.get(), typeManager));
     }
 
@@ -566,7 +566,7 @@ private ConnectorTableMetadata doGetTableMetadata(ConnectorSession session, Sche
 
         Function<HiveColumnHandle, ColumnMetadata> metadataGetter = columnMetadataGetter(table);
         ImmutableList.Builder<ColumnMetadata> columns = ImmutableList.builder();
-        for (HiveColumnHandle columnHandle : hiveColumnHandles(table, typeManager, getTimestampPrecision(session).getPrecision())) {
+        for (HiveColumnHandle columnHandle : hiveColumnHandles(table, typeManager, getTimestampPrecision(session))) {
             columns.add(metadataGetter.apply(columnHandle));
         }
 
@@ -710,7 +710,7 @@ public Map<String, ColumnHandle> getColumnHandles(ConnectorSession session, Conn
         SchemaTableName tableName = ((HiveTableHandle) tableHandle).getSchemaTableName();
         Table table = metastore.getTable(new HiveIdentity(session), tableName.getSchemaName(), tableName.getTableName())
                 .orElseThrow(() -> new TableNotFoundException(tableName));
-        return hiveColumnHandles(table, typeManager, getTimestampPrecision(session).getPrecision()).stream()
+        return hiveColumnHandles(table, typeManager, getTimestampPrecision(session)).stream()
                 .collect(toImmutableMap(HiveColumnHandle::getName, identity()));
     }
 
@@ -849,7 +849,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
             throw new PrestoException(NOT_SUPPORTED, "Bucketing/Partitioning columns not supported when Avro schema url is set");
         }
 
-        validateTimestampColumns(tableMetadata.getColumns());
+        validateTimestampColumns(tableMetadata.getColumns(), getTimestampPrecision(session));
         List<HiveColumnHandle> columnHandles = getColumnHandles(tableMetadata, ImmutableSet.copyOf(partitionedBy));
         HiveStorageFormat hiveStorageFormat = getHiveStorageFormat(tableMetadata.getProperties());
         Map<String, String> tableProperties = getEmptyTableProperties(tableMetadata, bucketProperty, new HdfsContext(session, schemaName, tableName));
@@ -1236,11 +1236,11 @@ public void finishStatisticsCollection(ConnectorSession session, ConnectorTableH
         List<String> partitionColumnNames = partitionColumns.stream()
                 .map(Column::getName)
                 .collect(toImmutableList());
-        // TODO: revisit when handling write path
-        List<HiveColumnHandle> hiveColumnHandles = hiveColumnHandles(table, typeManager, TimestampType.DEFAULT_PRECISION);
+        HiveTimestampPrecision timestampPrecision = getTimestampPrecision(session);
+        List<HiveColumnHandle> hiveColumnHandles = hiveColumnHandles(table, typeManager, timestampPrecision);
         Map<String, Type> columnTypes = hiveColumnHandles.stream()
                 .filter(columnHandle -> !columnHandle.isHidden())
-                .collect(toImmutableMap(HiveColumnHandle::getName, column -> column.getHiveType().getType(typeManager)));
+                .collect(toImmutableMap(HiveColumnHandle::getName, column -> column.getHiveType().getType(typeManager, timestampPrecision)));
 
         Map<List<String>, ComputedStatistics> computedStatisticsMap = createComputedStatisticsToPartitionMap(computedStatistics, partitionColumnNames, columnTypes);
 
@@ -1568,7 +1568,7 @@ public HiveInsertTableHandle beginInsert(ConnectorSession session, ConnectorTabl
             }
         }
 
-        List<HiveColumnHandle> handles = hiveColumnHandles(table, typeManager, getTimestampPrecision(session).getPrecision()).stream()
+        List<HiveColumnHandle> handles = hiveColumnHandles(table, typeManager, getTimestampPrecision(session)).stream()
                 .filter(columnHandle -> !columnHandle.isHidden())
                 .collect(toList());
 
@@ -2447,7 +2447,7 @@ else if (hiveTableHandle.isInAcidTransaction()) {
     @Override
     public Optional<ConnectorNewTableLayout> getNewTableLayout(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {
-        validateTimestampColumns(tableMetadata.getColumns());
+        validateTimestampColumns(tableMetadata.getColumns(), getTimestampPrecision(session));
         validatePartitionColumns(tableMetadata);
         validateBucketColumns(tableMetadata);
         validateColumns(tableMetadata);
@@ -2502,11 +2502,12 @@ public TableStatisticsMetadata getStatisticsCollectionMetadata(ConnectorSession
 
     private TableStatisticsMetadata getStatisticsCollectionMetadata(List<ColumnMetadata> columns, List<String> partitionedBy, Optional<Set<String>> analyzeColumns, boolean includeRowCount)
     {
-        validateTimestampColumns(columns);
         Set<ColumnStatisticMetadata> columnStatistics = columns.stream()
                 .filter(column -> !partitionedBy.contains(column.getName()))
                 .filter(column -> !column.isHidden())
                 .filter(column -> analyzeColumns.isEmpty() || analyzeColumns.get().contains(column.getName()))
+                // TODO: we only support stats collection at millis precision for now (https://github.com/prestosql/presto/issues/5170)
+                .filter(column -> !(column.getType() instanceof TimestampType) || column.getType() == TIMESTAMP_MILLIS)
                 .map(this::getColumnStatisticMetadata)
                 .flatMap(List::stream)
                 .collect(toImmutableSet());
@@ -2720,14 +2721,14 @@ private static void validateColumns(ConnectorTableMetadata tableMetadata)
         }
     }
 
-    // temporary, until variable precision timestamps are supported on write
-    private static void validateTimestampColumns(List<ColumnMetadata> columns)
+    // TODO validate timestamps in structural types (https://github.com/prestosql/presto/issues/5195)
+    private static void validateTimestampColumns(List<ColumnMetadata> columns, HiveTimestampPrecision timestampPrecision)
     {
         for (ColumnMetadata column : columns) {
             Type type = column.getType();
             if (type instanceof TimestampType) {
-                if (type != TIMESTAMP_MILLIS) {
-                    throw new PrestoException(NOT_SUPPORTED, "CREATE TABLE, INSERT and ANALYZE are not supported with requested timestamp precision: " + type);
+                if (((TimestampType) type).getPrecision() != timestampPrecision.getPrecision()) {
+                    throw new PrestoException(NOT_SUPPORTED, format("Incorrect timestamp precision for %s; the configured precision is %s", type, timestampPrecision));
                 }
             }
         }

presto-hive/src/main/java/io/prestosql/plugin/hive/HiveType.java

@@ -127,12 +127,12 @@ public Type getType(TypeManager typeManager)
         return typeManager.getType(getTypeSignature());
     }
 
-    public Type getType(TypeManager typeManager, int timestampPrecision)
+    public Type getType(TypeManager typeManager, HiveTimestampPrecision timestampPrecision)
     {
         Type tentativeType = typeManager.getType(getTypeSignature());
         // TODO: handle timestamps in structural types (https://github.com/prestosql/presto/issues/5195)
         if (tentativeType instanceof TimestampType) {
-            return TimestampType.createTimestampType(timestampPrecision);
+            return TimestampType.createTimestampType(timestampPrecision.getPrecision());
         }
         return tentativeType;
     }

presto-hive/src/main/java/io/prestosql/plugin/hive/HiveWriterFactory.java

@@ -84,6 +84,7 @@
 import static io.prestosql.plugin.hive.HiveSessionProperties.getCompressionCodec;
 import static io.prestosql.plugin.hive.HiveSessionProperties.getInsertExistingPartitionsBehavior;
 import static io.prestosql.plugin.hive.HiveSessionProperties.getTemporaryStagingDirectoryPath;
+import static io.prestosql.plugin.hive.HiveSessionProperties.getTimestampPrecision;
 import static io.prestosql.plugin.hive.HiveSessionProperties.isTemporaryStagingDirectoryEnabled;
 import static io.prestosql.plugin.hive.LocationHandle.WriteMode.DIRECT_TO_TARGET_EXISTING_DIRECTORY;
 import static io.prestosql.plugin.hive.metastore.MetastoreUtil.getHiveSchema;
@@ -544,7 +545,7 @@ else if (insertExistingPartitionsBehavior == InsertExistingPartitionsBehavior.ER
             }
 
             List<Type> types = dataColumns.stream()
-                    .map(column -> column.getHiveType().getType(typeManager))
+                    .map(column -> column.getHiveType().getType(typeManager, getTimestampPrecision(session)))
                     .collect(toImmutableList());
 
             Map<String, Integer> columnIndexes = new HashMap<>();

presto-hive/src/main/java/io/prestosql/plugin/hive/RcFileFileWriterFactory.java

@@ -48,6 +48,7 @@
 import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_WRITE_VALIDATION_FAILED;
 import static io.prestosql.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME;
 import static io.prestosql.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME;
+import static io.prestosql.plugin.hive.HiveSessionProperties.getTimestampPrecision;
 import static io.prestosql.plugin.hive.HiveSessionProperties.isRcfileOptimizedWriterValidate;
 import static io.prestosql.plugin.hive.rcfile.RcFilePageSourceFactory.createTextVectorEncoding;
 import static io.prestosql.plugin.hive.util.HiveUtil.getColumnNames;
@@ -122,7 +123,7 @@ else if (ColumnarSerDe.class.getName().equals(storageFormat.getSerDe())) {
         // an index to rearrange columns in the proper order
         List<String> fileColumnNames = getColumnNames(schema);
         List<Type> fileColumnTypes = getColumnTypes(schema).stream()
-                .map(hiveType -> hiveType.getType(typeManager))
+                .map(hiveType -> hiveType.getType(typeManager, getTimestampPrecision(session)))
                 .collect(toList());
 
         int[] fileInputColumnIndexes = fileColumnNames.stream()

presto-hive/src/main/java/io/prestosql/plugin/hive/RecordFileWriter.java

@@ -45,6 +45,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_WRITER_CLOSE_ERROR;
 import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_WRITER_DATA_ERROR;
+import static io.prestosql.plugin.hive.HiveSessionProperties.getTimestampPrecision;
 import static io.prestosql.plugin.hive.util.HiveUtil.getColumnNames;
 import static io.prestosql.plugin.hive.util.HiveUtil.getColumnTypes;
 import static io.prestosql.plugin.hive.util.HiveWriteUtils.createRecordWriter;
@@ -90,7 +91,7 @@ public RecordFileWriter(
         // existing tables may have columns in a different order
         List<String> fileColumnNames = getColumnNames(schema);
         List<Type> fileColumnTypes = getColumnTypes(schema).stream()
-                .map(hiveType -> hiveType.getType(typeManager))
+                .map(hiveType -> hiveType.getType(typeManager, getTimestampPrecision(session)))
                 .collect(toList());
 
         fieldCount = fileColumnNames.size();

presto-hive/src/main/java/io/prestosql/plugin/hive/orc/OrcFileWriterFactory.java

@@ -64,6 +64,7 @@
 import static io.prestosql.plugin.hive.HiveSessionProperties.getOrcOptimizedWriterMinStripeSize;
 import static io.prestosql.plugin.hive.HiveSessionProperties.getOrcOptimizedWriterValidateMode;
 import static io.prestosql.plugin.hive.HiveSessionProperties.getOrcStringStatisticsLimit;
+import static io.prestosql.plugin.hive.HiveSessionProperties.getTimestampPrecision;
 import static io.prestosql.plugin.hive.HiveSessionProperties.isOrcOptimizedWriterValidate;
 import static io.prestosql.plugin.hive.acid.AcidSchema.ACID_COLUMN_NAMES;
 import static io.prestosql.plugin.hive.acid.AcidSchema.createAcidColumnPrestoTypes;
@@ -148,7 +149,7 @@ public Optional<FileWriter> createFileWriter(
         // an index to rearrange columns in the proper order
         List<String> fileColumnNames = getColumnNames(schema);
         List<Type> fileColumnTypes = getColumnTypes(schema).stream()
-                .map(hiveType -> hiveType.getType(typeManager))
+                .map(hiveType -> hiveType.getType(typeManager, getTimestampPrecision(session)))
                 .collect(toList());
 
         int[] fileInputColumnIndexes = fileColumnNames.stream()

presto-hive/src/main/java/io/prestosql/plugin/hive/parquet/ParquetFileWriterFactory.java

@@ -42,6 +42,7 @@
 import java.util.concurrent.Callable;
 
 import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR;
+import static io.prestosql.plugin.hive.HiveSessionProperties.getTimestampPrecision;
 import static io.prestosql.plugin.hive.util.HiveUtil.getColumnNames;
 import static io.prestosql.plugin.hive.util.HiveUtil.getColumnTypes;
 import static java.util.Objects.requireNonNull;
@@ -91,7 +92,7 @@ public Optional<FileWriter> createFileWriter(
 
         List<String> fileColumnNames = getColumnNames(schema);
         List<Type> fileColumnTypes = getColumnTypes(schema).stream()
-                .map(hiveType -> hiveType.getType(typeManager))
+                .map(hiveType -> hiveType.getType(typeManager, getTimestampPrecision(session)))
                 .collect(toList());
 
         int[] fileInputColumnIndexes = fileColumnNames.stream()

presto-hive/src/main/java/io/prestosql/plugin/hive/rcfile/RcFilePageSourceFactory.java

@@ -24,6 +24,7 @@
 import io.prestosql.plugin.hive.HiveColumnHandle;
 import io.prestosql.plugin.hive.HiveConfig;
 import io.prestosql.plugin.hive.HivePageSourceFactory;
+import io.prestosql.plugin.hive.HiveTimestampPrecision;
 import io.prestosql.plugin.hive.ReaderProjections;
 import io.prestosql.plugin.hive.acid.AcidTransaction;
 import io.prestosql.plugin.hive.util.FSDataInputStreamTail;
@@ -184,7 +185,7 @@ else if (deserializerClassName.equals(ColumnarSerDe.class.getName())) {
 
         try {
             ImmutableMap.Builder<Integer, Type> readColumns = ImmutableMap.builder();
-            int timestampPrecision = getTimestampPrecision(session).getPrecision();
+            HiveTimestampPrecision timestampPrecision = getTimestampPrecision(session);
             for (HiveColumnHandle column : projectedReaderColumns) {
                 readColumns.put(column.getBaseHiveColumnIndex(), column.getHiveType().getType(typeManager, timestampPrecision));
             }

presto-hive/src/main/java/io/prestosql/plugin/hive/util/FieldSetterFactory.java

@@ -24,8 +24,10 @@
 import io.prestosql.spi.type.DecimalType;
 import io.prestosql.spi.type.DoubleType;
 import io.prestosql.spi.type.IntegerType;
+import io.prestosql.spi.type.LongTimestamp;
 import io.prestosql.spi.type.RealType;
 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.VarbinaryType;
@@ -56,10 +58,13 @@
 import static io.prestosql.plugin.hive.util.HiveUtil.isMapType;
 import static io.prestosql.plugin.hive.util.HiveUtil.isRowType;
 import static io.prestosql.plugin.hive.util.HiveWriteUtils.getHiveDecimal;
-import static io.prestosql.spi.type.TimestampType.TIMESTAMP_MILLIS;
-import static io.prestosql.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND;
+import static io.prestosql.spi.type.Timestamps.MICROSECONDS_PER_SECOND;
+import static io.prestosql.spi.type.Timestamps.MILLISECONDS_PER_SECOND;
+import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_MICROSECOND;
+import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND;
 import static java.lang.Float.intBitsToFloat;
 import static java.lang.Math.floorDiv;
+import static java.lang.Math.floorMod;
 import static java.lang.Math.toIntExact;
 import static java.util.Objects.requireNonNull;
 
@@ -118,8 +123,8 @@ public FieldSetter create(SettableStructObjectInspector rowInspector, Object row
             return new DateFieldSetter(rowInspector, row, field);
         }
 
-        if (type.equals(TIMESTAMP_MILLIS)) {
-            return new TimestampFieldSetter(rowInspector, row, field, timeZone);
+        if (type instanceof TimestampType) {
+            return new TimestampFieldSetter(rowInspector, row, field, (TimestampType) type, timeZone);
         }
 
         if (type instanceof DecimalType) {
@@ -365,22 +370,49 @@ private static class TimestampFieldSetter
             extends FieldSetter
     {
         private final DateTimeZone timeZone;
+        private final TimestampType type;
         private final TimestampWritableV2 value = new TimestampWritableV2();
 
-        public TimestampFieldSetter(SettableStructObjectInspector rowInspector, Object row, StructField field, DateTimeZone timeZone)
+        public TimestampFieldSetter(SettableStructObjectInspector rowInspector, Object row, StructField field, TimestampType type, DateTimeZone timeZone)
         {
             super(rowInspector, row, field);
+            this.type = requireNonNull(type, "type is null");
             this.timeZone = requireNonNull(timeZone, "timeZone is null");
         }
 
         @Override
         public void setField(Block block, int position)
         {
-            long epochMilli = floorDiv(TIMESTAMP_MILLIS.getLong(block, position), MICROSECONDS_PER_MILLISECOND);
-            epochMilli = timeZone.convertLocalToUTC(epochMilli, false);
-            value.set(Timestamp.ofEpochMilli(epochMilli));
+            long epochMicros;
+            int picosOfMicro;
+            if (type.isShort()) {
+                epochMicros = type.getLong(block, position);
+                picosOfMicro = 0;
+            }
+            else {
+                LongTimestamp longTimestamp = (LongTimestamp) type.getObject(block, position);
+                epochMicros = longTimestamp.getEpochMicros();
+                picosOfMicro = longTimestamp.getPicosOfMicro();
+            }
+
+            long epochSeconds = floorDiv(epochMicros, MICROSECONDS_PER_SECOND);
+            long picosOfSecond = (long) floorMod(epochMicros, MICROSECONDS_PER_SECOND) * PICOSECONDS_PER_MICROSECOND + picosOfMicro;
+
+            epochSeconds = convertLocalEpochSecondsToUtc(epochSeconds);
+            // no rounding since the the data has nanosecond precision, at most
+            int nanosOfSecond = toIntExact(picosOfSecond / PICOSECONDS_PER_NANOSECOND);
+
+            Timestamp timestamp = Timestamp.ofEpochSecond(epochSeconds, nanosOfSecond);
+            value.set(timestamp);
             rowInspector.setStructFieldData(row, field, value);
         }
+
+        private long convertLocalEpochSecondsToUtc(long epochSeconds)
+        {
+            long epochMillis = epochSeconds * MILLISECONDS_PER_SECOND;
+            epochMillis = timeZone.convertLocalToUTC(epochMillis, false);
+            return epochMillis / MILLISECONDS_PER_SECOND;
+        }
     }
 
     private static class DecimalFieldSetter

presto-hive/src/main/java/io/prestosql/plugin/hive/util/HiveBucketing.java

@@ -22,6 +22,7 @@
 import io.prestosql.plugin.hive.HiveBucketProperty;
 import io.prestosql.plugin.hive.HiveColumnHandle;
 import io.prestosql.plugin.hive.HiveTableHandle;
+import io.prestosql.plugin.hive.HiveTimestampPrecision;
 import io.prestosql.plugin.hive.HiveType;
 import io.prestosql.plugin.hive.metastore.Column;
 import io.prestosql.plugin.hive.metastore.Table;
@@ -184,7 +185,7 @@ public static Optional<HiveBucketHandle> getHiveBucketHandle(ConnectorSession se
             return Optional.empty();
         }
 
-        int timestampPrecision = getTimestampPrecision(session).getPrecision();
+        HiveTimestampPrecision timestampPrecision = getTimestampPrecision(session);
         Map<String, HiveColumnHandle> map = getRegularColumnHandles(table, typeManager, timestampPrecision).stream()
                 .collect(Collectors.toMap(HiveColumnHandle::getName, identity()));