Push responsibility up to client code for deciding what is a valid Avro enum

Author: martin-traverse

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

@@ -24,6 +24,7 @@
 import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer;
 import org.apache.arrow.adapter.avro.producers.AvroBooleanProducer;
 import org.apache.arrow.adapter.avro.producers.AvroBytesProducer;
+import org.apache.arrow.adapter.avro.producers.AvroEnumProducer;
 import org.apache.arrow.adapter.avro.producers.AvroFixedSizeBinaryProducer;
 import org.apache.arrow.adapter.avro.producers.AvroFixedSizeListProducer;
 import org.apache.arrow.adapter.avro.producers.AvroFloat2Producer;
@@ -62,6 +63,7 @@
 import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampSecProducer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampSecTzProducer;
 import org.apache.arrow.util.Preconditions;
+import org.apache.arrow.vector.BaseIntVector;
 import org.apache.arrow.vector.BigIntVector;
 import org.apache.arrow.vector.BitVector;
 import org.apache.arrow.vector.DateDayVector;
@@ -100,7 +102,6 @@
 import org.apache.arrow.vector.complex.MapVector;
 import org.apache.arrow.vector.complex.StructVector;
 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.FloatingPointPrecision;
 import org.apache.arrow.vector.types.TimeUnit;
@@ -331,12 +332,8 @@ private static <T> T buildBaseTypeSchema(
             String[] symbols = dictionarySymbols(dictionary);
             return builder.enumeration(field.getName()).symbols(symbols);
           } else {
-            Field decodedField =
-                new Field(
-                    field.getName(),
-                    dictionary.getVector().getField().getFieldType(),
-                    dictionary.getVector().getField().getChildren());
-            return buildBaseTypeSchema(builder, decodedField, namespace, dictionaries);
+            throw new IllegalArgumentException(
+                "Dictionary-encoded field is not a valid enum: " + field.getName());
           }
         }
 
@@ -449,7 +446,8 @@ private static LogicalType timestampLogicalType(ArrowType.Timestamp timestampTyp
 
   private static boolean dictionaryIsValidEnum(Dictionary dictionary) {
 
-    if (dictionary.getVectorType().getTypeID() != ArrowType.ArrowTypeID.Utf8) {
+    if (dictionary.getVectorType().getTypeID() != ArrowType.ArrowTypeID.Utf8
+        || !(dictionary.getVector() instanceof VarCharVector)) {
       return false;
     }
 
@@ -502,33 +500,24 @@ private static String[] dictionarySymbols(Dictionary dictionary) {
    * @param vectors The vectors that will be used to produce Avro data
    * @return The resulting composite Avro producer
    */
-  public static CompositeAvroProducer createCompositeProducer(
-      List<FieldVector> vectors, DictionaryProvider dictionaries) {
+  public static CompositeAvroProducer createCompositeProducer(List<FieldVector> vectors) {
 
     List<Producer<? extends FieldVector>> producers = new ArrayList<>(vectors.size());
 
     for (FieldVector vector : vectors) {
-      BaseAvroProducer<? extends FieldVector> producer = createProducer(vector, dictionaries);
+      BaseAvroProducer<? extends FieldVector> producer = createProducer(vector);
       producers.add(producer);
     }
 
     return new CompositeAvroProducer(producers);
   }
 
-  /** Overload provided for convenience, sets dictionaries = null. */
-  public static CompositeAvroProducer createCompositeProducer(List<FieldVector> vectors) {
-
-    return createCompositeProducer(vectors, null);
-  }
-
-  private static BaseAvroProducer<?> createProducer(
-      FieldVector vector, DictionaryProvider dictionaries) {
+  private static BaseAvroProducer<?> createProducer(FieldVector vector) {
     boolean nullable = vector.getField().isNullable();
-    return createProducer(vector, nullable, dictionaries);
+    return createProducer(vector, nullable);
   }
 
-  private static BaseAvroProducer<?> createProducer(
-      FieldVector vector, boolean nullable, DictionaryProvider dictionaries) {
+  private static BaseAvroProducer<?> createProducer(FieldVector vector, boolean nullable) {
 
     Preconditions.checkNotNull(vector, "Arrow vector object can't be null");
 
@@ -537,30 +526,13 @@ private static BaseAvroProducer<?> createProducer(
     // Avro understands nullable types as a union of type | null
     // Most nullable fields in a VSR will not be unions, so provide a special wrapper
     if (nullable && minorType != Types.MinorType.UNION) {
-      final BaseAvroProducer<?> innerProducer = createProducer(vector, false, dictionaries);
+      final BaseAvroProducer<?> innerProducer = createProducer(vector, false);
       return new AvroNullableProducer<>(innerProducer);
     }
 
     if (vector.getField().getDictionary() != null) {
-      if (dictionaries == null) {
-        throw new IllegalArgumentException(
-            "Field references a dictionary but no dictionaries were provided: "
-                + vector.getField().getName());
-      }
-      Dictionary dictionary = dictionaries.lookup(vector.getField().getDictionary().getId());
-      if (dictionary == null) {
-        throw new IllegalArgumentException(
-            "Field references a dictionary that does not exist: "
-                + vector.getField().getName()
-                + ", dictionary ID = "
-                + vector.getField().getDictionary().getId());
-      }
-      // If a field is dictionary-encoded but cannot be represented as an Avro enum,
-      // then decode it before writing
-      if (!dictionaryIsValidEnum(dictionary)) {
-        FieldVector decodedVector = (FieldVector) DictionaryEncoder.decode(vector, dictionary);
-        return createProducer(decodedVector, nullable, dictionaries);
-      }
+      BaseIntVector dictEncodedVector = (BaseIntVector) vector;
+      return new AvroEnumProducer(dictEncodedVector);
     }
 
     switch (minorType) {
@@ -640,23 +612,21 @@ private static BaseAvroProducer<?> createProducer(
         Producer<?>[] childProducers = new Producer<?>[childVectors.size()];
         for (int i = 0; i < childVectors.size(); i++) {
           FieldVector childVector = childVectors.get(i);
-          childProducers[i] =
-              createProducer(childVector, childVector.getField().isNullable(), dictionaries);
+          childProducers[i] = createProducer(childVector, childVector.getField().isNullable());
         }
         return new AvroStructProducer(structVector, childProducers);
 
       case LIST:
         ListVector listVector = (ListVector) vector;
         FieldVector itemVector = listVector.getDataVector();
-        Producer<?> itemProducer =
-            createProducer(itemVector, itemVector.getField().isNullable(), dictionaries);
+        Producer<?> itemProducer = createProducer(itemVector, itemVector.getField().isNullable());
         return new AvroListProducer(listVector, itemProducer);
 
       case FIXED_SIZE_LIST:
         FixedSizeListVector fixedListVector = (FixedSizeListVector) vector;
         FieldVector fixedItemVector = fixedListVector.getDataVector();
         Producer<?> fixedItemProducer =
-            createProducer(fixedItemVector, fixedItemVector.getField().isNullable(), dictionaries);
+            createProducer(fixedItemVector, fixedItemVector.getField().isNullable());
         return new AvroFixedSizeListProducer(fixedListVector, fixedItemProducer);
 
       case MAP:
@@ -670,7 +640,7 @@ private static BaseAvroProducer<?> createProducer(
         FieldVector valueVector = entryVector.getChildrenFromFields().get(1);
         Producer<?> keyProducer = new AvroStringProducer(keyVector);
         Producer<?> valueProducer =
-            createProducer(valueVector, valueVector.getField().isNullable(), dictionaries);
+            createProducer(valueVector, valueVector.getField().isNullable());
         Producer<?> entryProducer =
             new AvroStructProducer(entryVector, new Producer<?>[] {keyProducer, valueProducer});
         return new AvroMapProducer(mapVector, entryProducer);

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.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/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java

@@ -2872,8 +2872,7 @@ public void testWriteDictEnumEncoded() throws Exception {
       // 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);
+        CompositeAvroProducer producer = ArrowToAvroUtils.createCompositeProducer(vectors);
         for (int row = 0; row < rowCount; row++) {
           producer.produce(encoder);
         }
@@ -2898,107 +2897,4 @@ record = datumReader.read(record, decoder);
       }
     }
   }
-
-  @Test
-  public void testWriteEnumDecoded() throws Exception {
-
-    // Dict encoded fields that are not valid Avro enums should be decoded on write
-
-    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()); // spaced not allowed
-    dictionaryVector.set(1, "banana".getBytes());
-    dictionaryVector.set(2, "cherry".getBytes());
-    dictionaryVector.setValueCount(3);
-
-    Dictionary dictionary =
-        new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null));
-
-    FieldType dictionaryField2 = new FieldType(false, new ArrowType.Int(64, true), null);
-    BigIntVector dictionaryVector2 =
-        new BigIntVector(new Field("dictionary2", dictionaryField2, null), allocator);
-
-    dictionaryVector2.allocateNew(3);
-    dictionaryVector2.set(0, 0L);
-    dictionaryVector2.set(1, 1L);
-    dictionaryVector2.set(2, 2L);
-    dictionaryVector2.setValueCount(3);
-
-    Dictionary dictionary2 =
-        new Dictionary(dictionaryVector2, new DictionaryEncoding(2L, false, null));
-
-    DictionaryProvider dictionaries =
-        new DictionaryProvider.MapDictionaryProvider(dictionary, dictionary2);
-
-    // 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, "passion fruit".getBytes());
-    stringVector.setSafe(1, "banana".getBytes());
-    stringVector.setSafe(2, "cherry".getBytes());
-    stringVector.setSafe(3, "cherry".getBytes());
-    stringVector.setSafe(4, "passion fruit".getBytes());
-    stringVector.setSafe(5, "banana".getBytes());
-    stringVector.setSafe(6, "passion fruit".getBytes());
-    stringVector.setSafe(7, "cherry".getBytes());
-    stringVector.setSafe(8, "banana".getBytes());
-    stringVector.setSafe(9, "passion fruit".getBytes());
-    stringVector.setValueCount(10);
-
-    FieldType longField = new FieldType(false, new ArrowType.Int(64, true), null);
-    BigIntVector longVector = new BigIntVector(new Field("enumField2", longField, null), allocator);
-    longVector.allocateNew(10);
-    for (int i = 0; i < 10; i++) {
-      longVector.setSafe(i, (long) i % 3);
-    }
-    longVector.setValueCount(10);
-
-    IntVector encodedVector = (IntVector) DictionaryEncoder.encode(stringVector, dictionary);
-    IntVector encodedVector2 = (IntVector) DictionaryEncoder.encode(longVector, dictionary2);
-
-    // Set up VSR
-    List<FieldVector> vectors = Arrays.asList(encodedVector, encodedVector2);
-    int rowCount = 10;
-
-    try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
-
-      File dataFile = new File(TMP, "testWriteEnumDecodedavro");
-
-      // 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);
-          assertEquals(stringVector.getObject(row).toString(), record.get("enumField").toString());
-          assertEquals(longVector.getObject(row), record.get("enumField2"));
-        }
-      }
-    }
-  }
 }

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

@@ -17,11 +17,13 @@
 package org.apache.arrow.adapter.avro;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 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;
@@ -1437,4 +1439,70 @@ public void testWriteDictEnumEncoded() {
     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));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries));
+  }
+
+  @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));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries));
+  }
 }

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

@@ -120,7 +120,7 @@ private static void roundTripTest(
     try (FileOutputStream fos = new FileOutputStream(dataFile)) {
       BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null);
       CompositeAvroProducer producer =
-          ArrowToAvroUtils.createCompositeProducer(root.getFieldVectors(), dictionaries);
+          ArrowToAvroUtils.createCompositeProducer(root.getFieldVectors());
       for (int row = 0; row < rowCount; row++) {
         producer.produce(encoder);
       }