diff --git a/athena-timestream/src/main/java/com/amazonaws/athena/connectors/timestream/query/PredicateBuilder.java b/athena-timestream/src/main/java/com/amazonaws/athena/connectors/timestream/query/PredicateBuilder.java index dd191458d2..e66f3e3375 100644 --- a/athena-timestream/src/main/java/com/amazonaws/athena/connectors/timestream/query/PredicateBuilder.java +++ b/athena-timestream/src/main/java/com/amazonaws/athena/connectors/timestream/query/PredicateBuilder.java @@ -30,12 +30,23 @@ import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; +import java.time.LocalDateTime; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; +import java.time.format.DateTimeFormatterBuilder; import java.util.ArrayList; import java.util.List; import java.util.stream.Collectors; public class PredicateBuilder { + // We use a specific format to use the full precision provided by Timestream. + // Additionally, it prevents generating invalid format like `2024-12-31 00:11:22.`. + private static final DateTimeFormatter TIMESTAMP_FORMATTER = new DateTimeFormatterBuilder() + .appendPattern("yyyy-MM-dd HH:mm:ss.nnnnnnnnn") + .toFormatter() + .withZone(ZoneId.of("UTC")); + private PredicateBuilder() {} public static List buildConjucts( @@ -157,6 +168,8 @@ private static String quoteValue(Object value, ArrowType type) switch (Types.getMinorTypeForArrowType(type)) { case VARCHAR: return "\'" + value + "\'"; + case DATEMILLI: + return "\'" + ((LocalDateTime) value).format(TIMESTAMP_FORMATTER) + "\'"; default: return String.valueOf(value); } diff --git a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/query/SelectQueryBuilderTest.java b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/query/SelectQueryBuilderTest.java index 2ce992b582..2a4a86e7b8 100644 --- a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/query/SelectQueryBuilderTest.java +++ b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/query/SelectQueryBuilderTest.java @@ -36,9 +36,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; +import java.time.LocalDateTime; +import java.util.*; import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT; import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.VIEW_METADATA_FIELD; @@ -104,6 +103,37 @@ public void build() logger.info("build: exit"); } + @Test + public void buildWithTime() { + logger.info("build: enter"); + + String expected = "SELECT val FROM \"myDatabase\".\"myTable\" WHERE ((\"time1\" > '2024-04-05 09:31:12.000000000')) AND ((\"time0\" > '2024-04-05 09:31:12.142000000'))"; + + Map constraintsMap = new HashMap<>(); + constraintsMap.put("time0", SortedRangeSet.copyOf(Types.MinorType.DATEMILLI.getType(), + ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.DATEMILLI.getType(), + LocalDateTime.of(2024, 4, 5, 9, 31, 12, 142000000))), false)); + constraintsMap.put("time1", SortedRangeSet.copyOf(Types.MinorType.DATEMILLI.getType(), + ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.DATEMILLI.getType(), + LocalDateTime.of(2024, 4, 5, 9, 31, 12))), false)); + + Schema schema = SchemaBuilder.newBuilder() + .addField("val", Types.MinorType.DATEMILLI.getType()) + .build(); + + String actual = queryFactory.createSelectQueryBuilder(VIEW_METADATA_FIELD) + .withDatabaseName("myDatabase") + .withTableName("myTable") + .withProjection(schema) + .withConjucts(new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT)) + .build().replace("\n", ""); + + logger.info("build: actual[{}]", actual); + assertEquals(expected, actual); + + logger.info("build: exit"); + } + @Test public void buildWithView() {