GH-3223: Implement Variant parquet writer

parquet-avro/src/main/java/org/apache/parquet/avro/AvroWriteSupport.java

@@ -42,9 +42,12 @@
 import org.apache.parquet.io.api.Binary;
 import org.apache.parquet.io.api.RecordConsumer;
 import org.apache.parquet.schema.GroupType;
+import org.apache.parquet.schema.LogicalTypeAnnotation;
 import org.apache.parquet.schema.LogicalTypeAnnotation.UUIDLogicalTypeAnnotation;
 import org.apache.parquet.schema.MessageType;
 import org.apache.parquet.schema.Type;
+import org.apache.parquet.variant.Variant;
+import org.apache.parquet.variant.VariantValueWriter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -181,9 +184,58 @@ public void write(T record) {
   }
 
   private void writeRecord(GroupType schema, Schema avroSchema, Object record) {
-    recordConsumer.startGroup();
-    writeRecordFields(schema, avroSchema, record);
-    recordConsumer.endGroup();
+    if (schema.getLogicalTypeAnnotation() instanceof LogicalTypeAnnotation.VariantLogicalTypeAnnotation) {
+      writeVariantFields(schema, avroSchema, record);
+    } else {
+      recordConsumer.startGroup();
+      writeRecordFields(schema, avroSchema, record);
+      recordConsumer.endGroup();
+    }
+  }
+
+  private void writeVariantFields(GroupType schema, Schema avroSchema, Object record) {
+    List<Type> fields = schema.getFields();
+    List<Schema.Field> avroFields = avroSchema.getFields();
+    boolean binarySchema = true;
+    ByteBuffer metadata = null;
+    ByteBuffer value = null;
+    // Extract the value and metadata binary.
+    for (int index = 0; index < avroFields.size(); index++) {
+      Schema.Field avroField = avroFields.get(index);
+      Schema fieldSchema = AvroSchemaConverter.getNonNull(avroField.schema());
+      if (!fieldSchema.getType().equals(Schema.Type.BYTES)) {
+        binarySchema = false;
+        break;
+      }
+      Type fieldType = fields.get(index);
+      if (fieldType.getName() == "value") {
+        Object valueObj = model.getField(record, avroField.name(), index);
+        if (valueObj instanceof byte[]) {
+          value = ByteBuffer.wrap((byte[]) valueObj);
+        } else {
+          value = (ByteBuffer) valueObj;
+        }
+      } else if (fieldType.getName() == "metadata") {
+        Object metadataObj = model.getField(record, avroField.name(), index);
+        if (metadataObj instanceof byte[]) {
+          metadata = ByteBuffer.wrap((byte[]) metadataObj);
+        } else {
+          metadata = (ByteBuffer) metadataObj;
+        }
+      } else {
+        binarySchema = false;
+        break;
+      }
+    }
+
+    if (binarySchema) {
+      VariantValueWriter.write(recordConsumer, schema, new Variant(value, metadata));
+    } else {
+      // If the schema was something other than value and metaadata, treat the value as a non-variant record.
+      recordConsumer.startGroup();
+      writeRecordFields(schema, avroSchema, record);
+      recordConsumer.endGroup();
+    }
   }
 
   private void writeRecordFields(GroupType schema, Schema avroSchema, Object record) {

parquet-avro/src/test/java/org/apache/parquet/avro/AvroTestUtil.java

@@ -33,6 +33,7 @@
 import org.apache.parquet.hadoop.ParquetReader;
 import org.apache.parquet.hadoop.ParquetWriter;
 import org.apache.parquet.hadoop.util.HadoopInputFile;
+import org.apache.parquet.variant.Variant;
 import org.junit.Assert;
 import org.junit.rules.TemporaryFolder;
 
@@ -129,4 +130,36 @@ public static Configuration conf(String name, boolean value) {
     conf.setBoolean(name, value);
     return conf;
   }
+
+  /**
+   * Assert that to Variant values are logically equivalent.
+   * E.g. object fields may be ordered differently in the binary.
+   */
+  static void assertEquivalent(Variant expected, Variant actual) {
+    Assert.assertEquals(expected.getType(), actual.getType());
+    switch (expected.getType()) {
+      case STRING:
+        // Short strings may use the compact or extended representation.
+        Assert.assertEquals(expected.getString(), actual.getString());
+        break;
+      case ARRAY:
+        Assert.assertEquals(expected.numArrayElements(), actual.numArrayElements());
+        for (int i = 0; i < expected.numArrayElements(); ++i) {
+          assertEquivalent(expected.getElementAtIndex(i), actual.getElementAtIndex(i));
+        }
+        break;
+      case OBJECT:
+        Assert.assertEquals(expected.numObjectElements(), actual.numObjectElements());
+        for (int i = 0; i < expected.numObjectElements(); ++i) {
+          Variant.ObjectField expectedField = expected.getFieldAtIndex(i);
+          Variant.ObjectField actualField = actual.getFieldAtIndex(i);
+          Assert.assertEquals(expectedField.key, actualField.key);
+          assertEquivalent(expectedField.value, actualField.value);
+        }
+        break;
+      default:
+        // All other types have a single representation, and must be bit-for-bit identical.
+        Assert.assertEquals(expected.getValueBuffer(), actual.getValueBuffer());
+    }
+  }
 }

parquet-avro/src/test/java/org/apache/parquet/avro/TestReadVariant.java

@@ -2111,36 +2111,8 @@ void assertThrows(Callable callable, Class<? extends Exception> exception, Strin
   void assertEquivalent(ByteBuffer expectedMetadata, ByteBuffer expectedValue, GenericRecord actual) {
     assertEquals(expectedMetadata, (ByteBuffer) actual.get("metadata"));
     assertEquals(expectedMetadata, (ByteBuffer) actual.get("metadata"));
-    assertEquivalent(
+    AvroTestUtil.assertEquivalent(
         new Variant(expectedValue, expectedMetadata),
         new Variant(((ByteBuffer) actual.get("value")), expectedMetadata));
   }
-
-  void assertEquivalent(Variant expected, Variant actual) {
-    assertEquals(expected.getType(), actual.getType());
-    switch (expected.getType()) {
-      case STRING:
-        // Short strings may use the compact or extended representation.
-        assertEquals(expected.getString(), actual.getString());
-        break;
-      case ARRAY:
-        assertEquals(expected.numArrayElements(), actual.numArrayElements());
-        for (int i = 0; i < expected.numArrayElements(); ++i) {
-          assertEquivalent(expected.getElementAtIndex(i), actual.getElementAtIndex(i));
-        }
-        break;
-      case OBJECT:
-        assertEquals(expected.numObjectElements(), actual.numObjectElements());
-        for (int i = 0; i < expected.numObjectElements(); ++i) {
-          Variant.ObjectField expectedField = expected.getFieldAtIndex(i);
-          Variant.ObjectField actualField = actual.getFieldAtIndex(i);
-          assertEquals(expectedField.key, actualField.key);
-          assertEquivalent(expectedField.value, actualField.value);
-        }
-        break;
-      default:
-        // All other types have a single representation, and must be bit-for-bit identical.
-        assertEquals(expected.getValueBuffer(), actual.getValueBuffer());
-    }
-  }
 }