support 4 more logical types (time-micro, timestamp-micro, local-time…

…stamp-micros, local-timestamp-millis)
julianpeeters · Oct 3, 2023 · be8902c · be8902c
1 parent 68aa44d
commit be8902c
Show file tree

Hide file tree

Showing 13 changed files with 235 additions and 18 deletions.
diff --git a/avrohugger-core/src/main/scala/format/specific/converters/JavaConverter.scala b/avrohugger-core/src/main/scala/format/specific/converters/JavaConverter.scala
@@ -155,6 +155,18 @@ object JavaConverter {
         case JavaSqlTimestamp => BLOCK(tree.DOT("getTime").APPLY())
         case JavaTimeInstant  => BLOCK(tree.DOT("toEpochMilli"))
       }
+      case _: LogicalTypes.TimestampMicros => (typeMatcher.avroScalaTypes.timestampMicros match {
+        case JavaTimeZonedDateTime => BLOCK(tree.DOT("toEpochSecond").INFIX("*", LIT(1000000L)).INFIX("+", tree.DOT("getNano").INFIX("/", LIT(1000L))))
+      }) withComment "avro timestamp-micros long stores the number of microseconds from the unix epoch, 1 January 1970 00:00:00.000000 UTC"
+      case _: LogicalTypes.LocalTimestampMillis => (typeMatcher.avroScalaTypes.localTimestampMillis match {
+        case JavaTimeLocalDateTime => BLOCK(tree.DOT("toEpochSecond").APPLY(RootClass.newClass("java.time.ZoneOffset").DOT("UTC")).INFIX("*", LIT(1000L)).INFIX("+", tree.DOT("getNano").INFIX("/", LIT(1000000L))))
+      }) withComment "avro local-timestamp-millis long stores the number of millis, from 1 January 1970 00:00:00.000000"
+      case _: LogicalTypes.LocalTimestampMicros => (typeMatcher.avroScalaTypes.localTimestampMicros match {
+        case JavaTimeLocalDateTime => BLOCK(tree.DOT("toEpochSecond").APPLY(RootClass.newClass("java.time.ZoneOffset").DOT("UTC")).INFIX("*", LIT(1000000L)).INFIX("+", tree.DOT("getNano").INFIX("/", LIT(1000L))))
+      }) withComment "avro local-timestamp-micros long stores the number of microseconds, from 1 January 1970 00:00:00.000000"
+      case _: LogicalTypes.TimeMicros => (typeMatcher.avroScalaTypes.timeMicros match {
+        case JavaTimeLocalTime => BLOCK(tree.DOT("toNanoOfDay").INFIX("/", LIT(1000L)))
+      }) withComment "avro time-micros long stores the number of microseconds after midnight, 00:00:00.000000"
       case _ => tree
     }
     case Schema.Type.INT => schema.getLogicalType match {

diff --git a/avrohugger-core/src/main/scala/format/specific/converters/ScalaConverter.scala b/avrohugger-core/src/main/scala/format/specific/converters/ScalaConverter.scala
@@ -197,9 +197,17 @@ object ScalaConverter {
         }
       }
       case Schema.Type.LONG => {
+        val caseLWithTypeLong = CASE(ID("l") withType (LongClass))
         Option(schema.getLogicalType()) match {
           case Some(logicalType) => {
-            if (logicalType.getName == "timestamp-millis") {
+            if (logicalType.getName == "time-micros") {
+              (typeMatcher.avroScalaTypes.timeMicros match {
+                case JavaTimeLocalTime =>
+                  val LocalTimeClass = RootClass.newClass("java.time.LocalTime")
+                  val resultExpr = BLOCK(LocalTimeClass.DOT("ofNanoOfDay").APPLY(REF("l").INFIX("*", LIT(1000L))))
+                  tree MATCH caseLWithTypeLong ==> resultExpr
+              }) withComment "avro time-micros long stores the number of microseconds after midnight, 00:00:00.000000"
+            } else if (logicalType.getName == "timestamp-millis") {
               typeMatcher.avroScalaTypes.timestampMillis match {
                 case JavaSqlTimestamp => {
                   val TimestampClass = RootClass.newClass("java.sql.Timestamp")
@@ -214,6 +222,44 @@ object ScalaConverter {
                   tree MATCH longConversion
                 }
               }
+            } else if (logicalType.getName == "timestamp-micros") {
+              (typeMatcher.avroScalaTypes.timestampMicros match {
+                case JavaTimeZonedDateTime =>
+                  val ZonedDateTime = RootClass.newClass("java.time.ZonedDateTime")
+                  val LocalDateTime = RootClass.newClass("java.time.LocalDateTime")
+                  val ZoneOffset = RootClass.newClass("java.time.ZoneOffset")
+                  val ZoneId = RootClass.newClass("java.time.ZoneId")
+                  val resultExpr = BLOCK(ZonedDateTime.DOT("of").APPLY(LocalDateTime DOT "ofEpochSecond" APPLY(
+                    REF("l").INFIX("/", LIT(1000000L)),
+                    PAREN(REF("l").INFIX("%", LIT(1000000L))) DOT "toInt" INFIX("*", LIT(1000)),
+                    ZoneOffset DOT "UTC"
+                  ), ZoneId DOT "of" APPLY LIT("UTC")))
+                  tree MATCH CASE(ID("l") withType (LongClass)) ==> resultExpr
+              }) withComment "avro timestamp-micros long stores the number of microseconds from the unix epoch, 1 January 1970 00:00:00.000000 UTC"
+            } else if (logicalType.getName == "local-timestamp-millis") {
+              (typeMatcher.avroScalaTypes.localTimestampMillis match {
+                case JavaTimeLocalDateTime =>
+                  val LocalDateTime = RootClass.newClass("java.time.LocalDateTime")
+                  val ZoneOffset = RootClass.newClass("java.time.ZoneOffset")
+                  val resultExpr = BLOCK(LocalDateTime DOT "ofEpochSecond" APPLY(
+                    REF("l").INFIX("/", LIT(1000L)),
+                    PAREN(REF("l").INFIX("%", LIT(1000L))) DOT "toInt" INFIX("*", LIT(1000000)),
+                    ZoneOffset DOT "UTC"
+                  ))
+                  tree MATCH CASE(ID("l") withType (LongClass)) ==> resultExpr
+              }) withComment "avro local-timestamp-millis long stores the number of millis, from 1 January 1970 00:00:00.000000"
+            } else if (logicalType.getName == "local-timestamp-micros") {
+              (typeMatcher.avroScalaTypes.localTimestampMicros match {
+                case JavaTimeLocalDateTime =>
+                  val LocalDateTime = RootClass.newClass("java.time.LocalDateTime")
+                  val ZoneOffset = RootClass.newClass("java.time.ZoneOffset")
+                  val resultExpr = BLOCK(LocalDateTime DOT "ofEpochSecond" APPLY(
+                    REF("l").INFIX("/", LIT(1000000L)),
+                    PAREN(REF("l").INFIX("%", LIT(1000000L))) DOT "toInt" INFIX("*", LIT(1000)),
+                    ZoneOffset DOT "UTC"
+                  ))
+                  tree MATCH CASE(ID("l") withType (LongClass)) ==> resultExpr
+              }) withComment "avro local-timestamp-micros long stores the number of microseconds, from 1 January 1970 00:00:00.000000"
             }
             else tree
           }

diff --git a/avrohugger-core/src/main/scala/matchers/DefaultParamMatcher.scala b/avrohugger-core/src/main/scala/matchers/DefaultParamMatcher.scala
@@ -42,6 +42,18 @@ object DefaultParamMatcher {
           case TimestampMillis =>
             CustomDefaultParamMatcher.checkCustomTimestampMillisType(
               typeMatcher.avroScalaTypes.timestampMillis)
+          case TimestampMicros =>
+            CustomDefaultParamMatcher.checkCustomTimestampMicrosType(
+              typeMatcher.avroScalaTypes.timestampMicros)
+          case LocalTimestampMillis =>
+            CustomDefaultParamMatcher.checkCustomLocalTimestampMillisType(
+              typeMatcher.avroScalaTypes.localTimestampMillis)
+          case LocalTimestampMicros =>
+            CustomDefaultParamMatcher.checkCustomLocalTimestampMicrosType(
+              typeMatcher.avroScalaTypes.localTimestampMicros)
+          case TimeMicros =>
+            CustomDefaultParamMatcher.checkCustomTimeMicrosType(
+              typeMatcher.avroScalaTypes.timeMicros)
         }
       case Schema.Type.FLOAT   => LIT(0F)
       case Schema.Type.DOUBLE  => LIT(0D)

diff --git a/avrohugger-core/src/main/scala/matchers/TypeMatcher.scala b/avrohugger-core/src/main/scala/matchers/TypeMatcher.scala
@@ -49,6 +49,10 @@ class TypeMatcher(
             schema = schema,
             default = CustomTypeMatcher.checkCustomNumberType(avroScalaTypes.long)) {
             case TimestampMillis => CustomTypeMatcher.checkCustomTimestampMillisType(avroScalaTypes.timestampMillis)
+            case TimestampMicros => CustomTypeMatcher.checkCustomTimestampMicrosType(avroScalaTypes.timestampMicros)
+            case LocalTimestampMicros => CustomTypeMatcher.checkCustomLocalTimestampMicrosType(avroScalaTypes.localTimestampMicros)
+            case LocalTimestampMillis => CustomTypeMatcher.checkCustomLocalTimestampMillisType(avroScalaTypes.localTimestampMillis)
+            case TimeMicros => CustomTypeMatcher.checkCustomTimeMicrosType(avroScalaTypes.timeMicros)
           }
         case Schema.Type.INT      =>
           LogicalType.foldLogicalTypes(

diff --git a/avrohugger-core/src/main/scala/matchers/custom/CustomDefaultParamMatcher.scala b/avrohugger-core/src/main/scala/matchers/custom/CustomDefaultParamMatcher.scala
@@ -58,4 +58,24 @@ object CustomDefaultParamMatcher {
       case JavaSqlTime => NEW(REF("java.sql.Time"), LIT(0L))
       case JavaTimeLocalTime  => REF("java.time.LocalTime.now")
     }
+
+  def checkCustomTimeMicrosType(timeMillisType: AvroScalaTimeType): Tree =
+    timeMillisType match {
+      case JavaTimeLocalTime => REF("java.time.LocalTime.MIDNIGHT")
+    }
+
+  def checkCustomTimestampMicrosType(timeMillisType: AvroScalaTimestampType): Tree =
+    timeMillisType match {
+      case JavaTimeZonedDateTime => REF("java.time.ZonedDateTime.of").APPLY(REF("java.time.LocalDateTime") DOT "MIN", REF("java.time.ZoneId") DOT "of" APPLY LIT("UTC"))
+    }
+
+  def checkCustomLocalTimestampMillisType(timeMillisType: AvroScalaLocalTimestampType): Tree =
+    timeMillisType match {
+      case JavaTimeLocalDateTime => REF("java.time.LocalDateTime") DOT "MIN"
+    }
+
+  def checkCustomLocalTimestampMicrosType(timeMillisType: AvroScalaLocalTimestampType): Tree =
+    timeMillisType match {
+      case JavaTimeLocalDateTime => REF("java.time.LocalDateTime") DOT "MIN"
+    }
 }
diff --git a/avrohugger-core/src/main/scala/matchers/custom/CustomTypeMatcher.scala b/avrohugger-core/src/main/scala/matchers/custom/CustomTypeMatcher.scala
@@ -62,4 +62,20 @@ object CustomTypeMatcher {
             case ScalaBigDecimalWithPrecision(_) => decimalTaggedType(precision, scale)
           }
         }
+
+  def checkCustomTimeMicrosType(timeType: AvroScalaTimeType) = timeType match {
+    case JavaTimeLocalTime => RootClass.newClass(nme.createNameType("java.time.LocalTime"))
+  }
+
+  def checkCustomTimestampMicrosType(timeType: AvroScalaTimestampType) = timeType match {
+    case JavaTimeZonedDateTime => RootClass.newClass(nme.createNameType("java.time.ZonedDateTime"))
+  }
+
+  def checkCustomLocalTimestampMicrosType(timeType: AvroScalaLocalTimestampType) = timeType match {
+    case JavaTimeLocalDateTime => RootClass.newClass(nme.createNameType("java.time.LocalDateTime"))
+  }
+
+  def checkCustomLocalTimestampMillisType(timeType: AvroScalaLocalTimestampType) = timeType match {
+    case JavaTimeLocalDateTime => RootClass.newClass(nme.createNameType("java.time.LocalDateTime"))
+  }
 }
diff --git a/avrohugger-core/src/main/scala/types/AvroScalaTypes.scala b/avrohugger-core/src/main/scala/types/AvroScalaTypes.scala
@@ -24,6 +24,10 @@ case class AvroScalaTypes(
   date:     AvroScalaDateType     = JavaTimeLocalDate,
   timestampMillis: AvroScalaTimestampMillisType = JavaTimeInstant,
   timeMillis: AvroScalaTimeMillisType = JavaTimeLocalTime,
+  timeMicros: AvroScalaTimeType = JavaTimeLocalTime,
+  timestampMicros: AvroScalaTimestampType = JavaTimeZonedDateTime,
+  localTimestampMillis: AvroScalaLocalTimestampType = JavaTimeLocalDateTime,
+  localTimestampMicros: AvroScalaLocalTimestampType = JavaTimeLocalDateTime,
   uuid:     AvroUuidType          = JavaUuid
 )
 

diff --git a/avrohugger-core/src/main/scala/types/LogicalAvroScalaTypes.scala b/avrohugger-core/src/main/scala/types/LogicalAvroScalaTypes.scala
@@ -18,14 +18,25 @@ case object JavaTimeInstant extends AvroScalaTimestampMillisType
 sealed trait AvroUuidType extends Product with Serializable
 case object JavaUuid extends AvroUuidType
 
+sealed trait AvroScalaTimeType extends Serializable
 sealed trait AvroScalaTimeMillisType extends Product with Serializable
 case object JavaSqlTime extends AvroScalaTimeMillisType
-case object JavaTimeLocalTime extends AvroScalaTimeMillisType
+case object JavaTimeLocalTime extends AvroScalaTimeType with AvroScalaTimeMillisType
+
+sealed trait AvroScalaTimestampType extends Serializable
+case object JavaTimeZonedDateTime extends AvroScalaTimestampType
+
+sealed trait AvroScalaLocalTimestampType extends Serializable
+case object JavaTimeLocalDateTime extends AvroScalaLocalTimestampType
 
 sealed abstract class LogicalType(name: String)
 case class Decimal(precision: Int, scale: Int) extends LogicalType("decimal")
 case object Date extends LogicalType("date")
 case object TimestampMillis extends LogicalType("timestamp-millis")
+case object TimestampMicros extends LogicalType("timestamp-micros")
+case object LocalTimestampMicros extends LogicalType("local-timestamp-micros")
+case object LocalTimestampMillis extends LogicalType("local-timestamp-millis")
+case object TimeMicros extends LogicalType("timestamp-millis")
 case object TimeMillis extends LogicalType("time-millis")
 case object UUID extends LogicalType("uuid")
 
@@ -35,6 +46,10 @@ object LogicalType {
     case d: org.apache.avro.LogicalTypes.Decimal => Some(Decimal(d.getPrecision, d.getScale))
     case _: org.apache.avro.LogicalTypes.Date => Some(Date)
     case _: org.apache.avro.LogicalTypes.TimestampMillis => Some(TimestampMillis)
+    case _: org.apache.avro.LogicalTypes.TimestampMicros => Some(TimestampMicros)
+    case _: org.apache.avro.LogicalTypes.LocalTimestampMillis => Some(LocalTimestampMillis)
+    case _: org.apache.avro.LogicalTypes.LocalTimestampMicros => Some(LocalTimestampMicros)
+    case _: org.apache.avro.LogicalTypes.TimeMicros => Some(TimeMicros)
     case _: org.apache.avro.LogicalTypes.TimeMillis => Some(TimeMillis)
     case _ if logicalType.getName == "uuid" => Some(UUID)
     case _ => None

diff --git a/avrohugger-core/src/test/avro/logical.avsc b/avrohugger-core/src/test/avro/logical.avsc
@@ -59,6 +59,34 @@
         "type": "int",
         "logicalType": "time-millis"
       }
+    },
+    {
+      "name": "timeMicros",
+      "type": {
+        "type": "long",
+        "logicalType": "time-micros"
+      }
+    },
+    {
+      "name": "timestampMicros",
+      "type": {
+        "type": "long",
+        "logicalType": "timestamp-micros"
+      }
+    },
+    {
+      "name": "localTimestampMicros",
+      "type": {
+        "type": "long",
+        "logicalType": "local-timestamp-micros"
+      }
+    },
+    {
+      "name": "localTimestampMillis",
+      "type": {
+        "type": "long",
+        "logicalType": "local-timestamp-millis"
+      }
     }
   ]
 }