GH-731: Avro adapter, output dictionary-encoded fields as enums (#779)

## What's Changed

Updated ArrowToAvro to output dictionary-encoded string vectors as Avro
enums, where possible. Apologies for the delay - busy as usual!

To output dict encoded vectors as enums, a dictionary provider must be
supplied to the top level methods with all the required dictionaries.
All dictionary values must be present when the schema is written, i.e.
before the data blocks are produced. If data is being written as a
schema followed by multiple blocks, values added to a dictionary in
between blocks will not be included in the schema resulting in an
invalid Avro file (in general supply an invalid dictionary mapping will
result in invalid output).

Dictionary encoded fields are checked to ensure they are valid Avro
enums. If the dictionary encoded field is not a string field, or the
string values are not valid Avro enums, the field is decoded and output
as literal values. This is done by calling
DictionaryEncoder.decode(vector, dictionary), which will consume memory
for the vector. An alternative approach would be to decode values
one-by-one, however this would require a significant change to the
producer pattern since the current producers expect concrete vectors of
the output type. Another option would be to throw an error if there are
dictionary-encoded vectors that are not string types, i.e. push the
responsibility onto client code. I'm not sure which approach is best -
happy to take any guidance and I will update the code accordingly.

To read enums back the current approach for decoding is unchanged (the
AvroToArrow config has to be set up with a MapDictionaryProvider which
is populated when data is read). The last part of the Avro work is to
add the capability for reading / writing whole files block-by-block, so
there is an opportunity to do something with the top level APIs there,
for now the current API works and I've used it in the round trip tests.

Please let me know any feedback, happy to update as needed!

Closes #731.

Author: martin-traverse

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java

@@ -17,22 +17,22 @@
 package org.apache.arrow.adapter.avro.producers;
 
 import java.io.IOException;
-import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.BaseIntVector;
 import org.apache.avro.io.Encoder;
 
 /**
- * Producer that produces enum values from a dictionary-encoded {@link IntVector}, writes data to an
- * Avro encoder.
+ * Producer that produces enum values from a dictionary-encoded {@link BaseIntVector}, writes data
+ * to an Avro encoder.
  */
-public class AvroEnumProducer extends BaseAvroProducer<IntVector> {
+public class AvroEnumProducer extends BaseAvroProducer<BaseIntVector> {
 
   /** Instantiate an AvroEnumProducer. */
-  public AvroEnumProducer(IntVector vector) {
+  public AvroEnumProducer(BaseIntVector vector) {
     super(vector);
   }
 
   @Override
   public void produce(Encoder encoder) throws IOException {
-    encoder.writeEnum(vector.get(currentIndex++));
+    encoder.writeEnum((int) vector.getValueAsLong(currentIndex++));
   }
 }

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java

@@ -0,0 +1,47 @@
+/*
+ * 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.arrow.adapter.avro.producers;
+
+import java.io.IOException;
+import org.apache.arrow.vector.BaseIntVector;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.avro.io.Encoder;
+
+/**
+ * Producer that decodes values from a dictionary-encoded {@link FieldVector}, writes the resulting
+ * values to an Avro encoder.
+ *
+ * @param <T> Type of the underlying dictionary vector
+ */
+public class DictionaryDecodingProducer<T extends FieldVector>
+    extends BaseAvroProducer<BaseIntVector> {
+
+  private final Producer<T> dictProducer;
+
+  /** Instantiate a DictionaryDecodingProducer. */
+  public DictionaryDecodingProducer(BaseIntVector indexVector, Producer<T> dictProducer) {
+    super(indexVector);
+    this.dictProducer = dictProducer;
+  }
+
+  @Override
+  public void produce(Encoder encoder) throws IOException {
+    int dicIndex = (int) vector.getValueAsLong(currentIndex++);
+    dictProducer.setPosition(dicIndex);
+    dictProducer.produce(encoder);
+  }
+}

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java

@@ -76,10 +76,14 @@
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.complex.writer.BaseWriter;
 import org.apache.arrow.vector.complex.writer.FieldWriter;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryEncoder;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.types.DateUnit;
 import org.apache.arrow.vector.types.FloatingPointPrecision;
 import org.apache.arrow.vector.types.TimeUnit;
 import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.util.JsonStringArrayList;
@@ -2817,4 +2821,81 @@ record = datumReader.read(record, decoder);
       }
     }
   }
+
+  @Test
+  public void testWriteDictEnumEncoded() throws Exception {
+
+    BufferAllocator allocator = new RootAllocator();
+
+    // Create a dictionary
+    FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+    VarCharVector dictionaryVector =
+        new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+    dictionaryVector.allocateNew(3);
+    dictionaryVector.set(0, "apple".getBytes());
+    dictionaryVector.set(1, "banana".getBytes());
+    dictionaryVector.set(2, "cherry".getBytes());
+    dictionaryVector.setValueCount(3);
+
+    Dictionary dictionary =
+        new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null));
+    DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+    // Field definition
+    FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null);
+    VarCharVector stringVector =
+        new VarCharVector(new Field("enumField", stringField, null), allocator);
+    stringVector.allocateNew(10);
+    stringVector.setSafe(0, "apple".getBytes());
+    stringVector.setSafe(1, "banana".getBytes());
+    stringVector.setSafe(2, "cherry".getBytes());
+    stringVector.setSafe(3, "cherry".getBytes());
+    stringVector.setSafe(4, "apple".getBytes());
+    stringVector.setSafe(5, "banana".getBytes());
+    stringVector.setSafe(6, "apple".getBytes());
+    stringVector.setSafe(7, "cherry".getBytes());
+    stringVector.setSafe(8, "banana".getBytes());
+    stringVector.setSafe(9, "apple".getBytes());
+    stringVector.setValueCount(10);
+
+    IntVector encodedVector = (IntVector) DictionaryEncoder.encode(stringVector, dictionary);
+
+    // Set up VSR
+    List<FieldVector> vectors = Arrays.asList(encodedVector);
+    int rowCount = 10;
+
+    try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
+
+      File dataFile = new File(TMP, "testWriteEnumEncoded.avro");
+
+      // Write an AVRO block using the producer classes
+      try (FileOutputStream fos = new FileOutputStream(dataFile)) {
+        BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null);
+        CompositeAvroProducer producer =
+            ArrowToAvroUtils.createCompositeProducer(vectors, dictionaries);
+        for (int row = 0; row < rowCount; row++) {
+          producer.produce(encoder);
+        }
+        encoder.flush();
+      }
+
+      // Set up reading the AVRO block as a GenericRecord
+      Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields(), dictionaries);
+      GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>(schema);
+
+      try (InputStream inputStream = new FileInputStream(dataFile)) {
+
+        BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null);
+        GenericRecord record = null;
+
+        // Read and check values
+        for (int row = 0; row < rowCount; row++) {
+          record = datumReader.read(record, decoder);
+          // Values read from Avro should be the decoded enum values
+          assertEquals(stringVector.getObject(row).toString(), record.get("enumField").toString());
+        }
+      }
+    }
+  }
 }

adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java

@@ -20,11 +20,18 @@
 
 import java.util.Arrays;
 import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.types.DateUnit;
 import org.apache.arrow.vector.types.FloatingPointPrecision;
 import org.apache.arrow.vector.types.TimeUnit;
 import org.apache.arrow.vector.types.UnionMode;
 import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.avro.LogicalTypes;
@@ -1389,4 +1396,126 @@ public void testConvertUnionTypes() {
         Schema.Type.STRING,
         schema.getField("nullableDenseUnionField").schema().getTypes().get(3).getType());
   }
+
+  @Test
+  public void testWriteDictEnumEncoded() {
+
+    BufferAllocator allocator = new RootAllocator();
+
+    // Create a dictionary
+    FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+    VarCharVector dictionaryVector =
+        new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+    dictionaryVector.allocateNew(3);
+    dictionaryVector.set(0, "apple".getBytes());
+    dictionaryVector.set(1, "banana".getBytes());
+    dictionaryVector.set(2, "cherry".getBytes());
+    dictionaryVector.setValueCount(3);
+
+    Dictionary dictionary =
+        new Dictionary(
+            dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+    DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+    List<Field> fields =
+        Arrays.asList(
+            new Field(
+                "enumField",
+                new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null),
+                null));
+
+    Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
+
+    assertEquals(Schema.Type.RECORD, schema.getType());
+    assertEquals(1, schema.getFields().size());
+
+    Schema.Field enumField = schema.getField("enumField");
+
+    assertEquals(Schema.Type.ENUM, enumField.schema().getType());
+    assertEquals(3, enumField.schema().getEnumSymbols().size());
+    assertEquals("apple", enumField.schema().getEnumSymbols().get(0));
+    assertEquals("banana", enumField.schema().getEnumSymbols().get(1));
+    assertEquals("cherry", enumField.schema().getEnumSymbols().get(2));
+  }
+
+  @Test
+  public void testWriteDictEnumInvalid() {
+
+    BufferAllocator allocator = new RootAllocator();
+
+    // Create a dictionary
+    FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+    VarCharVector dictionaryVector =
+        new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+    dictionaryVector.allocateNew(3);
+    dictionaryVector.set(0, "passion fruit".getBytes());
+    dictionaryVector.set(1, "banana".getBytes());
+    dictionaryVector.set(2, "cherry".getBytes());
+    dictionaryVector.setValueCount(3);
+
+    Dictionary dictionary =
+        new Dictionary(
+            dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+    DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+    List<Field> fields =
+        Arrays.asList(
+            new Field(
+                "enumField",
+                new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null),
+                null));
+
+    // Dictionary field contains values that are not valid enums
+    // Should be decoded and output as a string field
+
+    Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
+
+    assertEquals(Schema.Type.RECORD, schema.getType());
+    assertEquals(1, schema.getFields().size());
+
+    Schema.Field enumField = schema.getField("enumField");
+    assertEquals(Schema.Type.STRING, enumField.schema().getType());
+  }
+
+  @Test
+  public void testWriteDictEnumInvalid2() {
+
+    BufferAllocator allocator = new RootAllocator();
+
+    // Create a dictionary
+    FieldType dictionaryField = new FieldType(false, new ArrowType.Int(64, true), null);
+    BigIntVector dictionaryVector =
+        new BigIntVector(new Field("dictionary", dictionaryField, null), allocator);
+
+    dictionaryVector.allocateNew(3);
+    dictionaryVector.set(0, 123L);
+    dictionaryVector.set(1, 456L);
+    dictionaryVector.set(2, 789L);
+    dictionaryVector.setValueCount(3);
+
+    Dictionary dictionary =
+        new Dictionary(
+            dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+    DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+    List<Field> fields =
+        Arrays.asList(
+            new Field(
+                "enumField",
+                new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null),
+                null));
+
+    // Dictionary field encodes LONG values rather than STRING
+    // Should be doecded and output as a LONG field
+
+    Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
+
+    assertEquals(Schema.Type.RECORD, schema.getType());
+    assertEquals(1, schema.getFields().size());
+
+    Schema.Field enumField = schema.getField("enumField");
+    assertEquals(Schema.Type.LONG, enumField.schema().getType());
+  }
 }