Skip to content

Commit

Permalink
Add support for precision for TIMESTAMP in PostgreSQL type mapping
Browse files Browse the repository at this point in the history 
Add support for mapping TIMESTAMP from Presto to
PostgreSQL (and vice versa) for precisions 1 to 6. It covers all
precisions currently supported by PostgreSQL.
  • Loading branch information
@kokosing @losipiuk
kokosing authored and losipiuk committed Sep 10, 2020
1 parent 91875da commit e7fdfe3
Show file tree
Hide file tree
Showing 6 changed files with 143 additions and 62 deletions.
54 changes: 38 additions & 16 deletions presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/StandardColumnMappings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import io.prestosql.spi.type.CharType;
import io.prestosql.spi.type.DecimalType;
import io.prestosql.spi.type.Decimals;
import io.prestosql.spi.type.TimestampType;
import io.prestosql.spi.type.VarcharType;
import org.joda.time.DateTimeZone;
import org.joda.time.chrono.ISOChronology;
Expand Down Expand Up @@ -55,9 +56,11 @@
import static io.prestosql.spi.type.RealType.REAL;
import static io.prestosql.spi.type.SmallintType.SMALLINT;
import static io.prestosql.spi.type.TimeType.TIME;
import static io.prestosql.spi.type.TimestampType.MAX_SHORT_PRECISION;
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.NANOSECONDS_PER_DAY;
import static io.prestosql.spi.type.Timestamps.NANOSECONDS_PER_MICROSECOND;
import static io.prestosql.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_DAY;
import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_MILLISECOND;
Expand All @@ -70,6 +73,7 @@
import static java.lang.Float.floatToRawIntBits;
import static java.lang.Float.intBitsToFloat;
import static java.lang.Math.floorDiv;
import static java.lang.Math.floorMod;
import static java.lang.Math.max;
import static java.lang.Math.min;
import static java.lang.Math.toIntExact;
Expand Down Expand Up @@ -351,28 +355,38 @@ public static LongWriteFunction timeWriteFunction()
* {@link #timestampColumnMapping} instead.
*/
@Deprecated
public static ColumnMapping timestampColumnMappingUsingSqlTimestamp()
public static ColumnMapping timestampColumnMappingUsingSqlTimestamp(TimestampType timestampType)
{
// TODO support higher precision
checkArgument(timestampType.getPrecision() <= MAX_SHORT_PRECISION, "Precision is out of range: %s", timestampType.getPrecision());
return ColumnMapping.longMapping(
TIMESTAMP_MILLIS,
(resultSet, columnIndex) -> {
Timestamp timestamp = resultSet.getTimestamp(columnIndex);
return toPrestoTimestamp(timestamp.toLocalDateTime());
return toPrestoTimestamp(timestampType, timestamp.toLocalDateTime());
},
timestampWriteFunctionUsingSqlTimestamp());
timestampWriteFunctionUsingSqlTimestamp(timestampType));
}

@Deprecated
public static ColumnMapping timestampColumnMapping()
{
return timestampColumnMapping(TIMESTAMP_MILLIS);
}

public static ColumnMapping timestampColumnMapping(TimestampType timestampType)
{
checkArgument(timestampType.getPrecision() <= MAX_SHORT_PRECISION, "Precision is out of range: %s", timestampType.getPrecision());
return ColumnMapping.longMapping(
TIMESTAMP_MILLIS,
timestampReadFunction(),
timestampWriteFunction());
timestampType,
timestampReadFunction(timestampType),
timestampWriteFunction(timestampType));
}

public static LongReadFunction timestampReadFunction()
public static LongReadFunction timestampReadFunction(TimestampType timestampType)
{
return (resultSet, columnIndex) -> toPrestoTimestamp(resultSet.getObject(columnIndex, LocalDateTime.class));
checkArgument(timestampType.getPrecision() <= MAX_SHORT_PRECISION, "Precision is out of range: %s", timestampType.getPrecision());
return (resultSet, columnIndex) -> toPrestoTimestamp(timestampType, resultSet.getObject(columnIndex, LocalDateTime.class));
}

/**
Expand All @@ -382,24 +396,32 @@ public static LongReadFunction timestampReadFunction()
* {@link #timestampWriteFunction} instead.
*/
@Deprecated
public static LongWriteFunction timestampWriteFunctionUsingSqlTimestamp()
public static LongWriteFunction timestampWriteFunctionUsingSqlTimestamp(TimestampType timestampType)
{
checkArgument(timestampType.getPrecision() <= MAX_SHORT_PRECISION, "Precision is out of range: %s", timestampType.getPrecision());
return (statement, index, value) -> statement.setTimestamp(index, Timestamp.valueOf(fromPrestoTimestamp(value)));
}

public static LongWriteFunction timestampWriteFunction()
public static LongWriteFunction timestampWriteFunction(TimestampType timestampType)
{
checkArgument(timestampType.getPrecision() <= MAX_SHORT_PRECISION, "Precision is out of range: %s", timestampType.getPrecision());
return (statement, index, value) -> statement.setObject(index, fromPrestoTimestamp(value));
}

public static long toPrestoTimestamp(LocalDateTime localDateTime)
public static long toPrestoTimestamp(TimestampType timestampType, LocalDateTime localDateTime)
{
return localDateTime.atZone(UTC).toInstant().toEpochMilli() * MICROSECONDS_PER_MILLISECOND;
long precision = timestampType.getPrecision();
checkArgument(precision <= MAX_SHORT_PRECISION, "Precision is out of range: %s", precision);
Instant instant = localDateTime.atZone(UTC).toInstant();
return instant.getEpochSecond() * MICROSECONDS_PER_SECOND + roundDiv(instant.getNano(), NANOSECONDS_PER_MICROSECOND);
}

public static LocalDateTime fromPrestoTimestamp(long value)
public static LocalDateTime fromPrestoTimestamp(long epochMicros)
{
return Instant.ofEpochMilli(floorDiv(value, MICROSECONDS_PER_MILLISECOND)).atZone(UTC).toLocalDateTime();
long epochSecond = floorDiv(epochMicros, MICROSECONDS_PER_SECOND);
int nanoFraction = floorMod(epochMicros, MICROSECONDS_PER_SECOND) * NANOSECONDS_PER_MICROSECOND;
Instant instant = Instant.ofEpochSecond(epochSecond, nanoFraction);
return LocalDateTime.ofInstant(instant, UTC);
}

public static LocalTime fromPrestoTime(long value)
Expand Down Expand Up @@ -473,7 +495,7 @@ public static Optional<ColumnMapping> jdbcTypeToPrestoType(JdbcTypeHandle type)

case Types.TIMESTAMP:
// TODO default to `timestampColumnMapping`
return Optional.of(timestampColumnMappingUsingSqlTimestamp());
return Optional.of(timestampColumnMappingUsingSqlTimestamp(TIMESTAMP_MILLIS));
}
return Optional.empty();
}
Expand Down
2 changes: 1 addition & 1 deletion presto-mysql/src/main/java/io/prestosql/plugin/mysql/MySqlClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
}
if (TIMESTAMP_MILLIS.equals(type)) {
// TODO use `timestampWriteFunction`
return WriteMapping.longMapping("datetime", timestampWriteFunctionUsingSqlTimestamp());
return WriteMapping.longMapping("datetime", timestampWriteFunctionUsingSqlTimestamp(TIMESTAMP_MILLIS));
}
if (VARBINARY.equals(type)) {
return WriteMapping.sliceMapping("mediumblob", varbinaryWriteFunction());
Expand Down
17 changes: 10 additions & 7 deletions presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
import io.prestosql.spi.type.LongTimestampWithTimeZone;
import io.prestosql.spi.type.MapType;
import io.prestosql.spi.type.StandardTypes;
import io.prestosql.spi.type.TimestampType;
import io.prestosql.spi.type.TimestampWithTimeZoneType;
import io.prestosql.spi.type.TinyintType;
import io.prestosql.spi.type.Type;
Expand Down Expand Up @@ -134,7 +135,7 @@
import static io.prestosql.spi.type.StandardTypes.JSON;
import static io.prestosql.spi.type.TimeType.TIME;
import static io.prestosql.spi.type.TimeZoneKey.UTC_KEY;
import static io.prestosql.spi.type.TimestampType.TIMESTAMP_MILLIS;
import static io.prestosql.spi.type.TimestampType.createTimestampType;
import static io.prestosql.spi.type.TimestampWithTimeZoneType.createTimestampWithTimeZoneType;
import static io.prestosql.spi.type.Timestamps.MILLISECONDS_PER_SECOND;
import static io.prestosql.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
Expand Down Expand Up @@ -360,10 +361,11 @@ public Optional<ColumnMapping> toPrestoType(ConnectorSession session, Connection
return Optional.of(timeColumnMappingWithTruncation());
}
if (typeHandle.getJdbcType() == Types.TIMESTAMP) {
TimestampType timestampType = createTimestampType(typeHandle.getDecimalDigits());
return Optional.of(ColumnMapping.longMapping(
TIMESTAMP_MILLIS,
timestampReadFunction(),
timestampWriteFunction()));
timestampType,
timestampReadFunction(timestampType),
timestampWriteFunction(timestampType)));
}
if (typeHandle.getJdbcType() == Types.NUMERIC && getDecimalRounding(session) == ALLOW_OVERFLOW) {
if (typeHandle.getColumnSize() == 131089) {
Expand Down Expand Up @@ -439,8 +441,9 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
if (TIME.equals(type)) {
return WriteMapping.longMapping("time", timeWriteFunction());
}
if (TIMESTAMP_MILLIS.equals(type)) {
return WriteMapping.longMapping("timestamp", timestampWriteFunction());
if (type instanceof TimestampType && ((TimestampType) type).getPrecision() <= MAX_SUPPORTED_TIMESTAMP_PRECISION) {
TimestampType timestampType = (TimestampType) type;
return WriteMapping.longMapping(format("timestamp(%s)", timestampType.getPrecision()), timestampWriteFunction(timestampType));
}
if (type instanceof TimestampWithTimeZoneType && ((TimestampWithTimeZoneType) type).getPrecision() <= MAX_SUPPORTED_TIMESTAMP_PRECISION) {
int precision = ((TimestampWithTimeZoneType) type).getPrecision();
Expand Down Expand Up @@ -496,7 +499,7 @@ public boolean isLimitGuaranteed(ConnectorSession session)
// When writing with setObject() using LocalDateTime, driver converts the value to string representing date-time in JVM zone,
// therefore cannot represent local date-time which is a "gap" in this zone.
// TODO replace this method with StandardColumnMappings#timestampWriteFunction when https://github.com/pgjdbc/pgjdbc/issues/1390 is done
private static LongWriteFunction timestampWriteFunction()
private static LongWriteFunction timestampWriteFunction(TimestampType timestampType)
{
return (statement, index, value) -> {
LocalDateTime localDateTime = fromPrestoTimestamp(value);
Expand Down
11 changes: 9 additions & 2 deletions presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/TypeUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import io.prestosql.spi.type.CharType;
import io.prestosql.spi.type.DecimalType;
import io.prestosql.spi.type.LongTimestampWithTimeZone;
import io.prestosql.spi.type.TimestampType;
import io.prestosql.spi.type.TimestampWithTimeZoneType;
import io.prestosql.spi.type.Type;
import io.prestosql.spi.type.VarcharType;
Expand Down Expand Up @@ -53,7 +54,6 @@
import static io.prestosql.spi.type.RealType.REAL;
import static io.prestosql.spi.type.SmallintType.SMALLINT;
import static io.prestosql.spi.type.TimeZoneKey.UTC_KEY;
import static io.prestosql.spi.type.TimestampType.TIMESTAMP_MILLIS;
import static io.prestosql.spi.type.Timestamps.MILLISECONDS_PER_SECOND;
import static io.prestosql.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND;
Expand Down Expand Up @@ -92,6 +92,12 @@ static String getArrayElementPgTypeName(ConnectorSession session, PostgreSqlClie
return "timestamptz";
}

// TypeInfoCache#getPGArrayType doesn't allow for specifying variable-length limits.
// Map all timestamp(x) types to unparametrized one which defaults to highest precision
if (elementType instanceof TimestampType) {
return "timestamp";
}

if (elementType instanceof DecimalType) {
return "decimal";
}
Expand Down Expand Up @@ -189,7 +195,8 @@ private static Object prestoNativeToJdbcObject(ConnectorSession session, Type pr
return new Date(UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millis));
}

if (TIMESTAMP_MILLIS.equals(prestoType)) {
if (prestoType instanceof TimestampType && ((TimestampType) prestoType).isShort()) {
TimestampType timestampType = (TimestampType) prestoType;
return toPgTimestamp(fromPrestoTimestamp((long) prestoNative));
}

Expand Down
Loading

0 comments on commit e7fdfe3

Please sign in to comment.