Variant: Refactor readers to use a parent handler.

Author: rdblue

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

@@ -32,7 +32,6 @@
 import org.apache.parquet.schema.GroupType;
 import org.apache.parquet.schema.PrimitiveStringifier;
 import org.apache.parquet.schema.PrimitiveType;
-import org.apache.parquet.variant.VariantColumnConverter;
 
 public class AvroConverters {
 
@@ -341,27 +340,4 @@ public String convert(Binary binary) {
       return stringifier.stringify(binary);
     }
   }
-
-  static final class FieldVariantConverter<T> extends VariantColumnConverter {
-    private final ParentValueContainer parent;
-    private final Schema avroSchema;
-    private final GenericData model;
-
-    public FieldVariantConverter(
-        ParentValueContainer parent, GroupType schema, Schema avroSchema, GenericData model) {
-      super(schema);
-      this.avroSchema = avroSchema;
-      this.model = model;
-      this.parent = parent;
-    }
-
-    @Override
-    @SuppressWarnings("unchecked")
-    public void addVariant(ByteBuffer value, ByteBuffer metadata) {
-      T currentRecord = (T) model.newRecord(null, avroSchema);
-      model.setField(currentRecord, "metadata", 0, metadata);
-      model.setField(currentRecord, "value", 1, value);
-      parent.add(currentRecord);
-    }
-  }
 }

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

@@ -170,7 +170,7 @@ private static Converter newConverter(Schema schema, Type type, GenericData mode
         return new MapConverter(parent, type.asGroupType(), schema, model);
       case RECORD:
         if (type.getLogicalTypeAnnotation() instanceof LogicalTypeAnnotation.VariantLogicalTypeAnnotation) {
-          return new AvroConverters.FieldVariantConverter(parent, type.asGroupType(), schema, model);
+          return new AvroVariantConverter(parent, type.asGroupType(), schema, model);
         } else {
           return new AvroIndexedRecordConverter<>(parent, type.asGroupType(), schema, model);
         }

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

@@ -396,7 +396,7 @@ private static Converter newConverter(
         return newStringConverter(schema, model, parent, validator);
       case RECORD:
         if (type.getLogicalTypeAnnotation() instanceof LogicalTypeAnnotation.VariantLogicalTypeAnnotation) {
-          return new AvroConverters.FieldVariantConverter(parent, type.asGroupType(), schema, model);
+          return new AvroVariantConverter(parent, type.asGroupType(), schema, model);
         } else {
           return new AvroRecordConverter(parent, type.asGroupType(), schema, model, validator);
         }
@@ -1004,7 +1004,7 @@ static boolean isElementType(Type repeatedType, Schema elementSchema) {
     } else if (elementSchema != null && elementSchema.getType() == Schema.Type.RECORD) {
       Schema schemaFromRepeated = CONVERTER.convert(repeatedType.asGroupType());
       if (checkReaderWriterCompatibility(elementSchema, schemaFromRepeated)
-              .getType()
+          .getType()
           == COMPATIBLE) {
         return true;
       }

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

@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.parquet.avro;
+
+import java.nio.ByteBuffer;
+import java.util.function.Consumer;
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
+import org.apache.parquet.Preconditions;
+import org.apache.parquet.io.api.Converter;
+import org.apache.parquet.io.api.GroupConverter;
+import org.apache.parquet.schema.GroupType;
+import org.apache.parquet.variant.ImmutableMetadata;
+import org.apache.parquet.variant.VariantBuilder;
+import org.apache.parquet.variant.VariantConverters;
+
+/**
+ * Converter for Variant values.
+ */
+class AvroVariantConverter extends GroupConverter implements VariantConverters.ParentConverter<VariantBuilder> {
+  private final ParentValueContainer parent;
+  private final Schema avroSchema;
+  private final GenericData model;
+  private final int metadataPos;
+  private final int valuePos;
+  private final GroupConverter wrappedConverter;
+
+  private VariantBuilder builder = null;
+  private ImmutableMetadata metadata = null;
+
+  AvroVariantConverter(ParentValueContainer parent, GroupType variantGroup, Schema avroSchema, GenericData model) {
+    this.parent = parent;
+    this.avroSchema = avroSchema;
+    this.metadataPos = avroSchema.getField("metadata").pos();
+    this.valuePos = avroSchema.getField("value").pos();
+    this.model = model;
+    this.wrappedConverter = VariantConverters.newVariantConverter(variantGroup, this::setMetadata, this);
+  }
+
+  @Override
+  public void build(Consumer<VariantBuilder> consumer) {
+    Preconditions.checkState(builder != null, "Cannot build variant: builder has not been initialized");
+    consumer.accept(builder);
+  }
+
+  @Override
+  public Converter getConverter(int fieldIndex) {
+    return wrappedConverter.getConverter(fieldIndex);
+  }
+
+  @Override
+  public void start() {
+    wrappedConverter.start();
+  }
+
+  @Override
+  public void end() {
+    wrappedConverter.end();
+
+    Preconditions.checkState(metadata != null, "Cannot build variant: missing metadata");
+
+    builder.appendNullIfEmpty();
+
+    Object record = model.newRecord(null, avroSchema);
+    model.setField(record, "metadata", metadataPos, metadata.getEncodedBuffer());
+    model.setField(record, "value", valuePos, builder.encodedValue());
+    parent.add(record);
+
+    this.builder = null;
+  }
+
+  void setMetadata(ByteBuffer metadataBuffer) {
+    // If the metadata hasn't changed, we don't need to rebuild the map.
+    if (metadata == null || metadata.getEncodedBuffer() != metadataBuffer) {
+      this.metadata = new ImmutableMetadata(metadataBuffer);
+    }
+
+    this.builder = new VariantBuilder(metadata);
+  }
+}