From c4803a052bcd683d9bb4baf9b5e42ec0dacbfd40 Mon Sep 17 00:00:00 2001
From: AbdelrahmanElhawary <abdulrahman.elhawary@incorta.com>
Date: Thu, 9 Jul 2026 14:13:45 +0300
Subject: [PATCH] Arrow: Fix TIMESTAMP_MILLIS scaling for dictionary-encoded
 columns

---
 .../vectorized/VectorizedArrowReader.java     |  56 ++++++++
 ...stVectorizedTimestampMillisDictionary.java | 122 ++++++++++++++++++
 2 files changed, 178 insertions(+)
 create mode 100644 arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedTimestampMillisDictionary.java

diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
index 11c73aea4a44..578e25ee1970 100644
--- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
+++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
@@ -384,6 +384,62 @@ public void setRowGroupInfo(PageReadStore source, Map<ColumnPath, ColumnChunkMet
         vectorizedColumnIterator.setRowGroupInfo(
             source.getPageReader(columnDescriptor),
             !ParquetUtil.hasNonDictionaryPages(chunkMetaData));
+
+    boolean isTimestampMillis = false;
+    if (columnDescriptor != null && columnDescriptor.getPrimitiveType() != null) {
+      org.apache.parquet.schema.LogicalTypeAnnotation annotation =
+          columnDescriptor.getPrimitiveType().getLogicalTypeAnnotation();
+
+      if (annotation
+          instanceof
+          org.apache.parquet.schema.LogicalTypeAnnotation.TimestampLogicalTypeAnnotation) {
+        org.apache.parquet.schema.LogicalTypeAnnotation.TimestampLogicalTypeAnnotation
+            timestampAnnotation =
+                (org.apache.parquet.schema.LogicalTypeAnnotation.TimestampLogicalTypeAnnotation)
+                    annotation;
+
+        isTimestampMillis =
+            timestampAnnotation.getUnit()
+                == org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MILLIS;
+      }
+    }
+
+    if (this.dictionary != null && isTimestampMillis) {
+      final Dictionary backingDictionary = this.dictionary;
+
+      this.dictionary =
+          new Dictionary(backingDictionary.getEncoding()) {
+            @Override
+            public long decodeToLong(int id) {
+              return backingDictionary.decodeToLong(id) * 1000L;
+            }
+
+            @Override
+            public int decodeToInt(int id) {
+              return backingDictionary.decodeToInt(id);
+            }
+
+            @Override
+            public float decodeToFloat(int id) {
+              return backingDictionary.decodeToFloat(id);
+            }
+
+            @Override
+            public double decodeToDouble(int id) {
+              return backingDictionary.decodeToDouble(id);
+            }
+
+            @Override
+            public org.apache.parquet.io.api.Binary decodeToBinary(int id) {
+              return backingDictionary.decodeToBinary(id);
+            }
+
+            @Override
+            public int getMaxId() {
+              return backingDictionary.getMaxId();
+            }
+          };
+    }
   }
 
   @Override
diff --git a/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedTimestampMillisDictionary.java b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedTimestampMillisDictionary.java
new file mode 100644
index 000000000000..35db8fe8d6fb
--- /dev/null
+++ b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedTimestampMillisDictionary.java
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.arrow.vectorized;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.lang.reflect.Field;
+import java.util.Map;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.iceberg.arrow.vectorized.parquet.VectorizedColumnIterator;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.Types;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.Dictionary;
+import org.apache.parquet.column.Encoding;
+import org.apache.parquet.column.page.PageReadStore;
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.schema.LogicalTypeAnnotation;
+import org.apache.parquet.schema.PrimitiveType;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.MockedConstruction;
+import org.mockito.Mockito;
+
+public class TestVectorizedTimestampMillisDictionary {
+
+  private BufferAllocator allocator;
+
+  @BeforeEach
+  public void before() {
+    this.allocator = new RootAllocator(Long.MAX_VALUE);
+  }
+
+  @AfterEach
+  public void after() {
+    if (allocator != null) {
+      allocator.close();
+    }
+  }
+
+  @Test
+  public void testTimestampMillisDictionaryProxyScaling() throws Exception {
+    // 1. Mock Parquet descriptor structures to avoid package-private constructor limitations
+    ColumnDescriptor mockDescriptor = Mockito.mock(ColumnDescriptor.class);
+    PrimitiveType mockPrimitiveType = Mockito.mock(PrimitiveType.class);
+    LogicalTypeAnnotation.TimestampLogicalTypeAnnotation mockAnnotation =
+        Mockito.mock(LogicalTypeAnnotation.TimestampLogicalTypeAnnotation.class);
+
+    Mockito.when(mockDescriptor.getPrimitiveType()).thenReturn(mockPrimitiveType);
+    Mockito.when(mockPrimitiveType.getLogicalTypeAnnotation()).thenReturn(mockAnnotation);
+    Mockito.when(mockAnnotation.getUnit()).thenReturn(LogicalTypeAnnotation.TimeUnit.MILLIS);
+    Mockito.when(mockDescriptor.getPath())
+        .thenReturn(new String[] {"reserved_incorta_extraction_timestamp"});
+
+    Types.NestedField icebergField =
+        Types.NestedField.optional(
+            1, "reserved_incorta_extraction_timestamp", Types.TimestampType.withZone());
+
+    PageReadStore mockPageReadStore = Mockito.mock(PageReadStore.class);
+    ColumnChunkMetaData mockChunkMeta = Mockito.mock(ColumnChunkMetaData.class);
+    Dictionary mockParquetDictionary = Mockito.mock(Dictionary.class);
+
+    ColumnPath columnPath = ColumnPath.get("reserved_incorta_extraction_timestamp");
+    Map<ColumnPath, ColumnChunkMetaData> metadataMap = Maps.newHashMap();
+    metadataMap.put(columnPath, mockChunkMeta);
+
+    Mockito.when(mockParquetDictionary.getEncoding()).thenReturn(Encoding.PLAIN_DICTIONARY);
+    Mockito.when(mockParquetDictionary.decodeToLong(0)).thenReturn(1783426418473L);
+
+    // 2. Intercept internal instantiation of VectorizedColumnIterator using construction mocking
+    try (MockedConstruction<VectorizedColumnIterator> mockedConstruction =
+        Mockito.mockConstruction(
+            VectorizedColumnIterator.class,
+            (mockIterator, context) -> {
+              Mockito.when(mockIterator.setRowGroupInfo(Mockito.any(), Mockito.anyBoolean()))
+                  .thenReturn(mockParquetDictionary);
+            })) {
+
+      VectorizedArrowReader reader =
+          new VectorizedArrowReader(mockDescriptor, icebergField, allocator, true);
+
+      // 3. Execute target method
+      reader.setRowGroupInfo(mockPageReadStore, metadataMap);
+
+      // 4. Extract private field using Java Reflection to bypass missing getters
+      Field dictionaryField = VectorizedArrowReader.class.getDeclaredField("dictionary");
+      dictionaryField.setAccessible(true);
+      Dictionary activeDictionary = (Dictionary) dictionaryField.get(reader);
+
+      // 5. Run Verification Assertions
+      assertThat(activeDictionary)
+          .as("The internal dictionary reference must be substituted by our proxy class")
+          .isNotNull()
+          .isNotEqualTo(mockParquetDictionary);
+
+      long expectedMicros = 1783426418473L * 1000L;
+      assertThat(activeDictionary.decodeToLong(0))
+          .as(
+              "The proxy hook must intercept lookups and multiply millisecond properties to microseconds")
+          .isEqualTo(expectedMicros);
+    }
+  }
+}