feat(arrow): store and read a timestamp as the wall clock it is

common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala

@@ -24,7 +24,7 @@ import org.apache.texera.amber.core.tuple.AttributeTypeUtils.AttributeTypeExcept
 import org.apache.texera.amber.core.tuple._
 import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
 import org.apache.arrow.vector.types.FloatingPointPrecision
-import org.apache.arrow.vector.types.TimeUnit.MILLISECOND
+import org.apache.arrow.vector.types.TimeUnit
 import org.apache.arrow.vector.types.pojo.ArrowType.PrimitiveType
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
 import org.apache.arrow.vector.{
@@ -40,6 +40,9 @@ import org.apache.arrow.vector.{
 }
 
 import java.nio.charset.StandardCharsets
+import java.sql.Timestamp
+import java.time.temporal.ChronoUnit
+import java.time.{Instant, LocalDateTime, ZoneId, ZoneOffset}
 import java.util
 import scala.jdk.CollectionConverters.CollectionHasAsScala
 import scala.language.implicitConversions
@@ -81,7 +84,10 @@ object ArrowUtils extends LazyLogging {
               // Use the attribute type from the schema (which includes metadata)
               // instead of deriving it from the Arrow type
               val attributeType = schema.getAttributes(index).getType
-              AttributeTypeUtils.parseField(value, attributeType)
+              fieldVector.getField.getType match {
+                case timestamp: ArrowType.Timestamp => wallClockOf(value, timestamp)
+                case _                              => AttributeTypeUtils.parseField(value, attributeType)
+              }
             } catch {
               case e: Exception =>
                 logger.warn("Caught error during parsing Arrow value back to Texera value", e)
@@ -93,6 +99,54 @@ object ArrowUtils extends LazyLogging {
       .build()
   }
 
+  /** The wall clock a timestamp column holds, read the way its own field states.
+    *
+    * A Texera TIMESTAMP carries no zone, so the wall clock is the whole of what
+    * crosses over, and the field says how to arrive at one: the zone its numbers
+    * are counted in, and the unit they are counted in. A zoned vector hands back
+    * that number bare, and letting `new Timestamp(number)` make the wall clock
+    * would take it for milliseconds in the JVM's zone: one file would then say
+    * different things on servers in different places, with nothing in the file to
+    * account for the difference. A zoneless vector hands back a LocalDateTime
+    * already, which is the wall clock itself.
+    */
+  private def wallClockOf(value: AnyRef, field: ArrowType.Timestamp): Timestamp =
+    value match {
+      case null               => null
+      case ldt: LocalDateTime => Timestamp.valueOf(ldt)
+      case number: java.lang.Long =>
+        Timestamp.valueOf(
+          LocalDateTime.ofInstant(instantOf(number, field.getUnit), zoneOf(field))
+        )
+      case other => AttributeTypeUtils.parseTimestamp(other)
+    }
+
+  /** The zone a timestamp field counts its numbers in. An unlabelled field counts
+    * from the epoch with no zone in the picture, which is what UTC arithmetic is.
+    */
+  private def zoneOf(field: ArrowType.Timestamp): ZoneId =
+    Option(field.getTimezone).map(ZoneId.of).getOrElse(ZoneOffset.UTC)
+
+  /** The instant a zoned vector's bare number stands for. Arrow leaves those
+    * unscaled, so the field's unit is what says how far from the epoch it reaches.
+    */
+  private def instantOf(number: Long, unit: TimeUnit): Instant =
+    unit match {
+      case TimeUnit.SECOND      => Instant.ofEpochSecond(number)
+      case TimeUnit.MILLISECOND => Instant.ofEpochMilli(number)
+      case TimeUnit.MICROSECOND => Instant.EPOCH.plus(number, ChronoUnit.MICROS)
+      case TimeUnit.NANOSECOND  => Instant.EPOCH.plusNanos(number)
+    }
+
+  /** The number a field of this unit records an instant as, inverting [[instantOf]]. */
+  private def numberOf(instant: Instant, unit: TimeUnit): Long =
+    unit match {
+      case TimeUnit.SECOND      => instant.getEpochSecond
+      case TimeUnit.MILLISECOND => instant.toEpochMilli
+      case TimeUnit.MICROSECOND => ChronoUnit.MICROS.between(Instant.EPOCH, instant)
+      case TimeUnit.NANOSECOND  => ChronoUnit.NANOS.between(Instant.EPOCH, instant)
+    }
+
   /**
     * Converts an Arrow Schema into Texera Schema.
     * Checks field metadata to recover types that share an Arrow representation
@@ -213,17 +267,29 @@ object ArrowUtils extends LazyLogging {
             .asInstanceOf[Float8Vector]
             .setSafe(index, !isNull, if (isNull) 0 else value.asInstanceOf[Double])
 
-        case _: ArrowType.Timestamp =>
+        // The wall clock written as the field's own zone and unit, so the number
+        // and the label beside it agree. Going through the value's own epoch
+        // would have read the wall clock in the JVM's zone instead, putting a
+        // machine's setting into the file: the same table written in two places
+        // would hold two different instants under one label. Inverts
+        // [[wallClockOf]].
+        case timestamp: ArrowType.Timestamp =>
           vector
             .asInstanceOf[TimeStampVector]
             .setSafe(
               index,
               !isNull,
               if (isNull) 0L
               else
-                AttributeTypeUtils
-                  .parseField(value, AttributeType.LONG)
-                  .asInstanceOf[Long]
+                numberOf(
+                  AttributeTypeUtils
+                    .parseField(value, AttributeType.TIMESTAMP)
+                    .asInstanceOf[Timestamp]
+                    .toLocalDateTime
+                    .atZone(zoneOf(timestamp))
+                    .toInstant,
+                  timestamp.getUnit
+                )
             )
 
         case _: ArrowType.Utf8 =>
@@ -299,7 +365,7 @@ object ArrowUtils extends LazyLogging {
         ArrowType.Bool.INSTANCE
 
       case AttributeType.TIMESTAMP =>
-        new ArrowType.Timestamp(MILLISECOND, "UTC")
+        new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")
 
       case AttributeType.BINARY =>
         new ArrowType.Binary

common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala

@@ -20,7 +20,7 @@
 package org.apache.texera.amber.util
 
 import org.apache.arrow.memory.RootAllocator
-import org.apache.arrow.vector.{VarCharVector, VectorSchemaRoot}
+import org.apache.arrow.vector.{TimeStampVector, VarCharVector, VectorSchemaRoot}
 import org.apache.arrow.vector.types.{FloatingPointPrecision, TimeUnit}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
 import org.apache.texera.amber.core.state.State
@@ -436,6 +436,112 @@ class ArrowUtilsSpec extends AnyFlatSpec with Matchers {
     }
   }
 
+  // ----- Timestamp fields that are not the UTC millisecond ones we write -----
+
+  // fromTexeraSchema only ever writes Timestamp(MILLISECOND, "UTC"), so the
+  // roots built above never exercise another zone or unit. An .arrow file handed
+  // to ArrowSourceOpDesc can carry either: pandas writes a tz-aware column as
+  // Timestamp(NANOSECOND, <its zone>). These build the field directly.
+  private def withFieldRoot(field: Field)(body: VectorSchemaRoot => Unit): Unit = {
+    val allocator = new RootAllocator()
+    val root = VectorSchemaRoot.create(
+      new org.apache.arrow.vector.types.pojo.Schema(util.Arrays.asList(field)),
+      allocator
+    )
+    try {
+      root.allocateNew()
+      body(root)
+    } finally {
+      root.close()
+      allocator.close()
+    }
+  }
+
+  private val timestampSchema = Schema(List(new Attribute("t", AttributeType.TIMESTAMP)))
+
+  private def timestampTuple(wallClock: String): Tuple =
+    Tuple
+      .builder(timestampSchema)
+      .addSequentially(Array[Any](Timestamp.valueOf(wallClock)))
+      .build()
+
+  "getTexeraTuple" should "read a zoned timestamp as the wall clock its field's zone gives" in {
+    // 1704603600000 is 2024-01-07 00:00 in New York and 05:00 in UTC. The field
+    // says New York, so New York is the reading that comes back.
+    val field =
+      arrowField("t", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "America/New_York"))
+    withFieldRoot(field) { root =>
+      root.getVector(0).asInstanceOf[TimeStampVector].setSafe(0, 1704603600000L)
+      root.setRowCount(1)
+      ArrowUtils.getTexeraTuple(0, root).getField[Timestamp]("t") shouldBe
+        Timestamp.valueOf("2024-01-07 00:00:00")
+    }
+  }
+
+  it should "count a zoned timestamp in the unit its field declares" in {
+    // Arrow hands a zoned vector its number unscaled, so the same reading is a
+    // million times the number in a nanosecond field that it is in a millisecond
+    // one. Taken for milliseconds, this one would land in 1970.
+    val field =
+      arrowField("t", new ArrowType.Timestamp(TimeUnit.NANOSECOND, "America/New_York"))
+    withFieldRoot(field) { root =>
+      root.getVector(0).asInstanceOf[TimeStampVector].setSafe(0, 1704603600000000000L)
+      root.setRowCount(1)
+      ArrowUtils.getTexeraTuple(0, root).getField[Timestamp]("t") shouldBe
+        Timestamp.valueOf("2024-01-07 00:00:00")
+    }
+  }
+
+  it should "read an unlabelled timestamp as the wall clock it already is" in {
+    // No zone to reconcile, and Arrow scales the number itself: the vector hands
+    // back a LocalDateTime, which is the wall clock, whatever the server's zone.
+    val field = arrowField("t", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))
+    withFieldRoot(field) { root =>
+      root.getVector(0).asInstanceOf[TimeStampVector].setSafe(0, 1704603600000L)
+      root.setRowCount(1)
+      ArrowUtils.getTexeraTuple(0, root).getField[Timestamp]("t") shouldBe
+        Timestamp.valueOf("2024-01-07 05:00:00")
+    }
+  }
+
+  "setTexeraTuple" should "store a wall clock as the number its field's zone calls for" in {
+    val field =
+      arrowField("t", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "America/New_York"))
+    withFieldRoot(field) { root =>
+      ArrowUtils.setTexeraTuple(timestampTuple("2024-01-07 00:00:00"), 0, root)
+      // The New York instant of that reading. Stored as the UTC one it would be
+      // 1704585600000, five hours off what the field's own label promises.
+      root.getVector(0).asInstanceOf[TimeStampVector].get(0) shouldBe 1704603600000L
+    }
+  }
+
+  it should "store a wall clock in the unit its field declares" in {
+    val field = arrowField("t", new ArrowType.Timestamp(TimeUnit.NANOSECOND, "America/New_York"))
+    withFieldRoot(field) { root =>
+      ArrowUtils.setTexeraTuple(timestampTuple("2024-01-07 00:00:00"), 0, root)
+      root.getVector(0).asInstanceOf[TimeStampVector].get(0) shouldBe 1704603600000000000L
+    }
+  }
+
+  "a timestamp round-trip" should "preserve the wall clock through a field of any zone and unit" in {
+    val fields = List(
+      new ArrowType.Timestamp(TimeUnit.SECOND, "Asia/Tokyo"),
+      new ArrowType.Timestamp(TimeUnit.MILLISECOND, "America/New_York"),
+      new ArrowType.Timestamp(TimeUnit.MICROSECOND, "Europe/Berlin"),
+      new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"),
+      new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)
+    )
+    fields.foreach { arrowType =>
+      withFieldRoot(arrowField("t", arrowType)) { root =>
+        ArrowUtils.setTexeraTuple(timestampTuple("2024-01-07 09:30:00"), 0, root)
+        withClue(s"$arrowType: ") {
+          ArrowUtils.getTexeraTuple(0, root).getField[Timestamp]("t") shouldBe
+            Timestamp.valueOf("2024-01-07 09:30:00")
+        }
+      }
+    }
+  }
+
   // ----- fromAttributeType (null input) -----
 
   "fromAttributeType" should "reject a null attribute type" in {

common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala

@@ -149,7 +149,11 @@ class ArrowUtilsSpec extends AnyFlatSpec {
           Long.box(1L),
           Boolean.box(true),
           Double.box(1.1),
-          new Timestamp(10000L),
+          // Stated as a wall clock, which is what a Texera TIMESTAMP holds. Built
+          // from an epoch instead, the wall clock would be whichever one the
+          // machine's zone gives that instant, and what gets stored below would
+          // move with it.
+          Timestamp.valueOf("1970-01-01 00:00:10"),
           "hello world"
         )
       )
@@ -169,7 +173,9 @@ class ArrowUtilsSpec extends AnyFlatSpec {
     assert(vectorSchemaRoot.getVector(2).getObject(index).asInstanceOf[Boolean] == true)
     assert(vectorSchemaRoot.getVector(3).getObject(index).asInstanceOf[Double] == 1.1)
 
-    // the arrow storage type of timestamp is Long
+    // The arrow storage type of timestamp is Long, and the field is labelled
+    // UTC, so the wall clock above is stored as the UTC instant of the same
+    // reading: ten seconds past the epoch, on a server anywhere.
     assert(vectorSchemaRoot.getVector(4).getObject(index).asInstanceOf[Long] == 10000L)
 
     // the arrow storage type of string is Text