Skip to content

Commit 425c3b0

Browse files
findepifindepi
committed
Support CTAS with timestamp(>6) in PostgreSQL
1 parent 50a87c7 commit 425c3b0

File tree

2 files changed

+111
-11
lines changed

2 files changed

+111
-11
lines changed

‎presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java‎

Lines changed: 32 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import io.prestosql.spi.type.ArrayType;
6161
import io.prestosql.spi.type.DecimalType;
6262
import io.prestosql.spi.type.Decimals;
63+
import io.prestosql.spi.type.LongTimestamp;
6364
import io.prestosql.spi.type.LongTimestampWithTimeZone;
6465
import io.prestosql.spi.type.MapType;
6566
import io.prestosql.spi.type.StandardTypes;
@@ -139,6 +140,7 @@
139140
import static io.prestosql.spi.type.TimestampWithTimeZoneType.createTimestampWithTimeZoneType;
140141
import static io.prestosql.spi.type.Timestamps.MILLISECONDS_PER_SECOND;
141142
import static io.prestosql.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND;
143+
import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_MICROSECOND;
142144
import static io.prestosql.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND;
143145
import static io.prestosql.spi.type.TypeSignature.mapType;
144146
import static io.prestosql.spi.type.VarbinaryType.VARBINARY;
@@ -158,7 +160,7 @@ public class PostgreSqlClient
158160
*/
159161
private static final int ARRAY_RESULT_SET_VALUE_COLUMN = 2;
160162
private static final String DUPLICATE_TABLE_SQLSTATE = "42P07";
161-
private static final int MAX_SUPPORTED_TIMESTAMP_PRECISION = 6;
163+
private static final int POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION = 6;
162164

163165
private final Type jsonType;
164166
private final Type uuidType;
@@ -367,7 +369,7 @@ public Optional<ColumnMapping> toPrestoType(ConnectorSession session, Connection
367369
return Optional.of(ColumnMapping.longMapping(
368370
timestampType,
369371
timestampReadFunction(timestampType),
370-
timestampWriteFunction(timestampType)));
372+
PostgreSqlClient::shortTimestampWriteFunction));
371373
}
372374
if (typeHandle.getJdbcType() == Types.NUMERIC && getDecimalRounding(session) == ALLOW_OVERFLOW) {
373375
if (typeHandle.getColumnSize() == 131089) {
@@ -444,11 +446,16 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
444446
if (TIME.equals(type)) {
445447
return WriteMapping.longMapping("time", timeWriteFunction());
446448
}
447-
if (type instanceof TimestampType && ((TimestampType) type).getPrecision() <= MAX_SUPPORTED_TIMESTAMP_PRECISION) {
449+
if (type instanceof TimestampType) {
448450
TimestampType timestampType = (TimestampType) type;
449-
return WriteMapping.longMapping(format("timestamp(%s)", timestampType.getPrecision()), timestampWriteFunction(timestampType));
451+
if (timestampType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
452+
verify(timestampType.getPrecision() <= TimestampType.MAX_SHORT_PRECISION);
453+
return WriteMapping.longMapping(format("timestamp(%s)", timestampType.getPrecision()), PostgreSqlClient::shortTimestampWriteFunction);
454+
}
455+
verify(timestampType.getPrecision() > TimestampType.MAX_SHORT_PRECISION);
456+
return WriteMapping.objectMapping(format("timestamp(%s)", POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION), longTimestampWriteFunction());
450457
}
451-
if (type instanceof TimestampWithTimeZoneType && ((TimestampWithTimeZoneType) type).getPrecision() <= MAX_SUPPORTED_TIMESTAMP_PRECISION) {
458+
if (type instanceof TimestampWithTimeZoneType && ((TimestampWithTimeZoneType) type).getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
452459
int precision = ((TimestampWithTimeZoneType) type).getPrecision();
453460
String postgresType = format("timestamptz(%d)", precision);
454461
if (precision <= TimestampWithTimeZoneType.MAX_SHORT_PRECISION) {
@@ -502,12 +509,26 @@ public boolean isLimitGuaranteed(ConnectorSession session)
502509
// When writing with setObject() using LocalDateTime, driver converts the value to string representing date-time in JVM zone,
503510
// therefore cannot represent local date-time which is a "gap" in this zone.
504511
// TODO replace this method with StandardColumnMappings#timestampWriteFunction when https://github.com/pgjdbc/pgjdbc/issues/1390 is done
505-
private static LongWriteFunction timestampWriteFunction(TimestampType timestampType)
512+
private static void shortTimestampWriteFunction(PreparedStatement statement, int index, long epochMicros)
513+
throws SQLException
506514
{
507-
return (statement, index, value) -> {
508-
LocalDateTime localDateTime = fromPrestoTimestamp(value);
509-
statement.setObject(index, toPgTimestamp(localDateTime));
510-
};
515+
LocalDateTime localDateTime = fromPrestoTimestamp(epochMicros);
516+
statement.setObject(index, toPgTimestamp(localDateTime));
517+
}
518+
519+
private static ObjectWriteFunction longTimestampWriteFunction()
520+
{
521+
return ObjectWriteFunction.of(LongTimestamp.class, ((statement, index, timestamp) -> {
522+
// PostgreSQL supports up to 6 digits of precision
523+
//noinspection ConstantConditions
524+
verify(POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION == 6);
525+
526+
long epochMicros = timestamp.getEpochMicros();
527+
if (timestamp.getPicosOfMicro() >= PICOSECONDS_PER_MICROSECOND / 2) {
528+
epochMicros++;
529+
}
530+
shortTimestampWriteFunction(statement, index, epochMicros);
531+
}));
511532
}
512533

513534
@Override
@@ -524,7 +545,7 @@ public void setColumnComment(JdbcIdentity identity, JdbcTableHandle handle, Jdbc
524545
private static ColumnMapping timestampWithTimeZoneColumnMapping(int precision)
525546
{
526547
// PosgreSQL supports timestamptz precision up to microseconds
527-
checkArgument(precision <= MAX_SUPPORTED_TIMESTAMP_PRECISION, "unsupported precision value %d", precision);
548+
checkArgument(precision <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION, "unsupported precision value %d", precision);
528549
TimestampWithTimeZoneType prestoType = createTimestampWithTimeZoneType(precision);
529550
if (precision <= TimestampWithTimeZoneType.MAX_SHORT_PRECISION) {
530551
return ColumnMapping.longMapping(

‎presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343

4444
import java.math.BigDecimal;
4545
import java.math.RoundingMode;
46+
import java.sql.SQLException;
4647
import java.text.NumberFormat;
4748
import java.time.LocalDate;
4849
import java.time.LocalDateTime;
@@ -100,6 +101,7 @@
100101
import static io.prestosql.testing.datatype.DataType.timestampDataType;
101102
import static io.prestosql.testing.datatype.DataType.varbinaryDataType;
102103
import static io.prestosql.testing.datatype.DataType.varcharDataType;
104+
import static io.prestosql.testing.sql.TestTable.randomTableSuffix;
103105
import static io.prestosql.type.JsonType.JSON;
104106
import static io.prestosql.type.UuidType.UUID;
105107
import static java.lang.String.format;
@@ -113,6 +115,7 @@
113115
import static java.util.function.Function.identity;
114116
import static java.util.stream.Collectors.joining;
115117
import static java.util.stream.Collectors.toList;
118+
import static org.assertj.core.api.Assertions.assertThat;
116119

117120
public class TestPostgreSqlTypeMapping
118121
extends AbstractTestQueryFramework
@@ -1007,6 +1010,9 @@ public void testTime(boolean insertWithPresto, ZoneId sessionZone)
10071010
}
10081011
}
10091012

1013+
/**
1014+
* @see #testTimestampCoercion
1015+
*/
10101016
@Test(dataProvider = "testTimestampDataProvider")
10111017
public void testTimestamp(boolean insertWithPresto, ZoneId sessionZone)
10121018
{
@@ -1052,6 +1058,60 @@ public void testTimestamp(boolean insertWithPresto, ZoneId sessionZone)
10521058
}
10531059
}
10541060

1061+
/**
1062+
* Additional test supplementing {@link #testTimestamp} with values that do not necessarily round-trip, including
1063+
* timestamp precision higher than expressible with {@code LocalDateTime}.
1064+
*
1065+
* @see #testTimestamp
1066+
*/
1067+
@Test
1068+
public void testTimestampCoercion()
1069+
throws SQLException
1070+
{
1071+
// precision 0 ends up as precision 0
1072+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00'", "TIMESTAMP '1970-01-01 00:00:00'");
1073+
1074+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.1'", "TIMESTAMP '1970-01-01 00:00:00.1'");
1075+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.9'", "TIMESTAMP '1970-01-01 00:00:00.9'");
1076+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.123'", "TIMESTAMP '1970-01-01 00:00:00.123'");
1077+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.123000'", "TIMESTAMP '1970-01-01 00:00:00.123000'");
1078+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.999'", "TIMESTAMP '1970-01-01 00:00:00.999'");
1079+
// max supported precision
1080+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.123456'", "TIMESTAMP '1970-01-01 00:00:00.123456'");
1081+
1082+
testCreateTableAsAndInsertConsistency("TIMESTAMP '2020-09-27 12:34:56.1'", "TIMESTAMP '2020-09-27 12:34:56.1'");
1083+
testCreateTableAsAndInsertConsistency("TIMESTAMP '2020-09-27 12:34:56.9'", "TIMESTAMP '2020-09-27 12:34:56.9'");
1084+
testCreateTableAsAndInsertConsistency("TIMESTAMP '2020-09-27 12:34:56.123'", "TIMESTAMP '2020-09-27 12:34:56.123'");
1085+
testCreateTableAsAndInsertConsistency("TIMESTAMP '2020-09-27 12:34:56.123000'", "TIMESTAMP '2020-09-27 12:34:56.123000'");
1086+
testCreateTableAsAndInsertConsistency("TIMESTAMP '2020-09-27 12:34:56.999'", "TIMESTAMP '2020-09-27 12:34:56.999'");
1087+
// max supported precision
1088+
testCreateTableAsAndInsertConsistency("TIMESTAMP '2020-09-27 12:34:56.123456'", "TIMESTAMP '2020-09-27 12:34:56.123456'");
1089+
1090+
// round down
1091+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.1234561'", "TIMESTAMP '1970-01-01 00:00:00.123456'");
1092+
1093+
// nanoc round up, end result rounds down
1094+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.123456499'", "TIMESTAMP '1970-01-01 00:00:00.123456'");
1095+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.123456499999'", "TIMESTAMP '1970-01-01 00:00:00.123456'");
1096+
1097+
// round up
1098+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.1234565'", "TIMESTAMP '1970-01-01 00:00:00.123457'");
1099+
1100+
// max precision
1101+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.111222333444'", "TIMESTAMP '1970-01-01 00:00:00.111222'");
1102+
1103+
// round up to next second
1104+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 00:00:00.9999995'", "TIMESTAMP '1970-01-01 00:00:01.000000'");
1105+
1106+
// round up to next day
1107+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1970-01-01 23:59:59.9999995'", "TIMESTAMP '1970-01-02 00:00:00.000000'");
1108+
1109+
// negative epoch
1110+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1969-12-31 23:59:59.9999995'", "TIMESTAMP '1970-01-01 00:00:00.000000'");
1111+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1969-12-31 23:59:59.999999499999'", "TIMESTAMP '1969-12-31 23:59:59.999999'");
1112+
testCreateTableAsAndInsertConsistency("TIMESTAMP '1969-12-31 23:59:59.9999994'", "TIMESTAMP '1969-12-31 23:59:59.999999'");
1113+
}
1114+
10551115
@Test(dataProvider = "testTimestampDataProvider")
10561116
public void testArrayTimestamp(boolean insertWithPresto, ZoneId sessionZone)
10571117
{
@@ -1421,6 +1481,25 @@ private void testUnsupportedDataTypeConvertedToVarchar(Session session, String d
14211481
}
14221482
}
14231483

1484+
private void testCreateTableAsAndInsertConsistency(String inputLiteral, String expectedResult)
1485+
throws SQLException
1486+
{
1487+
String tableName = "test_ctas_and_insert_" + randomTableSuffix();
1488+
1489+
// CTAS
1490+
assertUpdate("CREATE TABLE " + tableName + " AS SELECT " + inputLiteral + " a", 1);
1491+
assertThat(query("SELECT a FROM " + tableName))
1492+
.matches("VALUES " + expectedResult);
1493+
1494+
// INSERT as a control query, where the coercion is done by the engine
1495+
postgreSqlServer.execute("DELETE FROM tpch." + tableName);
1496+
assertUpdate("INSERT INTO " + tableName + " (a) VALUES (" + inputLiteral + ")", 1);
1497+
assertThat(query("SELECT a FROM " + tableName))
1498+
.matches("VALUES " + expectedResult);
1499+
1500+
assertUpdate("DROP TABLE " + tableName);
1501+
}
1502+
14241503
public static DataType<ZonedDateTime> timestampWithTimeZoneDataType(int precision, boolean insertWithPresto)
14251504
{
14261505
if (insertWithPresto) {

0 commit comments

Comments
 (0)