fix: ArrowWrapperTools Must Account For Timestamp Column's TimeUnit

deephaven · Mar 7, 2025 · 66e1d9b · 66e1d9b
1 parent 7093839
commit 66e1d9b
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 22 deletions.
diff --git a/extensions/arrow/src/main/java/io/deephaven/extensions/arrow/ArrowWrapperTools.java b/extensions/arrow/src/main/java/io/deephaven/extensions/arrow/ArrowWrapperTools.java
@@ -257,7 +257,7 @@ private static AbstractColumnSource<?> generateColumnSource(
             case TIMESTAMPNANOTZ:
             case TIMESTAMPSEC:
             case TIMESTAMPSECTZ:
-                return new ArrowInstantColumnSource(highBit, field, arrowHelper);
+                return new ArrowInstantColumnSource(vector.getMinorType(), highBit, field, arrowHelper);
             case NULL:
             case STRUCT:
             case DATEDAY:

diff --git a/...s/arrow/src/main/java/io/deephaven/extensions/arrow/sources/ArrowInstantColumnSource.java b/...s/arrow/src/main/java/io/deephaven/extensions/arrow/sources/ArrowInstantColumnSource.java
@@ -1,10 +1,6 @@
 //
 // Copyright (c) 2016-2025 Deephaven Data Labs and Patent Pending
 //
-// ****** AUTO-GENERATED CLASS - DO NOT EDIT MANUALLY
-// ****** Run GenerateArrowColumnSources or "./gradlew generateArrowColumnSources" to regenerate
-//
-// @formatter:off
 package io.deephaven.extensions.arrow.sources;
 
 import io.deephaven.chunk.WritableChunk;
@@ -15,28 +11,57 @@
 import io.deephaven.engine.table.impl.ImmutableColumnSourceGetDefaults;
 import io.deephaven.extensions.arrow.ArrowWrapperTools;
 import java.time.Instant;
+
+import io.deephaven.time.DateTimeUtils;
 import org.apache.arrow.vector.TimeStampVector;
+import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.jetbrains.annotations.NotNull;
 
 /**
- * Arrow Vector: {@link TimeStampVector}
- * Deephaven Type: java.time.Instant
+ * Arrow Vector: {@link TimeStampVector} Deephaven Type: java.time.Instant
  */
-public class ArrowInstantColumnSource extends AbstractArrowColumnSource<Instant> implements ImmutableColumnSourceGetDefaults.ForObject<Instant> {
-    public ArrowInstantColumnSource(final int highBit, final @NotNull Field field,
-            final ArrowWrapperTools. @NotNull ArrowTableContext arrowTableContext) {
+public class ArrowInstantColumnSource extends AbstractArrowColumnSource<Instant>
+        implements ImmutableColumnSourceGetDefaults.ForObject<Instant> {
+    private final long factor;
+
+    public ArrowInstantColumnSource(
+            final Types.MinorType minorType,
+            final int highBit,
+            final @NotNull Field field,
+            final @NotNull ArrowWrapperTools.ArrowTableContext arrowTableContext) {
         super(Instant.class, highBit, field, arrowTableContext);
+        switch (minorType) {
+            case TIMESTAMPNANO:
+            case TIMESTAMPNANOTZ:
+                factor = 1;
+                break;
+            case TIMESTAMPMICRO:
+            case TIMESTAMPMICROTZ:
+                factor = 1_000;
+                break;
+            case TIMESTAMPMILLI:
+            case TIMESTAMPMILLITZ:
+                factor = 1_000_000;
+                break;
+            case TIMESTAMPSEC:
+            case TIMESTAMPSECTZ:
+                factor = 1_000_000_000;
+                break;
+            default:
+                throw new IllegalArgumentException("Unsupported minor type: " + minorType);
+        }
     }
 
     @Override
-    public void fillChunk(final ChunkSource. @NotNull FillContext context,
+    public void fillChunk(final ChunkSource.@NotNull FillContext context,
             final @NotNull WritableChunk<? super Values> destination,
             final @NotNull RowSequence rowSequence) {
         final WritableObjectChunk<Instant, ? super Values> chunk = destination.asWritableObjectChunk();
         final ArrowWrapperTools.FillContext arrowContext = (ArrowWrapperTools.FillContext) context;
         chunk.setSize(0);
-        fillChunk(arrowContext, rowSequence, rowKey -> chunk.add(extract(getPositionInBlock(rowKey), arrowContext.getVector(field))));
+        fillChunk(arrowContext, rowSequence,
+                rowKey -> chunk.add(extract(getPositionInBlock(rowKey), arrowContext.getVector(field))));
     }
 
     @Override
@@ -48,6 +73,7 @@ public final Instant get(final long rowKey) {
     }
 
     private Instant extract(final int posInBlock, final TimeStampVector vector) {
-        return vector.isSet(posInBlock) == 0 ? null : io.deephaven.time.DateTimeUtils.epochNanosToInstant(vector.get(posInBlock));
+        final long value = vector.get(posInBlock);
+        return vector.isSet(posInBlock) == 0 ? null : DateTimeUtils.epochNanosToInstant(value * factor);
     }
 }
diff --git a/...ication/reflective/src/main/java/io/deephaven/replicators/GenerateArrowColumnSources.java b/...ication/reflective/src/main/java/io/deephaven/replicators/GenerateArrowColumnSources.java
@@ -122,15 +122,6 @@ public static void main(String[] args) {
                                 .addStatement("$T localDateTime = vector.getObject(posInBlock)", LocalDateTime.class)
                                 .addStatement("return localDateTime != null ? localDateTime.toLocalTime() : null")
                                 .build()));
-
-        final ClassName instant = ClassName.get("java.time", "Instant");
-        generateArrowColumnSource("ArrowInstantColumnSource", instant, instant, TimeStampVector.class,
-                "get", "ForObject", WritableObjectChunk.class, List.of(
-                        preparePrivateExtractMethod(instant, TimeStampVector.class)
-                                .addStatement(
-                                        "return vector.isSet(posInBlock) == 0 ? null : io.deephaven.time.DateTimeUtils.epochNanosToInstant(vector.get(posInBlock))",
-                                        instant)
-                                .build()));
     }
 
     @NotNull