Add timestamp support to Pinot connector

docs/src/main/sphinx/connector/pinot.rst

@@ -160,22 +160,23 @@ Data types
 
 Pinot does not allow null values in any data type and supports the following primitive types:
 
-==========================   ============
+==========================   ==============
 Pinot                        Trino
-==========================   ============
+==========================   ==============
 ``INT``                      ``INTEGER``
 ``LONG``                     ``BIGINT``
 ``FLOAT``                    ``REAL``
 ``DOUBLE``                   ``DOUBLE``
 ``STRING``                   ``VARCHAR``
 ``BYTES``                    ``VARBINARY``
 ``JSON``                     ``JSON``
+``TIMESTAMP``                ``TIMESTAMP``
 ``INT_ARRAY``                ``VARCHAR``
 ``LONG_ARRAY``               ``VARCHAR``
 ``FLOAT_ARRAY``              ``VARCHAR``
 ``DOUBLE_ARRAY``             ``VARCHAR``
 ``STRING_ARRAY``             ``VARCHAR``
-==========================   ============
+==========================   ==============
 
 .. _pinot-sql-support:
 

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotSegmentPageSource.java

@@ -17,13 +17,15 @@
 import io.airlift.slice.Slices;
 import io.trino.plugin.pinot.client.PinotDataFetcher;
 import io.trino.plugin.pinot.client.PinotDataTableWithSize;
+import io.trino.plugin.pinot.conversion.PinotTimestamps;
 import io.trino.spi.Page;
 import io.trino.spi.PageBuilder;
 import io.trino.spi.TrinoException;
 import io.trino.spi.block.Block;
 import io.trino.spi.block.BlockBuilder;
 import io.trino.spi.connector.ConnectorPageSource;
 import io.trino.spi.type.StandardTypes;
+import io.trino.spi.type.TimestampType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.VarbinaryType;
 import io.trino.spi.type.VarcharType;
@@ -167,7 +169,13 @@ private void writeBlock(BlockBuilder blockBuilder, Type columnType, int columnId
             writeBooleanBlock(blockBuilder, columnType, columnIdx);
         }
         else if (javaType.equals(long.class)) {
-            writeLongBlock(blockBuilder, columnType, columnIdx);
+            if (columnType instanceof TimestampType) {
+                // Pinot TimestampType is always ShortTimestampType.
+                writeShortTimestampBlock(blockBuilder, columnType, columnIdx);
+            }
+            else {
+                writeLongBlock(blockBuilder, columnType, columnIdx);
+            }
         }
         else if (javaType.equals(double.class)) {
             writeDoubleBlock(blockBuilder, columnType, columnIdx);
@@ -229,6 +237,15 @@ private void writeArrayBlock(BlockBuilder blockBuilder, Type columnType, int col
         }
     }
 
+    private void writeShortTimestampBlock(BlockBuilder blockBuilder, Type columnType, int columnIndex)
+    {
+        for (int i = 0; i < currentDataTable.getDataTable().getNumberOfRows(); i++) {
+            // Trino is using micros since epoch for ShortTimestampType, Pinot uses millis since epoch.
+            columnType.writeLong(blockBuilder, PinotTimestamps.toMicros(getLong(i, columnIndex)));
+            completedBytes += Long.BYTES;
+        }
+    }
+
     private Type getType(int columnIndex)
     {
         checkArgument(columnIndex < columnHandles.size(), "Invalid field index");
@@ -254,6 +271,7 @@ private long getLong(int rowIndex, int columnIndex)
             case FLOAT:
                 return floatToIntBits(currentDataTable.getDataTable().getFloat(rowIndex, columnIndex));
             case LONG:
+            case TIMESTAMP:
                 return currentDataTable.getDataTable().getLong(rowIndex, columnIndex);
             default:
                 throw new PinotException(PINOT_DECODE_ERROR, Optional.empty(), format("Unexpected pinot type: '%s'", dataType));

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotTypeConverter.java

@@ -20,6 +20,7 @@
 import io.trino.spi.type.DoubleType;
 import io.trino.spi.type.IntegerType;
 import io.trino.spi.type.RealType;
+import io.trino.spi.type.TimestampType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
 import io.trino.spi.type.TypeSignature;
@@ -95,6 +96,8 @@ private Type toTrinoType(FieldSpec.DataType dataType)
                 return jsonTypeSupplier.get();
             case BYTES:
                 return VarbinaryType.VARBINARY;
+            case TIMESTAMP:
+                return TimestampType.TIMESTAMP_MILLIS;
             default:
                 break;
         }

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/conversion/PinotTimestamps.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.pinot.conversion;
+
+import com.google.common.primitives.Longs;
+
+import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeParseException;
+
+import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND;
+import static java.time.ZoneOffset.UTC;
+import static java.time.format.DateTimeFormatter.ISO_INSTANT;
+
+public final class PinotTimestamps
+{
+    private static final DateTimeFormatter PINOT_TIMESTAMP_FORMATTER =
+            DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss[.SSS][.SS][.S]");
+
+    private PinotTimestamps() {}
+
+    public static long toMicros(long millis)
+    {
+        return millis * MICROSECONDS_PER_MILLISECOND;
+    }
+
+    public static long toMicros(Instant instant)
+    {
+        return toMicros(instant.toEpochMilli());
+    }
+
+    public static LocalDateTime tryParse(String value)
+    {
+        Long epochMillis = Longs.tryParse(value);
+        if (epochMillis != null) {
+            return LocalDateTime.ofInstant(Instant.ofEpochMilli(epochMillis), UTC);
+        }
+        // Try parsing using standard formats
+        LocalDateTime timestamp = tryParse(PINOT_TIMESTAMP_FORMATTER, value);
+        if (timestamp == null) {
+            timestamp = tryParse(ISO_INSTANT, value);
+        }
+        return timestamp;
+    }
+
+    private static LocalDateTime tryParse(DateTimeFormatter formatter, String value)
+    {
+        try {
+            return formatter.parse(value, LocalDateTime::from);
+        }
+        catch (DateTimeParseException e) {
+            // Ignore
+        }
+        return null;
+    }
+}

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/DecoderFactory.java

@@ -21,6 +21,7 @@
 import io.trino.spi.type.FixedWidthType;
 import io.trino.spi.type.IntegerType;
 import io.trino.spi.type.RealType;
+import io.trino.spi.type.TimestampType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.VarbinaryType;
 
@@ -55,6 +56,9 @@ public static Decoder createDecoder(Type type)
             if (type instanceof BooleanType) {
                 return new BooleanDecoder();
             }
+            if (type instanceof TimestampType) {
+                return new TimestampDecoder();
+            }
             throw new PinotException(PINOT_UNSUPPORTED_COLUMN_TYPE, Optional.empty(), "type '" + type + "' not supported");
         }
         if (type instanceof ArrayType) {

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/TimestampDecoder.java

@@ -0,0 +1,63 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.pinot.decoders;
+
+import io.trino.plugin.pinot.conversion.PinotTimestamps;
+import io.trino.spi.TrinoException;
+import io.trino.spi.block.BlockBuilder;
+
+import java.time.Instant;
+import java.time.LocalDateTime;
+import java.util.function.Supplier;
+
+import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
+import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS;
+import static java.lang.String.format;
+import static java.time.ZoneOffset.UTC;
+
+public class TimestampDecoder
+        implements Decoder
+{
+    @Override
+    public void decode(Supplier<Object> getter, BlockBuilder output)
+    {
+        Object value = getter.get();
+        if (value == null) {
+            output.appendNull();
+        }
+        else {
+            LocalDateTime timestamp;
+            if (value instanceof String stringValue) {
+                timestamp = PinotTimestamps.tryParse(stringValue);
+                if (timestamp == null) {
+                    throw new TrinoException(NOT_SUPPORTED, format(
+                            "Unable to parse string representation of type TIMESTAMP: %s [%s]",
+                            value,
+                            value.getClass().getSimpleName()));
+                }
+            }
+            else if (value instanceof Double || value instanceof Long) {
+                timestamp = LocalDateTime.ofInstant(Instant.ofEpochMilli(((Number) value).longValue()), UTC);
+            }
+            else {
+                throw new TrinoException(NOT_SUPPORTED, format(
+                        "Unsupported representation of type TIMESTAMP: %s [%s]",
+                        value,
+                        value.getClass().getSimpleName()));
+            }
+            long epochMicros = PinotTimestamps.toMicros(timestamp.atOffset(UTC).toInstant());
+            TIMESTAMP_MILLIS.writeLong(output, epochMicros);
+        }
+    }
+}

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/query/PinotQueryBuilder.java

@@ -23,6 +23,8 @@
 import io.trino.spi.predicate.Range;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.type.RealType;
+import io.trino.spi.type.TimestampType;
+import io.trino.spi.type.Timestamps;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.VarbinaryType;
 import io.trino.spi.type.VarcharType;
@@ -167,9 +169,20 @@ private static Object convertValue(Type type, Object value)
         if (type instanceof VarbinaryType) {
             return Hex.encodeHexString(((Slice) value).getBytes());
         }
+        if (type instanceof TimestampType) {
+            return toMillis((Long) value);
+        }
         return value;
     }
 
+    private static Long toMillis(Long value)
+    {
+        if (value == null) {
+            return null;
+        }
+        return Timestamps.epochMicrosToMillisWithRounding(value);
+    }
+
     private static String toConjunct(String columnName, String operator, Object value)
     {
         if (value instanceof Slice) {

plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/query/aggregation/ImplementMinMax.java

@@ -39,6 +39,7 @@
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.IntegerType.INTEGER;
 import static io.trino.spi.type.RealType.REAL;
+import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -48,7 +49,7 @@ public class ImplementMinMax
         implements AggregateFunctionRule<AggregateExpression, Void>
 {
     private static final Capture<Variable> ARGUMENT = newCapture();
-    private static final Set<Type> SUPPORTED_ARGUMENT_TYPES = ImmutableSet.of(INTEGER, BIGINT, REAL, DOUBLE);
+    private static final Set<Type> SUPPORTED_ARGUMENT_TYPES = ImmutableSet.of(INTEGER, BIGINT, REAL, DOUBLE, TIMESTAMP_MILLIS);
 
     private final Function<String, String> identifierQuote;