diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala index af14ae9acd0..3617e689918 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala +++ b/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,7 +267,13 @@ 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( @@ -221,9 +281,15 @@ object ArrowUtils extends LazyLogging { !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 diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala index c91f09540ff..e3f8014462d 100644 --- a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala +++ b/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, ). 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 { diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala index 6a97df8ca3e..94abd91b8ae 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala +++ b/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