Column projection of union type

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

@@ -385,6 +385,7 @@ public static class ReadBuilder {
     };
     private Long start = null;
     private Long length = null;
+    private Schema fileSchema = null;
 
     private ReadBuilder(InputFile file) {
       Preconditions.checkNotNull(file, "Input file cannot be null");
@@ -446,6 +447,11 @@ public ReadBuilder classLoader(ClassLoader classLoader) {
       return this;
     }
 
+    public ReadBuilder setFileSchema(Schema fileSchema) {
+      this.fileSchema = fileSchema;
+      return this;
+    }
+
     public <D> AvroIterable<D> build() {
       Preconditions.checkNotNull(schema, "Schema is required");
       Function<Schema, DatumReader<?>> readerFunc;
@@ -458,7 +464,7 @@ public <D> AvroIterable<D> build() {
       }
 
       return new AvroIterable<>(file,
-          new ProjectionDatumReader<>(readerFunc, schema, renames, nameMapping),
+          new ProjectionDatumReader<>(readerFunc, schema, renames, nameMapping, fileSchema),
           start, length, reuseContainers);
     }
   }

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.
+  Also an extra tag field is added into the struct of Iceberg schema during the conversion.
+  The fields in the struct of Iceberg schema are expected to be stored in the same order
+  as the corresponding types in the union of Avro schema.
+  Except the tag field, the fields in the struct of Iceberg schema are the same as the types in the union of Avro schema
+  in the general case. In case of field projection, the fields in the struct of Iceberg schema only contains
+  the fields to be projected which equals to a subset of the types in the union of Avro schema.
+  Therefore, this function visits the complex union with the consideration of both cases.
+   */
+  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));
+      } else {
+        boolean relatedFieldInStructFound = false;
+        Types.StructType struct = type.asStructType();
+        if (fieldIndexInStruct < struct.fields().size() &&
+            UNION_TAG_FIELD_NAME.equals(struct.fields().get(fieldIndexInStruct).name())) {
+          fieldIndexInStruct++;
+        }
+
+        if (fieldIndexInStruct < struct.fields().size()) {
+          // 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;
+          String structFieldName = type.asStructType().fields().get(fieldIndexInStruct).name();
+          int indexFromStructFieldName = Integer.valueOf(structFieldName.substring(5));
+          if (actualTypeIndex == indexFromStructFieldName) {
+            relatedFieldInStructFound = true;
+            options.add(visit(type.asStructType().fields().get(fieldIndexInStruct).type(), schema, visitor));
+            fieldIndexInStruct++;
+          }
+        }
+
+        // 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) {

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

@@ -50,6 +50,18 @@ public ProjectionDatumReader(Function<Schema, DatumReader<?>> getReader,
     this.nameMapping = nameMapping;
   }
 
+  public ProjectionDatumReader(Function<Schema, DatumReader<?>> getReader,
+      org.apache.iceberg.Schema expectedSchema,
+      Map<String, String> renames,
+      NameMapping nameMapping,
+      Schema fileSchema) {
+    this.getReader = getReader;
+    this.expectedSchema = expectedSchema;
+    this.renames = renames;
+    this.nameMapping = nameMapping;
+    this.fileSchema = fileSchema;
+  }
+
   @Override
   public void setRowPositionSupplier(Supplier<Long> posSupplier) {
     if (wrapped instanceof SupportsRowPosition) {
@@ -59,7 +71,9 @@ public void setRowPositionSupplier(Supplier<Long> posSupplier) {
 
   @Override
   public void setSchema(Schema newFileSchema) {
-    this.fileSchema = newFileSchema;
+    if (this.fileSchema == null) {
+      this.fileSchema = newFileSchema;
+    }
     if (nameMapping == null && !AvroSchemaUtil.hasIds(fileSchema)) {
       nameMapping = MappingUtil.create(expectedSchema);
     }

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

@@ -78,6 +78,7 @@ public void writeAndValidateRequiredComplexUnion() throws IOException {
     try (AvroIterable<InternalRow> reader = Avro.read(Files.localInput(testFile))
         .createReaderFunc(SparkAvroReader::new)
         .project(expectedSchema)
+        .setFileSchema(avroSchema)
         .build()) {
       rows = Lists.newArrayList(reader);
 
@@ -533,4 +534,95 @@ 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));
+    }
+  }
+
+  @Test(expected = ArrayIndexOutOfBoundsException.class)
+  public void writeAndReadRequiredComplexUnionWithSchemaMismatch() throws IOException {
+    org.apache.avro.Schema avroWriteSchema = SchemaBuilder.record("root")
+        .fields()
+        .name("unionCol")
+        .type()
+        .unionOf()
+        .intType()
+        .and()
+        .stringType()
+        .endUnion()
+        .noDefault()
+        .endRecord();
+
+    GenericData.Record unionRecord1 = new GenericData.Record(avroWriteSchema);
+    unionRecord1.put("unionCol", "foo");
+    GenericData.Record unionRecord2 = new GenericData.Record(avroWriteSchema);
+    unionRecord2.put("unionCol", 1);
+
+    File testFile = temp.newFile();
+    try (DataFileWriter<GenericData.Record> writer = new DataFileWriter<>(new GenericDatumWriter<>())) {
+      writer.create(avroWriteSchema, testFile);
+      writer.append(unionRecord1);
+      writer.append(unionRecord2);
+    }
+
+    org.apache.avro.Schema avroReadSchema = SchemaBuilder.record("root")
+        .fields()
+        .name("unionCol")
+        .type()
+        .unionOf()
+        .stringType()
+        .and()
+        .intType()
+        .endUnion()
+        .noDefault()
+        .endRecord();
+
+    Schema expectedSchema = AvroSchemaUtil.toIceberg(avroReadSchema);
+
+    List<InternalRow> rows;
+    try (AvroIterable<InternalRow> reader = Avro.read(Files.localInput(testFile))
+        .createReaderFunc(SparkAvroReader::new)
+        .project(expectedSchema)
+        .setFileSchema(avroReadSchema)
+        .build()) {
+      rows = Lists.newArrayList(reader);
+    }
+  }
 }