Fix a bunch of issues, still WIP

cashmand · cashmand · commit a65d25e8760a · 2025-05-08T14:35:17.000-04:00
diff --git a/parquet-avro/src/test/java/org/apache/parquet/avro/TestWriteVariant.java b/parquet-avro/src/test/java/org/apache/parquet/avro/TestWriteVariant.java
@@ -19,6 +19,7 @@
 package org.apache.parquet.avro;
 
 import static org.apache.parquet.avro.AvroTestUtil.*;
+import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.*;
 import static org.junit.Assert.*;
 
 import java.io.File;
@@ -99,7 +100,7 @@ private static ByteBuffer variant(String s) {
 
   private static MessageType parquetSchema(GroupType variantGroup) {
     return Types.buildMessage()
-        .required(PrimitiveTypeName.INT32)
+        .required(INT32)
         .named("id")
         .addField(variantGroup)
         .named("table");
@@ -231,7 +232,7 @@ public void testUnshreddedValues() throws IOException {
   public void testShreddedValues() throws IOException {
     for (Variant v : VARIANTS) {
       GenericRecord record = createRecord(1, v);
-      MessageType writeSchema = shreddingFromValue(v);
+      MessageType writeSchema = shreddingSchema(v);
       TestSchema testSchema = new TestSchema(writeSchema, readSchema);
 
       GenericRecord actual = writeAndRead(testSchema, record);
@@ -241,24 +242,21 @@ public void testShreddedValues() throws IOException {
     }
   }
 
-//  @Test
-//  public void testMixedShredding() throws IOException {
-//    for (Variant v : VARIANTS) {
-//      List<Record> expected =
-//          IntStream.range(0, VARIANTS.length)
-//              .mapToObj(i -> RECORD.copy("id", i, "var", VARIANTS[i]))
-//              .collect(Collectors.toList());
-//
-//      List<Record> actual =
-//          writeAndRead((id, name) -> ParquetVariantUtil.toParquetSchema(v.value()), expected);
-//
-//      assertThat(actual.size()).isEqualTo(expected.size());
-//
-//      for (int i = 0; i < expected.size(); i += 1) {
-//        InternalTestHelpers.assertEquals(SCHEMA.asStruct(), expected.get(i), actual.get(i));
-//      }
-//    }
-//  }
+  @Test
+  public void testMixedShredding() throws IOException {
+    for (Variant v : VARIANTS) {
+      List<GenericRecord> expected = new ArrayList<>();
+      for (int i = 0; i < VARIANTS.length; i++) {
+        expected.add(createRecord(i, VARIANTS[i]));
+      }
+
+      MessageType writeSchema = shreddingSchema(v);
+      TestSchema testSchema = new TestSchema(writeSchema, readSchema);
+
+      List<GenericRecord> actual = writeAndRead(testSchema, expected);
+      // TODO: CHECK RESULTS
+    }
+  }
 
   // Write schema contains the full shredding schema. Read schema should just be a value/metadata pair.
   private static class TestSchema {
@@ -336,4 +334,149 @@ private List<GenericRecord> writeAndRead(
     }
     return result;
   }
+
+  /**
+   * Build a shredding schema that will perfectly shred the provided value.
+   */
+  private static MessageType shreddingSchema(Variant v) {
+    Type shreddedType = shreddedType(v);
+    Types.GroupBuilder<GroupType> partialType = Types.buildGroup(Type.Repetition.OPTIONAL)
+        .as(LogicalTypeAnnotation.variantType((byte) 1))
+        .required(BINARY)
+        .named("metadata")
+        .optional(BINARY)
+        .named("value");
+    Type variantType;
+    if (shreddedType == null) {
+      variantType = partialType.named("var");
+    } else {
+      variantType = partialType.addField(shreddedType).named("var");
+    }
+    return Types.buildMessage()
+        .required(INT32)
+        .named("id")
+        .addField(variantType)
+        .named("table");
+  }
+
+  private static GroupType shreddedGroup(Variant v, String name) {
+    Type shreddedType = shreddedType(v);
+    if (shreddedType == null) {
+      return Types.buildGroup(Type.Repetition.OPTIONAL)
+          .optional(BINARY)
+          .named("value")
+          .named(name);
+    } else {
+      return Types.buildGroup(Type.Repetition.OPTIONAL)
+          .optional(BINARY)
+          .named("value")
+          .addField(shreddedType)
+          .named(name);
+    }
+  }
+
+  /**
+   * @return A shredded type, or null if there is no valid shredded type.
+   */
+  private static Type shreddedType(Variant v) {
+    switch (v.getType()) {
+      case NULL:
+        return null;
+      case BOOLEAN:
+        return Types.optional(BOOLEAN).named("typed_value");
+      case BYTE:
+        return Types.optional(INT32).as(LogicalTypeAnnotation.intType(8)).named("typed_value");
+      case SHORT:
+        return Types.optional(INT32).as(LogicalTypeAnnotation.intType(16)).named("typed_value");
+      case INT:
+        return Types.optional(INT32).named("typed_value");
+      case LONG:
+        return Types.optional(INT64).named("typed_value");
+      case FLOAT:
+        return Types.optional(FLOAT).named("typed_value");
+      case DOUBLE:
+        return Types.optional(DOUBLE).named("typed_value");
+      case DECIMAL4:
+        return Types.optional(INT32)
+            .as(LogicalTypeAnnotation.decimalType(v.getDecimal().scale(), 9))
+            .named("typed_value");
+      case DECIMAL8:
+        return Types.optional(INT64)
+            .as(LogicalTypeAnnotation.decimalType(v.getDecimal().scale(), 18))
+            .named("typed_value");
+      case DECIMAL16:
+        return Types.optional(BINARY)
+            .as(LogicalTypeAnnotation.decimalType(v.getDecimal().scale(), 38))
+            .named("typed_value");
+      case DATE:
+        return Types.optional(INT32)
+            .as(LogicalTypeAnnotation.dateType())
+            .named("typed_value");
+      case TIMESTAMP_TZ:
+        return Types.optional(INT64)
+            .as(LogicalTypeAnnotation.timestampType(true, TimeUnit.MICROS))
+            .named("typed_value");
+      case TIMESTAMP_NTZ:
+        return Types.optional(INT64)
+            .as(LogicalTypeAnnotation.timestampType(false, TimeUnit.MICROS))
+            .named("typed_value");
+      case BINARY:
+        return Types.optional(BINARY).named("typed_value");
+      case STRING:
+        return Types.optional(BINARY)
+            .as(LogicalTypeAnnotation.stringType())
+            .named("typed_value");
+      case TIME:
+        return Types.optional(INT64)
+            .as(LogicalTypeAnnotation.timeType(false, TimeUnit.MICROS))
+            .named("typed_value");
+      case TIMESTAMP_NANOS_TZ:
+        return Types.optional(INT64)
+            .as(LogicalTypeAnnotation.timestampType(true, TimeUnit.NANOS))
+            .named("typed_value");
+      case TIMESTAMP_NANOS_NTZ:
+        return Types.optional(INT64)
+            .as(LogicalTypeAnnotation.timestampType(false, TimeUnit.NANOS))
+            .named("typed_value");
+      case UUID:
+        return Types.optional(FIXED_LEN_BYTE_ARRAY)
+            .as(LogicalTypeAnnotation.uuidType())
+            .named("typed_value");
+      case OBJECT:
+        return shreddedObjectType(v);
+      case ARRAY:
+        return shreddedArrayType(v);
+      default:
+        throw new UnsupportedOperationException("Unsupported shredding type: " + v.getType());
+    }
+  }
+
+  private static Type shreddedObjectType(Variant v) {
+    if (v.numObjectElements() == 0) {
+      // Parquet can't represent empty groups.
+      return null;
+    }
+    Types.GroupBuilder<GroupType> builder = Types.optionalGroup();
+    for (int i = 0; i < v.numObjectElements(); i++) {
+      Variant.ObjectField field = v.getFieldAtIndex(i);
+      Types.GroupBuilder<GroupType> fieldBuilder = Types.optionalGroup();
+      Type fieldType = shreddedGroup(field.value, field.key);
+      builder.addField(fieldType);
+    }
+    return builder.named("typed_value");
+  }
+
+  private static Type shreddedArrayType(Variant v) {
+    // Use the first element to determine the array element type
+    Variant firstElement;
+    if (v.numArrayElements() > 0) {
+      firstElement = v.getElementAtIndex(0);
+    } else {
+      // Use null as a dummy value, which will omit typed_value from the schema.
+      firstElement = fullVariant(b -> b.appendNull());
+    }
+
+    Type elementType = shreddedGroup(firstElement, "element");
+    return Types.optionalList().setElementType(elementType).named("typed_value");
+  }
 }
diff --git a/parquet-variant/src/main/java/org/apache/parquet/variant/VariantBuilder.java b/parquet-variant/src/main/java/org/apache/parquet/variant/VariantBuilder.java
@@ -151,9 +151,9 @@ public byte[] valueWithoutMetadata() {
    */
   void shallowAppendVariant(Binary value) {
     onAppend();
-    int size = value.length();
-    checkCapacity(size);
     byte[] buf = value.getBytes();
+    int size = VariantUtil.valueSize(ByteBuffer.wrap(buf), 0);
+    checkCapacity(size);
     System.arraycopy(buf, 0, writeBuffer, writePos, size);
     writePos += size;
   }
diff --git a/parquet-variant/src/main/java/org/apache/parquet/variant/VariantUtil.java b/parquet-variant/src/main/java/org/apache/parquet/variant/VariantUtil.java
@@ -843,4 +843,73 @@ static HashMap<String, Integer> getMetadataMap(ByteBuffer metadata) {
     }
     return result;
   }
+
+  /**
+   * Computes the actual size (in bytes) of the Variant value at `value[pos...]`
+   * @param value The Variant value
+   * @param pos The starting index of the Variant value
+   * @return The size (in bytes) of the Variant value
+   */
+  public static int valueSize(ByteBuffer value, int pos) {
+    checkIndex(pos, value.limit());
+    int basicType = value.get(pos) & BASIC_TYPE_MASK;
+    int typeInfo = (value.get(pos) >> BASIC_TYPE_BITS) & PRIMITIVE_TYPE_MASK;
+    switch (basicType) {
+      case SHORT_STR:
+        return 1 + typeInfo;
+      case OBJECT: {
+        VariantUtil.ObjectInfo info = VariantUtil.getObjectInfo(slice(value, pos));
+        return info.dataStartOffset
+            + readUnsigned(
+            value,
+            pos + info.offsetStartOffset + info.numElements * info.offsetSize,
+            info.offsetSize);
+      }
+      case ARRAY: {
+        VariantUtil.ArrayInfo info = VariantUtil.getArrayInfo(slice(value, pos));
+        return info.dataStartOffset
+            + readUnsigned(
+            value,
+            pos + info.offsetStartOffset + info.numElements * info.offsetSize,
+            info.offsetSize);
+      }
+      default:
+        switch (typeInfo) {
+          case NULL:
+          case TRUE:
+          case FALSE:
+            return 1;
+          case INT8:
+            return 2;
+          case INT16:
+            return 3;
+          case INT32:
+          case DATE:
+          case FLOAT:
+            return 5;
+          case INT64:
+          case DOUBLE:
+          case TIMESTAMP_TZ:
+          case TIMESTAMP_NTZ:
+          case TIME:
+          case TIMESTAMP_NANOS_TZ:
+          case TIMESTAMP_NANOS_NTZ:
+            return 9;
+          case DECIMAL4:
+            return 6;
+          case DECIMAL8:
+            return 10;
+          case DECIMAL16:
+            return 18;
+          case BINARY:
+          case LONG_STR:
+            return 1 + U32_SIZE + readUnsigned(value, pos + 1, U32_SIZE);
+          case UUID:
+            return 1 + UUID_SIZE;
+          default:
+            throw new UnsupportedOperationException(
+                String.format("Unknown type in Variant. primitive type: %d", typeInfo));
+        }
+    }
+  }
 }
diff --git a/parquet-variant/src/main/java/org/apache/parquet/variant/VariantValueWriter.java b/parquet-variant/src/main/java/org/apache/parquet/variant/VariantValueWriter.java
