HIVE-20889: Support timestamp-micros in AvroSerDe

Author: araika

common/src/java/org/apache/hadoop/hive/common/type/Timestamp.java

@@ -162,6 +162,11 @@ public long toEpochMilli() {
     return localDateTime.toInstant(ZoneOffset.UTC).toEpochMilli();
   }
 
+  public long toEpochMicro() {
+    return localDateTime.toEpochSecond(ZoneOffset.UTC) * 1_000_000
+            + localDateTime.getNano() / 1000;
+  }
+
   public long toEpochMilli(ZoneId id) {
     return localDateTime.atZone(id).toInstant().toEpochMilli();
   }
@@ -236,6 +241,18 @@ public static Timestamp ofEpochMilli(long epochMilli, int nanos) {
         .withNano(nanos));
   }
 
+  public static Timestamp ofEpochMicro(long epochMicro) {
+    int nanos = (int) ((epochMicro % 1000000) * 1000);
+    epochMicro -= nanos / 1_000_000;
+
+    Instant instant = Instant.ofEpochSecond(
+            epochMicro / 1_000_000,
+            nanos
+    );
+
+    return new Timestamp(LocalDateTime.ofInstant(instant, ZoneOffset.UTC));
+  }
+
   public void setNanos(int nanos) {
     localDateTime = localDateTime.withNano(nanos);
   }

ql/src/test/queries/clientpositive/avro_timestamp_micros.q

@@ -0,0 +1,3 @@
+CREATE EXTERNAL TABLE hive_test(`dt` timestamp) STORED AS AVRO;
+INSERT INTO hive_test VALUES (cast('2024-08-09 14:08:26.326107' as timestamp));
+SELECT * FROM hive_test;

ql/src/test/results/clientpositive/llap/avro_timestamp_micros.q.out

@@ -0,0 +1,26 @@
+PREHOOK: query: CREATE EXTERNAL TABLE hive_test(`dt` timestamp) STORED AS AVRO
+PREHOOK: type: CREATETABLE
+PREHOOK: Output: database:default
+PREHOOK: Output: default@hive_test
+POSTHOOK: query: CREATE EXTERNAL TABLE hive_test(`dt` timestamp) STORED AS AVRO
+POSTHOOK: type: CREATETABLE
+POSTHOOK: Output: database:default
+POSTHOOK: Output: default@hive_test
+PREHOOK: query: INSERT INTO hive_test VALUES (cast('2024-08-09 14:08:26.326107' as timestamp))
+PREHOOK: type: QUERY
+PREHOOK: Input: _dummy_database@_dummy_table
+PREHOOK: Output: default@hive_test
+POSTHOOK: query: INSERT INTO hive_test VALUES (cast('2024-08-09 14:08:26.326107' as timestamp))
+POSTHOOK: type: QUERY
+POSTHOOK: Input: _dummy_database@_dummy_table
+POSTHOOK: Output: default@hive_test
+POSTHOOK: Lineage: hive_test.dt SCRIPT []
+PREHOOK: query: SELECT * FROM hive_test
+PREHOOK: type: QUERY
+PREHOOK: Input: default@hive_test
+PREHOOK: Output: hdfs://### HDFS PATH ###
+POSTHOOK: query: SELECT * FROM hive_test
+POSTHOOK: type: QUERY
+POSTHOOK: Input: default@hive_test
+POSTHOOK: Output: hdfs://### HDFS PATH ###
+2024-08-09 14:08:26.326107

serde/src/java/org/apache/hadoop/hive/serde2/avro/AvroDeserializer.java

@@ -31,6 +31,7 @@
 import java.util.Map;
 import java.util.TimeZone;
 
+import org.apache.avro.LogicalType;
 import org.apache.avro.Schema;
 import org.apache.avro.Schema.Type;
 import org.apache.avro.generic.GenericData;
@@ -388,11 +389,21 @@ private Object deserializePrimitive(Object datum, Schema fileSchema, Schema reco
           skipProlepticConversion = HiveConf.ConfVars.HIVE_AVRO_PROLEPTIC_GREGORIAN_DEFAULT.defaultBoolVal;
         }
       }
-      Timestamp timestamp = TimestampTZUtil.convertTimestampToZone(
-          Timestamp.ofEpochMilli((Long) datum), ZoneOffset.UTC, convertToTimeZone, legacyConversion);
-      if (!skipProlepticConversion) {
+      LogicalType logicalType = recordSchema.getLogicalType();
+      Timestamp timestamp;
+      if (logicalType != null && logicalType.getName().equals(AvroSerDe.TIMESTAMP_TYPE_NAME_MICROS)) {
+        timestamp = Timestamp.ofEpochMicro((Long) datum);
+      } else {
+        timestamp = Timestamp.ofEpochMilli((Long) datum);
+      }
+      timestamp = TimestampTZUtil.convertTimestampToZone(
+              timestamp, ZoneOffset.UTC, convertToTimeZone, legacyConversion);
+      if (!skipProlepticConversion && logicalType.getName().equals(AvroSerDe.TIMESTAMP_TYPE_NAME_MICROS)) {
+        timestamp = Timestamp.ofEpochMicro(
+            CalendarUtils.convertTimeToProlepticMicros(timestamp.toEpochMicro()));
+      } else if (!skipProlepticConversion) {
         timestamp = Timestamp.ofEpochMilli(
-            CalendarUtils.convertTimeToProleptic(timestamp.toEpochMilli()));
+                CalendarUtils.convertTimeToProleptic(timestamp.toEpochMilli()));
       }
       return timestamp;
     }

serde/src/java/org/apache/hadoop/hive/serde2/avro/AvroSerDe.java

@@ -57,7 +57,8 @@ public class AvroSerDe extends AbstractSerDe {
   public static final String CHAR_TYPE_NAME = "char";
   public static final String VARCHAR_TYPE_NAME = "varchar";
   public static final String DATE_TYPE_NAME = "date";
-  public static final String TIMESTAMP_TYPE_NAME = "timestamp-millis";
+  public static final String TIMESTAMP_TYPE_NAME_MILLIS = "timestamp-millis";
+  public static final String TIMESTAMP_TYPE_NAME_MICROS = "timestamp-micros";
   public static final String WRITER_TIME_ZONE = "writer.time.zone";
   public static final String WRITER_PROLEPTIC = "writer.proleptic";
   public static final String WRITER_ZONE_CONVERSION_LEGACY = "writer.zone.conversion.legacy";

serde/src/java/org/apache/hadoop/hive/serde2/avro/AvroSerializer.java

@@ -17,13 +17,16 @@
  */
 package org.apache.hadoop.hive.serde2.avro;
 
+import java.time.Instant;
+import java.time.LocalDateTime;
 import java.time.ZoneOffset;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TimeZone;
 
+import org.apache.avro.LogicalType;
 import org.apache.avro.Schema;
 import org.apache.avro.Schema.Field;
 import org.apache.avro.Schema.Type;
@@ -231,10 +234,18 @@ private Object serializePrimitive(TypeInfo typeInfo, PrimitiveObjectInspector fi
     case TIMESTAMP:
       Timestamp timestamp =
         ((TimestampObjectInspector) fieldOI).getPrimitiveJavaObject(structFieldData);
+      LogicalType logicalType = schema.getLogicalType();
+      if (logicalType != null && logicalType.getName().equals(AvroSerDe.TIMESTAMP_TYPE_NAME_MICROS)) {
+        long micros = defaultProleptic ? timestamp.toEpochMicro() :
+                CalendarUtils.convertTimeToProlepticMicros(timestamp.toEpochMicro());
+        timestamp = TimestampTZUtil.convertTimestampToZone(
+                Timestamp.ofEpochMicro(micros), TimeZone.getDefault().toZoneId(), ZoneOffset.UTC, legacyConversion);
+        return timestamp.toEpochMicro();
+      }
       long millis = defaultProleptic ? timestamp.toEpochMilli() :
-          CalendarUtils.convertTimeToHybrid(timestamp.toEpochMilli());
+              CalendarUtils.convertTimeToHybrid(timestamp.toEpochMilli());
       timestamp = TimestampTZUtil.convertTimestampToZone(
-          Timestamp.ofEpochMilli(millis), TimeZone.getDefault().toZoneId(), ZoneOffset.UTC, legacyConversion);
+              Timestamp.ofEpochMilli(millis), TimeZone.getDefault().toZoneId(), ZoneOffset.UTC, legacyConversion);
       return timestamp.toEpochMilli();
     case UNKNOWN:
       throw new AvroSerdeException("Received UNKNOWN primitive category.");

serde/src/java/org/apache/hadoop/hive/serde2/avro/SchemaToTypeInfo.java

@@ -180,7 +180,8 @@ public static TypeInfo generateTypeInfo(Schema schema,
     }
 
     if (type == LONG &&
-      AvroSerDe.TIMESTAMP_TYPE_NAME.equals(schema.getProp(AvroSerDe.AVRO_PROP_LOGICAL_TYPE))) {
+            (AvroSerDe.TIMESTAMP_TYPE_NAME_MILLIS.equals(schema.getProp(AvroSerDe.AVRO_PROP_LOGICAL_TYPE)) ||
+                    AvroSerDe.TIMESTAMP_TYPE_NAME_MICROS.equals(schema.getProp(AvroSerDe.AVRO_PROP_LOGICAL_TYPE)))) {
       return TypeInfoFactory.timestampTypeInfo;
     }
 

serde/src/java/org/apache/hadoop/hive/serde2/avro/TypeInfoToSchema.java

@@ -159,7 +159,7 @@ private Schema createAvroPrimitive(TypeInfo typeInfo) {
       case TIMESTAMP:
         schema = AvroSerdeUtils.getSchemaFor("{" +
           "\"type\":\"" + AvroSerDe.AVRO_LONG_TYPE_NAME + "\"," +
-          "\"logicalType\":\"" + AvroSerDe.TIMESTAMP_TYPE_NAME + "\"}");
+          "\"logicalType\":\"" + AvroSerDe.TIMESTAMP_TYPE_NAME_MICROS + "\"}");
         break;
       case VOID:
         schema = Schema.create(Schema.Type.NULL);

serde/src/test/org/apache/hadoop/hive/serde2/avro/TestAvroObjectInspectorGenerator.java

@@ -227,7 +227,7 @@ public class TestAvroObjectInspectorGenerator {
       "  \"fields\" : [\n" +
       "    {\"name\":\"timestampField\", " +
       "     \"type\":\"" + AvroSerDe.AVRO_LONG_TYPE_NAME + "\", " +
-      "         \"logicalType\":\"" + AvroSerDe.TIMESTAMP_TYPE_NAME + "\"}" +
+      "         \"logicalType\":\"" + AvroSerDe.TIMESTAMP_TYPE_NAME_MILLIS + "\"}" +
       "  ]\n" +
       "}";
 

serde/src/test/org/apache/hadoop/hive/serde2/avro/TestAvroSerializer.java

@@ -125,7 +125,7 @@ public void canSerializeDoubles() throws SerDeException, IOException {
   public void canSerializeTimestamps() throws SerDeException, IOException {
     singleFieldTest("timestamp1", Timestamp.valueOf("2011-01-01 00:00:00").toEpochMilli(),
         "\"" + AvroSerDe.AVRO_LONG_TYPE_NAME + "\"," +
-        "\"logicalType\":\"" + AvroSerDe.TIMESTAMP_TYPE_NAME + "\"");
+        "\"logicalType\":\"" + AvroSerDe.TIMESTAMP_TYPE_NAME_MILLIS + "\"");
   }
 
   @Test

storage-api/src/java/org/apache/hadoop/hive/common/type/CalendarUtils.java

@@ -157,6 +157,27 @@ public static long convertTimeToProleptic(long hybrid) {
     return proleptic;
   }
 
+  /**
+   * Convert epoch microseconds from the hybrid Julian/Gregorian calendar to the
+   * proleptic Gregorian.
+   * @param hybridMicros Microseconds of epoch in the hybrid Julian/Gregorian
+   * @return Microseconds of epoch in the proleptic Gregorian
+   */
+  public static long convertTimeToProlepticMicros(long hybridMicros) {
+    long prolepticMicros = hybridMicros;
+    long hybridMillis = hybridMicros / 1_000L; // Convert micros to millis
+
+    if (hybridMillis < SWITCHOVER_MILLIS) {
+      String dateStr = HYBRID_TIME_FORMAT.get().format(new Date(hybridMillis));
+      try {
+        prolepticMicros = PROLEPTIC_TIME_FORMAT.get().parse(dateStr).getTime() * 1_000L; // Convert millis back to micros
+      } catch (ParseException e) {
+        throw new IllegalArgumentException("Can't parse " + dateStr, e);
+      }
+    }
+    return prolepticMicros;
+  }
+
   /**
    * Convert epoch millis from the proleptic Gregorian calendar to the hybrid
    * Julian/Gregorian.
@@ -176,6 +197,27 @@ public static long convertTimeToHybrid(long proleptic) {
     return hybrid;
   }
 
+  /**
+   * Convert epoch microseconds from the proleptic Gregorian calendar to the
+   * hybrid Julian/Gregorian.
+   * @param prolepticMicros Microseconds of epoch in the proleptic Gregorian
+   * @return Microseconds of epoch in the hybrid Julian/Gregorian
+   */
+  public static long convertTimeToHybridMicros(long prolepticMicros) {
+    long hybridMicros = prolepticMicros;
+    long prolepticMillis = prolepticMicros / 1_000L; // Convert micros to millis
+
+    if (prolepticMillis < SWITCHOVER_MILLIS) {
+      String dateStr = PROLEPTIC_TIME_FORMAT.get().format(new Date(prolepticMillis));
+      try {
+        hybridMicros = HYBRID_TIME_FORMAT.get().parse(dateStr).getTime() * 1_000L; // Convert millis back to micros
+      } catch (ParseException e) {
+        throw new IllegalArgumentException("Can't parse " + dateStr, e);
+      }
+    }
+    return hybridMicros;
+  }
+
   /**
    *
    * Formats epoch day to date according to proleptic or hybrid calendar