GH-3627: Add missing onAppend() to VariantBuilder#appendUUIDBytes (#3624)

Author: nastra

parquet-variant/src/main/java/org/apache/parquet/variant/VariantBuilder.java

@@ -420,6 +420,7 @@ public void appendUUID(java.util.UUID uuid) {
    * @param bytes a 16-byte value.
    */
   void appendUUIDBytes(ByteBuffer bytes) {
+    onAppend();
     checkCapacity(1 + VariantUtil.UUID_SIZE);
     writeBuffer[writePos++] = VariantUtil.primitiveHeader(VariantUtil.UUID);
     if (bytes.remaining() < VariantUtil.UUID_SIZE) {

parquet-variant/src/test/java/org/apache/parquet/variant/TestVariantArrayBuilder.java

@@ -18,6 +18,9 @@
  */
 package org.apache.parquet.variant;
 
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.UUID;
 import org.junit.Assert;
 import org.junit.Test;
 import org.slf4j.Logger;
@@ -26,6 +29,34 @@
 public class TestVariantArrayBuilder {
   private static final Logger LOG = LoggerFactory.getLogger(TestVariantArrayBuilder.class);
 
+  @Test
+  public void testArrayBuilderWithUUIDBytes() {
+    byte[] uuid = new byte[] {0, 17, 34, 51, 68, 85, 102, 119, -120, -103, -86, -69, -52, -35, -18, -1};
+    long msb = ByteBuffer.wrap(uuid, 0, 8).order(ByteOrder.BIG_ENDIAN).getLong();
+    long lsb = ByteBuffer.wrap(uuid, 8, 8).order(ByteOrder.BIG_ENDIAN).getLong();
+    UUID expected = new UUID(msb, lsb);
+
+    VariantBuilder builder = new VariantBuilder();
+    VariantArrayBuilder array = builder.startArray();
+    array.appendInt(1);
+    // appendUUIDBytes must go through onAppend() so that the element offset is recorded and
+    // numValues is incremented. Otherwise the offset list is wrong and the UUID element is lost.
+    array.appendUUIDBytes(ByteBuffer.wrap(uuid));
+    array.appendInt(2);
+    builder.endArray();
+
+    VariantTestUtil.testVariant(builder.build(), v -> {
+      VariantTestUtil.checkType(v, VariantUtil.ARRAY, Variant.Type.ARRAY);
+      Assert.assertEquals(3, v.numArrayElements());
+      VariantTestUtil.checkType(v.getElementAtIndex(0), VariantUtil.PRIMITIVE, Variant.Type.INT);
+      Assert.assertEquals(1, v.getElementAtIndex(0).getInt());
+      VariantTestUtil.checkType(v.getElementAtIndex(1), VariantUtil.PRIMITIVE, Variant.Type.UUID);
+      Assert.assertEquals(expected, v.getElementAtIndex(1).getUUID());
+      VariantTestUtil.checkType(v.getElementAtIndex(2), VariantUtil.PRIMITIVE, Variant.Type.INT);
+      Assert.assertEquals(2, v.getElementAtIndex(2).getInt());
+    });
+  }
+
   @Test
   public void testEmptyArrayBuilder() {
     VariantBuilder b = new VariantBuilder();

parquet-variant/src/test/java/org/apache/parquet/variant/TestVariantObjectBuilder.java

@@ -19,6 +19,8 @@
 package org.apache.parquet.variant;
 
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.UUID;
 import org.junit.Assert;
 import org.junit.Ignore;
 import org.junit.Test;
@@ -28,6 +30,29 @@
 public class TestVariantObjectBuilder {
   private static final Logger LOG = LoggerFactory.getLogger(TestVariantObjectBuilder.class);
 
+  @Test
+  public void testObjectBuilderWithUUIDBytes() {
+    byte[] uuid = new byte[] {0, 17, 34, 51, 68, 85, 102, 119, -120, -103, -86, -69, -52, -35, -18, -1};
+    long msb = ByteBuffer.wrap(uuid, 0, 8).order(ByteOrder.BIG_ENDIAN).getLong();
+    long lsb = ByteBuffer.wrap(uuid, 8, 8).order(ByteOrder.BIG_ENDIAN).getLong();
+    UUID expected = new UUID(msb, lsb);
+
+    VariantBuilder builder = new VariantBuilder();
+    VariantObjectBuilder object = builder.startObject();
+    object.appendKey("id");
+    // appendUUIDBytes must go through onAppend() so that numValues stays in sync with the
+    // appended keys. Otherwise endObject() throws because keys (1) != values (0).
+    object.appendUUIDBytes(ByteBuffer.wrap(uuid));
+    builder.endObject();
+
+    VariantTestUtil.testVariant(builder.build(), v -> {
+      VariantTestUtil.checkType(v, VariantUtil.OBJECT, Variant.Type.OBJECT);
+      Assert.assertEquals(1, v.numObjectElements());
+      VariantTestUtil.checkType(v.getFieldByKey("id"), VariantUtil.PRIMITIVE, Variant.Type.UUID);
+      Assert.assertEquals(expected, v.getFieldByKey("id").getUUID());
+    });
+  }
+
   @Test
   public void testEmptyObjectBuilder() {
     VariantBuilder b = new VariantBuilder();

parquet-variant/src/test/java/org/apache/parquet/variant/TestVariantScalarBuilder.java

@@ -512,4 +512,28 @@ public void testUUIDBuilder() {
       // expected
     }
   }
+
+  @Test
+  public void testUUIDBytesBuilder() {
+    byte[] uuid = new byte[] {0, 17, 34, 51, 68, 85, 102, 119, -120, -103, -86, -69, -52, -35, -18, -1};
+    long msb = ByteBuffer.wrap(uuid, 0, 8).order(ByteOrder.BIG_ENDIAN).getLong();
+    long lsb = ByteBuffer.wrap(uuid, 8, 8).order(ByteOrder.BIG_ENDIAN).getLong();
+    UUID expected = new UUID(msb, lsb);
+
+    VariantBuilder vb = new VariantBuilder();
+    vb.appendUUIDBytes(ByteBuffer.wrap(uuid));
+    VariantTestUtil.testVariant(vb.build(), v -> {
+      VariantTestUtil.checkType(v, VariantUtil.PRIMITIVE, Variant.Type.UUID);
+      Assert.assertEquals(expected, v.getUUID());
+    });
+
+    // appendUUIDBytes must go through onAppend(), so a second append on the root builder
+    // (which already holds a value) must be rejected instead of producing a multi-value buffer.
+    try {
+      vb.appendUUIDBytes(ByteBuffer.wrap(uuid));
+      Assert.fail("Expected Exception when appending multiple values");
+    } catch (Exception e) {
+      // expected
+    }
+  }
 }