Use dictionary decoding producer to handle dictionary fields that are not valid avro enums

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;
@@ -44,6 +45,7 @@
 import org.apache.arrow.adapter.avro.producers.AvroUint8Producer;
 import org.apache.arrow.adapter.avro.producers.BaseAvroProducer;
 import org.apache.arrow.adapter.avro.producers.CompositeAvroProducer;
+import org.apache.arrow.adapter.avro.producers.DictionaryDecodingProducer;
 import org.apache.arrow.adapter.avro.producers.Producer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroDateDayProducer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroDateMilliProducer;
@@ -62,6 +64,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 +103,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;
@@ -557,9 +559,11 @@ private static BaseAvroProducer<?> createProducer(
       }
       // 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);
+      if (dictionaryIsValidEnum(dictionary)) {
+        return new AvroEnumProducer((BaseIntVector) vector);
+      } else {
+        BaseAvroProducer<?> dictProducer = createProducer(dictionary.getVector(), false, null);
+        return new DictionaryDecodingProducer<>((BaseIntVector) vector, dictProducer);
       }
     }
 

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
+ * 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/DictionaryDecodingProducer.java

@@ -0,0 +1,45 @@
+/*
+ * 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 produces decoded values from a dictionary-encoded {@link BaseIntVector}, writes data to an
+ * Avro encoder.
+ */
+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);
+  }
+}