GH-3223: Implement Variant parquet writer (#3221)

Author: cashmand

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,79 @@ 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();
+    }
+  }
+
+  // Return true if schema and avroSchema have the same field names, in the same order.
+  private static boolean schemaMatches(GroupType schema, Schema avroSchema) {
+    List<Schema.Field> avroFields = avroSchema.getFields();
+    if (schema.getFieldCount() != avroFields.size()) {
+      return false;
+    }
+
+    for (int i = 0; i < avroFields.size(); i += 1) {
+      if (!avroFields.get(i).name().equals(schema.getFieldName(i))) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+
+  private void writeVariantFields(GroupType schema, Schema avroSchema, Object record) {
+    List<Type> fields = schema.getFields();
+    List<Schema.Field> avroFields = avroSchema.getFields();
+
+    if (schemaMatches(schema, avroSchema)) {
+      // If the Avro schema matches the Parquet schema, the shredding matches and writeRecordFields can be used.
+      // writeRecordFields will validate that the field types match.
+      recordConsumer.startGroup();
+      writeRecordFields(schema, avroSchema, record);
+      recordConsumer.endGroup();
+      return;
+    }
+
+    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().equals("value")) {
+        Object valueObj = model.getField(record, avroField.name(), index);
+        Preconditions.checkArgument(
+            valueObj instanceof ByteBuffer,
+            "Expected ByteBuffer for value, but got " + valueObj.getClass());
+        value = (ByteBuffer) valueObj;
+      } else if (fieldType.getName().equals("metadata")) {
+        Object metadataObj = model.getField(record, avroField.name(), index);
+        Preconditions.checkArgument(
+            metadataObj instanceof ByteBuffer,
+            "Expected metadata to be a ByteBuffer, but got " + metadataObj.getClass());
+        metadata = (ByteBuffer) metadataObj;
+      } else {
+        binarySchema = false;
+        break;
+      }
+    }
+
+    if (binarySchema) {
+      VariantValueWriter.write(recordConsumer, schema, new Variant(value, metadata));
+    } else {
+      throw new RuntimeException("Invalid Avro schema for Variant logical type: " + schema.getName());
+    }
   }
 
   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. values in an object 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());
-    }
-  }
 }