Column projection of union type

core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java

@@ -28,6 +28,8 @@
 import org.apache.iceberg.types.Types;
 
 public abstract class AvroSchemaWithTypeVisitor<T> {
+  private static final String UNION_TAG_FIELD_NAME = "tag";
+
   public static <T> T visit(org.apache.iceberg.Schema iSchema, Schema schema, AvroSchemaWithTypeVisitor<T> visitor) {
     return visit(iSchema.asStruct(), schema, visitor);
   }
@@ -97,17 +99,65 @@ private static <T> T visitUnion(Type type, Schema union, AvroSchemaWithTypeVisit
         options.add(visit(type, branch, visitor));
       }
     } else { // complex union case
-      int index = 1;
-      for (Schema branch : types) {
-        if (branch.getType() == Schema.Type.NULL) {
-          options.add(visit((Type) null, branch, visitor));
-        } else {
-          options.add(visit(type.asStructType().fields().get(index).type(), branch, visitor));
-          index += 1;
+      visitComplexUnion(type, union, visitor, options);
+    }
+    return visitor.union(type, union, options);
+  }
+
+  /*
+  A complex union with multiple types of Avro schema is converted into a struct with multiple fields of Iceberg schema.
+  A field is related to a type in the order defined in Avro schema. Also, an extra tag field is added into the struct of
+  Iceberg schema. The user can query the column of union type with the field projected (e.g. colUnion.field0) in which
+  the maximum number of the fields to be projected equals to the number of fields of the complete struct converted
+  from the union. The case of without field projection equals to the case of full fields projection.
+  Therefore, this function visits the complex union by assuming the field projection always happens.
+   */
+  private static <T> void visitComplexUnion(Type type, Schema union,
+                                            AvroSchemaWithTypeVisitor<T> visitor, List<T> options) {
+    boolean nullTypeFound = false;
+    int typeIndex = 0;
+    int fieldIndexInStruct = 0;
+    while (typeIndex < union.getTypes().size()) {
+      Schema schema = union.getTypes().get(typeIndex);
+      // in some cases, a NULL type exists in the union of Avro schema besides the actual types,
+      // and it affects the index of the actual types of the order in the union
+      if (schema.getType() == Schema.Type.NULL) {
+        nullTypeFound = true;
+        options.add(visit((Type) null, schema, visitor));
+        typeIndex++;
+        continue;
+      }
+
+      // If a NULL type is found before current type, the type index is one larger than the actual type index which
+      // can be used to track the corresponding field in the struct of Iceberg schema.
+      int actualTypeIndex = nullTypeFound ? typeIndex - 1 : typeIndex;
+      boolean relatedFieldInStructFound = false;
+      while (fieldIndexInStruct < type.asStructType().fields().size()) {
+        String structFieldName = type.asStructType().fields().get(fieldIndexInStruct).name();
+        if (UNION_TAG_FIELD_NAME.equals(structFieldName)) {
+          fieldIndexInStruct++;
+          continue;
+        }
+
+        int indexFromStructFieldName = Integer.valueOf(structFieldName.substring(5));
+        if (actualTypeIndex == indexFromStructFieldName) {
+          relatedFieldInStructFound = true;
+          options.add(visit(type.asStructType().fields().get(fieldIndexInStruct).type(), schema, visitor));
+          fieldIndexInStruct++;
         }
+        break;
+      }
+
+      // If a field is not projected, a corresponding field in the struct of Iceberg schema cannot be found
+      // for current type of union in Avro schema, a reader for current type still needs to be created and
+      // used to make the reading of Avro file successfully. In this case, a null field type is used to
+      // create the option for the reader of the current type which still can read the corresponding content
+      // in Avro file successfully.
+      if (!relatedFieldInStructFound) {
+        options.add(visit((Type) null, schema, visitor));
       }
+      typeIndex++;
     }
-    return visitor.union(type, union, options);
   }
 
   private static <T> T visitArray(Type type, Schema array, AvroSchemaWithTypeVisitor<T> visitor) {

spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java

@@ -83,7 +83,7 @@ public ValueReader<?> union(Type expected, Schema union, List<ValueReader<?>> op
       if (AvroSchemaUtil.isOptionSchema(union) || AvroSchemaUtil.isSingleTypeUnion(union)) {
         return ValueReaders.union(options);
       } else {
-        return SparkValueReaders.union(union, options);
+        return SparkValueReaders.union(union, options, expected);
       }
     }
 

spark/src/main/java/org/apache/iceberg/spark/data/SparkValueReaders.java

@@ -25,6 +25,7 @@
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -34,6 +35,7 @@
 import org.apache.iceberg.avro.ValueReader;
 import org.apache.iceberg.avro.ValueReaders;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.UUIDUtil;
 import org.apache.spark.sql.catalyst.InternalRow;
@@ -86,8 +88,8 @@ static ValueReader<InternalRow> struct(List<ValueReader<?>> readers, Types.Struc
     return new StructReader(readers, struct, idToConstant);
   }
 
-  static ValueReader<InternalRow> union(Schema schema, List<ValueReader<?>> readers) {
-    return new UnionReader(schema, readers);
+  static ValueReader<InternalRow> union(Schema schema, List<ValueReader<?>> readers, Type expected) {
+    return new UnionReader(schema, readers, expected);
   }
 
   private static class StringReader implements ValueReader<UTF8String> {
@@ -302,54 +304,75 @@ protected void set(InternalRow struct, int pos, Object value) {
   }
 
   private static class UnionReader implements ValueReader<InternalRow> {
+    private static final String UNION_TAG_FIELD_NAME = "tag";
     private final Schema schema;
     private final ValueReader[] readers;
+    private final int[] projectedFieldIdsToIdxInReturnedRow;
+    private boolean isTagFieldProjected;
+    private int numOfFieldsInReturnedRow;
+    private int nullTypeIndex;
 
-    private UnionReader(Schema schema, List<ValueReader<?>> readers) {
+    private UnionReader(Schema schema, List<ValueReader<?>> readers, Type expected) {
       this.schema = schema;
       this.readers = new ValueReader[readers.size()];
       for (int i = 0; i < this.readers.length; i += 1) {
         this.readers[i] = readers.get(i);
       }
-    }
 
-    @Override
-    public InternalRow read(Decoder decoder, Object reuse) throws IOException {
-      // first we need to filter out NULL alternative if it exists in the union schema
-      int nullIndex = -1;
-      List<Schema> alts = schema.getTypes();
-      for (int i = 0; i < alts.size(); i++) {
-        Schema alt = alts.get(i);
+      // checking if NULL type exists in Avro union schema
+      this.nullTypeIndex = -1;
+      for (int i = 0; i < this.schema.getTypes().size(); i++) {
+        Schema alt = this.schema.getTypes().get(i);
         if (Objects.equals(alt.getType(), Schema.Type.NULL)) {
-          nullIndex = i;
+          this.nullTypeIndex = i;
           break;
         }
       }
 
+      // Creating an integer array to track the mapping between the index of fields to be projected
+      // and the index of the value for the field stored in the returned row,
+      // if the value for a field equals to -1, it means the value of this field should not be stored
+      // in the returned row
+      int numberOfTypes = this.nullTypeIndex == -1 ?
+          this.schema.getTypes().size() : this.schema.getTypes().size() - 1;
+      this.projectedFieldIdsToIdxInReturnedRow = new int[numberOfTypes];
+      Arrays.fill(this.projectedFieldIdsToIdxInReturnedRow, -1);
+      this.numOfFieldsInReturnedRow = 0;
+      this.isTagFieldProjected = false;
+      for (Types.NestedField expectedStructField : expected.asStructType().fields()) {
+        String fieldName = expectedStructField.name();
+        if (fieldName.equals(UNION_TAG_FIELD_NAME)) {
+          this.isTagFieldProjected = true;
+          this.numOfFieldsInReturnedRow++;
+          continue;
+        }
+        int projectedFieldIndex = Integer.valueOf(fieldName.substring(5));
+        this.projectedFieldIdsToIdxInReturnedRow[projectedFieldIndex] = this.numOfFieldsInReturnedRow++;
+      }
+    }
+
+    @Override
+    public InternalRow read(Decoder decoder, Object reuse) throws IOException {
       int index = decoder.readIndex();
-      if (index == nullIndex) {
+      if (index == nullTypeIndex) {
         // if it is a null data, directly return null as the whole union result
         // we know for sure it is a null so the casting will always work.
-        return (InternalRow) readers[nullIndex].read(decoder, reuse);
+        return (InternalRow) readers[nullTypeIndex].read(decoder, reuse);
       }
 
       // otherwise, we need to return an InternalRow as a struct data
-      InternalRow struct = new GenericInternalRow(nullIndex >= 0 ? alts.size() : alts.size() + 1);
+      InternalRow struct = new GenericInternalRow(numOfFieldsInReturnedRow);
       for (int i = 0; i < struct.numFields(); i += 1) {
         struct.setNullAt(i);
       }
+      int fieldIndex = (nullTypeIndex < 0 || index < nullTypeIndex) ? index : index - 1;
+      if (isTagFieldProjected) {
+        struct.setInt(0, fieldIndex);
+      }
 
       Object value = readers[index].read(decoder, reuse);
-
-      if (nullIndex < 0) {
-        struct.update(index + 1, value);
-        struct.setInt(0, index);
-      } else if (index < nullIndex) {
-        struct.update(index + 1, value);
-        struct.setInt(0, index);
-      } else {
-        struct.update(index, value);
-        struct.setInt(0, index - 1);
+      if (projectedFieldIdsToIdxInReturnedRow[fieldIndex] != -1) {
+        struct.update(projectedFieldIdsToIdxInReturnedRow[fieldIndex], value);
       }
 
       return struct;

spark/src/test/java/org/apache/iceberg/spark/data/TestSparkAvroUnions.java

@@ -533,4 +533,45 @@ public void testDeeplyNestedUnionSchema3() throws IOException {
     // making sure the rows can be read successfully
     Assert.assertEquals(2, rows.size());
   }
+
+  @Test
+  public void writeAndValidateRequiredComplexUnionWithProjection() throws IOException {
+    org.apache.avro.Schema avroSchema = SchemaBuilder.record("root")
+        .fields()
+        .name("unionCol")
+        .type()
+        .unionOf()
+        .intType()
+        .and()
+        .stringType()
+        .endUnion()
+        .noDefault()
+        .endRecord();
+
+    GenericData.Record unionRecord1 = new GenericData.Record(avroSchema);
+    unionRecord1.put("unionCol", "foo");
+    GenericData.Record unionRecord2 = new GenericData.Record(avroSchema);
+    unionRecord2.put("unionCol", 1);
+
+    File testFile = temp.newFile();
+    try (DataFileWriter<GenericData.Record> writer = new DataFileWriter<>(new GenericDatumWriter<>())) {
+      writer.create(avroSchema, testFile);
+      writer.append(unionRecord1);
+      writer.append(unionRecord2);
+    }
+
+    Schema expectedSchema = AvroSchemaUtil.toIceberg(avroSchema).select("unionCol.field0");
+
+    List<InternalRow> rows;
+    try (AvroIterable<InternalRow> reader = Avro.read(Files.localInput(testFile))
+        .createReaderFunc(SparkAvroReader::new)
+        .project(expectedSchema)
+        .build()) {
+      rows = Lists.newArrayList(reader);
+
+      Assert.assertEquals(1, rows.get(0).getStruct(0, 1).numFields());
+      Assert.assertTrue(rows.get(0).getStruct(0, 1).isNullAt(0));
+      Assert.assertEquals(1, rows.get(1).getStruct(0, 1).getInt(0));
+    }
+  }
 }