address PR comments

Author: divjotarora

parquet-column/src/main/java/org/apache/parquet/schema/PrimitiveComparator.java

@@ -20,6 +20,7 @@
 
 import java.io.Serializable;
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.Comparator;
 import org.apache.parquet.io.api.Binary;
 
@@ -288,29 +289,16 @@ int compareBinary(Binary b1, Binary b2) {
         throw new IllegalArgumentException(
             "Timestamp binary length must be 12 bytes, got " + b1.length() + " and " + b2.length());
       }
-
-      ByteBuffer bb1 = b1.toByteBuffer();
-      ByteBuffer bb2 = b2.toByteBuffer();
-      // The buffers may be slices with a non-zero position (e.g. a ByteBuffer-backed Binary), so
-      // index relative to position() rather than absolute offset 0. Byte 11 (position + 11) is the
-      // most significant byte and carries the sign; byte 0 (position) is least significant.
-      int p1 = bb1.position();
-      int p2 = bb2.position();
-
-      // If one value is negative and the other is positive, one is trivially greater.
-      boolean neg1 = bb1.get(p1 + 11) < 0;
-      boolean neg2 = bb2.get(p2 + 11) < 0;
-      if (neg1 != neg2) {
-        return neg1 ? -1 : 1;
+      ByteBuffer bb1 = b1.toByteBuffer().slice();
+      bb1.order(ByteOrder.LITTLE_ENDIAN);
+      ByteBuffer bb2 = b2.toByteBuffer().slice();
+      bb2.order(ByteOrder.LITTLE_ENDIAN);
+      // Signed comparison of the high 4 bytes followed by unsigned comparison of the low 8 bytes.
+      int hiResult = Integer.compare(bb1.getInt(8), bb2.getInt(8));
+      if (hiResult != 0) {
+        return hiResult;
       }
-
-      for (int i = 11; i >= 0; --i) {
-        int diff = toUnsigned(bb1.get(p1 + i)) - toUnsigned(bb2.get(p2 + i));
-        if (diff != 0) {
-          return diff;
-        }
-      }
-      return 0;
+      return Long.compareUnsigned(bb1.getLong(0), bb2.getLong(0));
     }
 
     @Override

parquet-column/src/main/java/org/apache/parquet/schema/Types.java

@@ -563,10 +563,13 @@ public Optional<Boolean> visit(
               @Override
               public Optional<Boolean> visit(
                   LogicalTypeAnnotation.TimestampLogicalTypeAnnotation timestampLogicalType) {
-                if (primitiveType == PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY) {
-                  return checkFixedPrimitiveType(12, timestampLogicalType);
-                }
-                return checkInt64PrimitiveType(timestampLogicalType);
+                Preconditions.checkState(
+                    primitiveType == PrimitiveTypeName.INT64
+                        || (primitiveType == PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY
+                            && length == 12),
+                    "%s can only annotate INT64 or FIXED_LEN_BYTE_ARRAY(12)",
+                    timestampLogicalType);
+                return Optional.of(true);
               }
 
               @Override

parquet-column/src/test/java/org/apache/parquet/schema/TestPrimitiveComparator.java

@@ -455,6 +455,10 @@ public void testBinaryAsSignedIntegerLE12Comparator() {
     // an LSB-first comparator would wrongly rank +256 < +1 (seeing byte[0]=0 < 1 first)
     byte[] posTwoFiftySix = new byte[12];
     posTwoFiftySix[1] = 1;
+    // +2^63: low 8 bytes have bit 63 set (LE: bytes 0..6=0x00, byte 7=0x80), high 4 bytes=0x00.
+    // Sits between small positives and largePos; fails if the lower 8 bytes are compared signed.
+    byte[] posTwo63 = new byte[12];
+    posTwo63[7] = (byte) 0x80;
     // large positive: 0x7F FF..FF (most positive 96-bit value)
     byte[] largePos = new byte[12];
     for (int i = 0; i < 12; i++) largePos[i] = (byte) 0xFF;
@@ -467,6 +471,7 @@ public void testBinaryAsSignedIntegerLE12Comparator() {
         Binary.fromConstantByteArray(zero),
         Binary.fromConstantByteArray(posOne),
         Binary.fromConstantByteArray(posTwoFiftySix),
+        Binary.fromConstantByteArray(posTwo63),
         Binary.fromConstantByteArray(largePos));
   }
 

parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuilders.java

@@ -18,7 +18,6 @@
  */
 package org.apache.parquet.schema;
 
-import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MICROS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MILLIS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.timestampType;
@@ -535,7 +534,8 @@ public void testInt32AnnotationsRejectNonInt32() {
 
   @Test
   public void testInt64Annotations() {
-    OriginalType[] types = new OriginalType[] {TIME_MICROS, TIMESTAMP_MILLIS, TIMESTAMP_MICROS, UINT_64, INT_64};
+    // Test the non-timestamp annotations for INT64. Timestamps are tested separately below.
+    OriginalType[] types = new OriginalType[] {TIME_MICROS, UINT_64, INT_64};
     for (OriginalType logicalType : types) {
       PrimitiveType expected = new PrimitiveType(REQUIRED, INT64, "col", logicalType);
       PrimitiveType date = Types.required(INT64).as(logicalType).named("col");
@@ -545,7 +545,8 @@ public void testInt64Annotations() {
 
   @Test
   public void testInt64AnnotationsRejectNonInt64() {
-    OriginalType[] types = new OriginalType[] {TIME_MICROS, TIMESTAMP_MILLIS, TIMESTAMP_MICROS, UINT_64, INT_64};
+    // Test the non-timestamp annotations for INT64. Timestamps are tested separately below.
+    OriginalType[] types = new OriginalType[] {TIME_MICROS, UINT_64, INT_64};
     for (final OriginalType logicalType : types) {
       PrimitiveTypeName[] nonInt64 = new PrimitiveTypeName[] {BOOLEAN, INT32, INT96, DOUBLE, FLOAT, BINARY};
       for (final PrimitiveTypeName type : nonInt64) {
@@ -554,19 +555,54 @@ public void testInt64AnnotationsRejectNonInt64() {
             .hasMessage(
                 LogicalTypeAnnotation.fromOriginalType(logicalType, null) + " can only annotate INT64");
       }
-      // TIMESTAMP allows FLBA(12); other lengths are still rejected.
-      // Non-timestamp types still only accept INT64 for any FLBA length.
-      boolean isTimestamp = logicalType == TIMESTAMP_MILLIS || logicalType == TIMESTAMP_MICROS;
-      String flbaErrMsg = isTimestamp
-          ? LogicalTypeAnnotation.fromOriginalType(logicalType, null)
-              + " can only annotate FIXED_LEN_BYTE_ARRAY(12)"
-          : LogicalTypeAnnotation.fromOriginalType(logicalType, null) + " can only annotate INT64";
       assertThatThrownBy(() -> Types.required(FIXED_LEN_BYTE_ARRAY)
               .length(1)
               .as(logicalType)
               .named("col"))
           .isInstanceOf(IllegalStateException.class)
-          .hasMessage(flbaErrMsg);
+          .hasMessage(LogicalTypeAnnotation.fromOriginalType(logicalType, null) + " can only annotate INT64");
+    }
+  }
+
+  @Test
+  public void testTimestampAnnotations() {
+    // Test timestamp annotations for both INT64 and FLBA(12).
+    OriginalType[] types = new OriginalType[] {TIMESTAMP_MILLIS, TIMESTAMP_MICROS};
+    for (OriginalType logicalType : types) {
+      PrimitiveType expectedInt64 = new PrimitiveType(REQUIRED, INT64, "col", logicalType);
+      PrimitiveType dateInt64 = Types.required(INT64).as(logicalType).named("col");
+      assertThat(dateInt64).isEqualTo(expectedInt64);
+
+      PrimitiveType expectedFlba12 = new PrimitiveType(REQUIRED, FIXED_LEN_BYTE_ARRAY, 12, "col", logicalType);
+      PrimitiveType dateFlba12 = Types.required(FIXED_LEN_BYTE_ARRAY)
+          .length(12)
+          .as(logicalType)
+          .named("col");
+      assertThat(dateFlba12).isEqualTo(expectedFlba12);
+    }
+  }
+
+  @Test
+  public void testTimestampAnnotationsRejectNonTimestamp() {
+    // Test timestamp annotations for both INT64 and FLBA(12).
+    OriginalType[] types = new OriginalType[] {TIMESTAMP_MILLIS, TIMESTAMP_MICROS};
+    for (OriginalType logicalType : types) {
+      // Invalid primitive types are rejected.
+      PrimitiveTypeName[] nonTimestamp = new PrimitiveTypeName[] {BOOLEAN, INT32, INT96, DOUBLE, FLOAT, BINARY};
+      for (PrimitiveTypeName type : nonTimestamp) {
+        assertThatThrownBy(() -> Types.required(type).as(logicalType).named("col"))
+            .isInstanceOf(IllegalStateException.class)
+            .hasMessage(LogicalTypeAnnotation.fromOriginalType(logicalType, null)
+                + " can only annotate INT64 or FIXED_LEN_BYTE_ARRAY(12)");
+      }
+      // Invalid FLBA lengths are rejected.
+      assertThatThrownBy(() -> Types.required(FIXED_LEN_BYTE_ARRAY)
+              .length(11)
+              .as(logicalType)
+              .named("col"))
+          .isInstanceOf(IllegalStateException.class)
+          .hasMessage(LogicalTypeAnnotation.fromOriginalType(logicalType, null)
+              + " can only annotate INT64 or FIXED_LEN_BYTE_ARRAY(12)");
     }
   }
 
@@ -1451,23 +1487,6 @@ public void testTimestampLogicalTypeWithUTCParameter() {
     assertThat(nonUtcMicrosActual).isEqualTo(nonUtcMicrosExpected);
   }
 
-  @Test
-  public void testTimestampFlba12LogicalType() {
-    for (TimeUnit unit : TimeUnit.values()) {
-      // FLBA(12) with TIMESTAMP annotation is valid.
-      Types.required(FIXED_LEN_BYTE_ARRAY)
-          .length(12)
-          .as(timestampType(true, unit))
-          .named("ts");
-      // Other FLBA lengths must be rejected.
-      assertThatThrownBy(() -> Types.required(FIXED_LEN_BYTE_ARRAY)
-              .length(11)
-              .as(timestampType(true, unit))
-              .named("ts"))
-          .isInstanceOf(IllegalStateException.class);
-    }
-  }
-
   @Test
   public void testVariantLogicalType() {
     byte specVersion = 1;

parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java

@@ -318,7 +318,14 @@ public void visit(PrimitiveType primitiveType) {
         element.setRepetition_type(toParquetRepetition(primitiveType.getRepetition()));
         element.setType(getType(primitiveType.getPrimitiveTypeName()));
         if (primitiveType.getLogicalTypeAnnotation() != null) {
-          element.setConverted_type(convertToConvertedType(primitiveType.getLogicalTypeAnnotation()));
+          // The TimestampType logical type may have a converted type, but only for the INT64
+          // physical type.
+          boolean suppressConvertedType = primitiveType.getLogicalTypeAnnotation()
+                  instanceof LogicalTypeAnnotation.TimestampLogicalTypeAnnotation
+              && primitiveType.getPrimitiveTypeName() != PrimitiveTypeName.INT64;
+          if (!suppressConvertedType) {
+            element.setConverted_type(convertToConvertedType(primitiveType.getLogicalTypeAnnotation()));
+          }
           element.setLogicalType(convertToLogicalType(primitiveType.getLogicalTypeAnnotation()));
         }
         if (primitiveType.getDecimalMetadata() != null) {

parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java

@@ -449,6 +449,18 @@ public void testTimeLogicalTypes() {
         .required(PrimitiveTypeName.INT64)
         .as(timestampType(true, NANOS))
         .named("aTimestampUtcNanos")
+        .required(PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY)
+        .length(12)
+        .as(timestampType(true, MILLIS))
+        .named("aTimestampFlbaMillis")
+        .required(PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY)
+        .length(12)
+        .as(timestampType(true, MICROS))
+        .named("aTimestampFlbaMicros")
+        .required(PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY)
+        .length(12)
+        .as(timestampType(true, NANOS))
+        .named("aTimestampFlbaNanos")
         .required(PrimitiveTypeName.INT32)
         .as(timeType(false, MILLIS))
         .named("aTimeNonUtcMillis")
@@ -469,6 +481,19 @@ public void testTimeLogicalTypes() {
         .named("aTimeUtcNanos")
         .named("Message");
     List<SchemaElement> parquetSchema = parquetMetadataConverter.toParquetSchema(expected);
+    // FLBA(12) MILLIS/MICROS must not write a legacy converted_type (it is INT64-only).
+    SchemaElement flbaMillis = parquetSchema.stream()
+        .filter(e -> "aTimestampFlbaMillis".equals(e.getName()))
+        .findFirst()
+        .get();
+    assertThat(flbaMillis.isSetConverted_type()).isFalse();
+    assertThat(flbaMillis.isSetLogicalType()).isTrue();
+    SchemaElement flbaMicros = parquetSchema.stream()
+        .filter(e -> "aTimestampFlbaMicros".equals(e.getName()))
+        .findFirst()
+        .get();
+    assertThat(flbaMicros.isSetConverted_type()).isFalse();
+    assertThat(flbaMicros.isSetLogicalType()).isTrue();
     MessageType schema = parquetMetadataConverter.fromParquetSchema(parquetSchema, null);
     assertThat(schema).isEqualTo(expected);
   }