Skip to content

Commit

Permalink
Add timestamp to date coercion for hive
Browse files Browse the repository at this point in the history
  • Loading branch information
@findinpath
findinpath committed Nov 9, 2023
1 parent ecbc5b1 commit 2bf4f5f
Show file tree
Hide file tree
Showing 8 changed files with 167 additions and 26 deletions.
11 changes: 9 additions & 2 deletions plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/CoercionUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import io.trino.plugin.hive.HiveType;
import io.trino.plugin.hive.coercions.BooleanCoercer.BooleanToVarcharCoercer;
import io.trino.plugin.hive.coercions.DateCoercer.VarcharToDateCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToDateCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToVarcharCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.VarcharToLongTimestampCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.VarcharToShortTimestampCoercer;
Expand Down Expand Up @@ -189,8 +190,14 @@ public static Type createTypeFromCoercer(TypeManager typeManager, HiveType fromH
if (fromType == REAL && toType instanceof DecimalType toDecimalType) {
return Optional.of(createRealToDecimalCoercer(toDecimalType));
}
if (fromType instanceof TimestampType && toType instanceof VarcharType varcharType) {
return Optional.of(new LongTimestampToVarcharCoercer(TIMESTAMP_NANOS, varcharType));
if (fromType instanceof TimestampType) {
if (toType instanceof VarcharType varcharType) {
return Optional.of(new LongTimestampToVarcharCoercer(TIMESTAMP_NANOS, varcharType));
}
if (toType instanceof DateType toDateType) {
return Optional.of(new LongTimestampToDateCoercer(TIMESTAMP_NANOS, toDateType));
}
return Optional.empty();
}
if (fromType == DOUBLE && toType instanceof VarcharType toVarcharType) {
return Optional.of(new DoubleToVarcharCoercer(toVarcharType, coercionContext.treatNaNAsNull()));
Expand Down
22 changes: 22 additions & 0 deletions plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/TimestampCoercer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import io.trino.spi.TrinoException;
import io.trino.spi.block.Block;
import io.trino.spi.block.BlockBuilder;
import io.trino.spi.type.DateType;
import io.trino.spi.type.LongTimestamp;
import io.trino.spi.type.TimestampType;
import io.trino.spi.type.VarcharType;
Expand Down Expand Up @@ -94,6 +95,27 @@ protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int pos
}
}

public static class LongTimestampToDateCoercer
extends TypeCoercer<TimestampType, DateType>
{
public LongTimestampToDateCoercer(TimestampType fromType, DateType toType)
{
super(fromType, toType);
}

@Override
protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position)
{
LongTimestamp timestamp = (LongTimestamp) fromType.getObject(block, position);

long epochSecond = floorDiv(timestamp.getEpochMicros(), MICROSECONDS_PER_SECOND);
if (epochSecond < START_OF_MODERN_ERA_SECONDS) {
throw new TrinoException(HIVE_INVALID_TIMESTAMP_COERCION, "Coercion on historical dates is not supported");
}
toType.writeLong(blockBuilder, floorDiv(epochSecond, SECONDS_PER_DAY));
}
}

public static class VarcharToShortTimestampCoercer
extends TypeCoercer<VarcharType, TimestampType>
{
Expand Down
11 changes: 9 additions & 2 deletions plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcTypeTranslator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import io.trino.plugin.hive.coercions.DateCoercer.VarcharToDateCoercer;
import io.trino.plugin.hive.coercions.DoubleToVarcharCoercer;
import io.trino.plugin.hive.coercions.IntegerNumberToDoubleCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToDateCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToVarcharCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.VarcharToLongTimestampCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.VarcharToShortTimestampCoercer;
Expand Down Expand Up @@ -53,8 +54,14 @@ private OrcTypeTranslator() {}

public static Optional<TypeCoercer<? extends Type, ? extends Type>> createCoercer(OrcTypeKind fromOrcType, Type toTrinoType)
{
if (fromOrcType == TIMESTAMP && toTrinoType instanceof VarcharType varcharType) {
return Optional.of(new LongTimestampToVarcharCoercer(TIMESTAMP_NANOS, varcharType));
if (fromOrcType == TIMESTAMP) {
if (toTrinoType instanceof VarcharType varcharType) {
return Optional.of(new LongTimestampToVarcharCoercer(TIMESTAMP_NANOS, varcharType));
}
if (toTrinoType instanceof DateType toDateType) {
return Optional.of(new LongTimestampToDateCoercer(TIMESTAMP_NANOS, toDateType));
}
return Optional.empty();
}
if (isVarcharType(fromOrcType)) {
if (toTrinoType instanceof TimestampType timestampType) {
Expand Down
3 changes: 3 additions & 0 deletions plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveCoercionPolicy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,9 @@ private boolean canCoerce(HiveType fromHiveType, HiveType toHiveType, HiveTimest
fromHiveType.equals(HIVE_DOUBLE) ||
fromType instanceof DecimalType;
}
if (toHiveType.equals(HIVE_DATE)) {
return fromHiveType.equals(HIVE_TIMESTAMP);
}
if (fromHiveType.equals(HIVE_BYTE)) {
return toHiveType.equals(HIVE_SHORT) || toHiveType.equals(HIVE_INT) || toHiveType.equals(HIVE_LONG) || toHiveType.equals(HIVE_DOUBLE);
}
Expand Down
44 changes: 44 additions & 0 deletions plugin/trino-hive/src/test/java/io/trino/plugin/hive/coercions/TestTimestampCoercer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import io.trino.spi.type.VarcharType;
import org.junit.jupiter.api.Test;

import java.time.LocalDate;
import java.time.LocalDateTime;

import static io.airlift.slice.Slices.utf8Slice;
Expand All @@ -33,6 +34,7 @@
import static io.trino.plugin.hive.coercions.CoercionUtils.createCoercer;
import static io.trino.spi.predicate.Utils.blockToNativeValue;
import static io.trino.spi.predicate.Utils.nativeValueToBlock;
import static io.trino.spi.type.DateType.DATE;
import static io.trino.spi.type.TimestampType.TIMESTAMP_MICROS;
import static io.trino.spi.type.TimestampType.TIMESTAMP_PICOS;
import static io.trino.spi.type.VarcharType.createUnboundedVarcharType;
Expand All @@ -45,6 +47,41 @@

public class TestTimestampCoercer
{
@Test
public void testTimestampToDate()
{
// before epoch
assertTimestampToDateCoercion("1900-01-01T00:00:00.000", "1900-01-01");
assertTimestampToDateCoercion("1958-01-01T13:18:03.123", "1958-01-01");
// after epoch
assertTimestampToDateCoercion("2019-03-18T10:01:17.987", "2019-03-18");
assertTimestampToDateCoercion("2021-12-31T23:59:59.000", "2021-12-31");
assertTimestampToDateCoercion("2021-12-31T23:59:59.999", "2021-12-31");
assertTimestampToDateCoercion("2021-12-31T23:59:59.999999", "2021-12-31");
assertTimestampToDateCoercion("2021-12-31T23:59:59.999999999", "2021-12-31");
// time doubled in JVM zone
assertTimestampToDateCoercion("2018-10-28T01:33:17.456", "2018-10-28");
// epoch
assertTimestampToDateCoercion("1970-01-01T00:00:00.000", "1970-01-01");
// time gap in JVM zone
assertTimestampToDateCoercion("1970-01-01T00:13:42.000", "1970-01-01");
assertTimestampToDateCoercion("2018-04-01T02:13:55.123", "2018-04-01");
// time gap in Vilnius
assertTimestampToDateCoercion("2018-03-25T03:17:17.000", "2018-03-25");
// time gap in Kathmandu
assertTimestampToDateCoercion("1986-01-01T00:13:07.000", "1986-01-01");
// before epoch with second fraction
assertTimestampToDateCoercion("1969-12-31T23:59:59.123456", "1969-12-31");
}

@Test
public void testHistoricalLongTimestampToDate()
{
assertThatThrownBy(() -> assertTimestampToDateCoercion("1899-12-31T23:59:59.999999999", "1899-12-31"))
.isInstanceOf(TrinoException.class)
.hasMessageContaining("Coercion on historical dates is not supported");
}

@Test
public void testTimestampToVarchar()
{
Expand Down Expand Up @@ -277,4 +314,11 @@ public static void assertCoercions(Type fromType, Object valueToBeCoerced, Type
assertThat(blockToNativeValue(toType, coercedValue))
.isEqualTo(expectedValue);
}

private static void assertTimestampToDateCoercion(String timestampAsString, String expectedDate)
{
LocalDateTime localDateTime = LocalDateTime.parse(timestampAsString);
SqlTimestamp timestamp = SqlTimestamp.fromSeconds(TIMESTAMP_PICOS.getPrecision(), localDateTime.toEpochSecond(UTC), localDateTime.get(NANO_OF_SECOND));
assertCoercions(TIMESTAMP_PICOS, new LongTimestamp(timestamp.getEpochMicros(), timestamp.getPicosOfMicros()), DATE, LocalDate.parse(expectedDate).toEpochDay(), NANOSECONDS);
}
}
Loading

0 comments on commit 2bf4f5f

Please sign in to comment.