[LI] Fix SPARK ORC projection read for deeply nested union schema and… (

#154)

* [LI] Fix SPARK ORC projection read for deeply nested union schema and out-of-order projection fields

* fix style

* resolve comment

orc/src/main/java/org/apache/iceberg/orc/OrcSchemaWithTypeVisitor.java

@@ -19,15 +19,17 @@
 
 package org.apache.iceberg.orc;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 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.types.Types.StructType;
 import org.apache.orc.TypeDescription;
 
 public abstract class OrcSchemaWithTypeVisitor<T> {
-  private static final String PSEUDO_ICEBERG_FIELD_ID = "-1";
 
   public static <T> T visit(
       org.apache.iceberg.Schema iSchema, TypeDescription schema, OrcSchemaWithTypeVisitor<T> visitor) {
@@ -96,39 +98,38 @@ protected T visitUnion(Type type, TypeDescription union, OrcSchemaWithTypeVisito
   as the corresponding types in the union of Orc schema.
   Except the tag field, the fields in the struct of Iceberg schema are the same as the types in the union of Orc 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 ORC schema.
-  Therefore, this function visits the complex union with the consideration of both cases.
+  a subset of the types in the union of ORC schema, but all the readers for the union branch types must be constructed,
+  it's up to the reader code logic to determine to return what value for the given projection schema.
+  Therefore, this function visits the complex union with the consideration of either whole projection
+  or partially projected schema.
   Noted that null value and default value for complex union is not a consideration in case of ORC
    */
   private <T> void visitComplexUnion(Type type, TypeDescription union, OrcSchemaWithTypeVisitor<T> visitor,
                                      List<T> options) {
-    int typeIndex = 0;
-    int fieldIndexInStruct = 0;
-    while (typeIndex < union.getChildren().size()) {
-      TypeDescription schema = union.getChildren().get(typeIndex);
-      boolean relatedFieldInStructFound = false;
-      Types.StructType struct = type.asStructType();
-      if (fieldIndexInStruct < struct.fields().size() &&
-              ORCSchemaUtil.ICEBERG_UNION_TAG_FIELD_NAME
-                      .equals(struct.fields().get(fieldIndexInStruct).name())) {
-        fieldIndexInStruct++;
-      }
 
-      if (fieldIndexInStruct < struct.fields().size()) {
-        String structFieldName = type.asStructType().fields().get(fieldIndexInStruct).name();
-        int indexFromStructFieldName = Integer.parseInt(structFieldName
-                .substring(ORCSchemaUtil.ICEBERG_UNION_TYPE_FIELD_NAME_PREFIX_LENGTH));
-        if (typeIndex == indexFromStructFieldName) {
-          relatedFieldInStructFound = true;
-          T option = visit(type.asStructType().fields().get(fieldIndexInStruct).type(), schema, visitor);
-          options.add(option);
-          fieldIndexInStruct++;
-        }
+    StructType structType = type.asStructType();
+    List<TypeDescription> unionTypes = union.getChildren();
+    Map<Integer, Integer> idxInOrcUnionToIdxInType = new HashMap<>();
+    // Construct idxInOrcUnionToIdxInType
+    for (int i = 0; i < structType.fields().size(); i += 1) {
+      String fieldName = structType.fields().get(i).name();
+      if (!fieldName.equals(ORCSchemaUtil.ICEBERG_UNION_TAG_FIELD_NAME)) {
+        int idxInOrcUnion = Integer.parseInt(fieldName
+            .substring(ORCSchemaUtil.ICEBERG_UNION_TYPE_FIELD_NAME_PREFIX_LENGTH));
+        idxInOrcUnionToIdxInType.put(idxInOrcUnion, i);
       }
-      if (!relatedFieldInStructFound) {
-        visitNotProjectedTypeInComplexUnion(schema, visitor, options, typeIndex);
+    }
+
+    for (int i = 0; i < union.getChildren().size(); i += 1) {
+      if (idxInOrcUnionToIdxInType.containsKey(i)) {
+        options.add(visit(structType.fields().get(idxInOrcUnionToIdxInType.get(i)).type(), unionTypes.get(i), visitor));
+      } else {
+        // even if the type is not projected in the iceberg schema, a reader for the underlying orc type branch
+        // still needs to be created, we use a OrcToIcebergVisitorWithPseudoId to re-construct the iceberg type
+        // from the orc union branch type and add it to the options,
+        // with a pseudo iceberg-id "-1" to avoid failures with the remaining iceberg code infra
+        visitNotProjectedTypeInComplexUnion(unionTypes.get(i), visitor, options, i);
       }
-      typeIndex++;
     }
   }
 
@@ -137,16 +138,15 @@ private <T> void visitComplexUnion(Type type, TypeDescription union, OrcSchemaWi
   // used to make the reading of Orc file successfully. In this case, a pseudo Iceberg type is converted from
   // the Orc schema and is used to create the option for the reader of the current type which still can
   // read the corresponding content in Orc file successfully.
-  private static <T> void visitNotProjectedTypeInComplexUnion(TypeDescription schema,
+  private static <T> void visitNotProjectedTypeInComplexUnion(TypeDescription orcType,
                                                               OrcSchemaWithTypeVisitor<T> visitor,
                                                               List<T> options,
                                                               int typeIndex) {
-    OrcToIcebergVisitor schemaConverter = new OrcToIcebergVisitor();
-    schemaConverter.beforeField("field" + typeIndex, schema);
-    schema.setAttribute(org.apache.iceberg.orc.ORCSchemaUtil.ICEBERG_ID_ATTRIBUTE, PSEUDO_ICEBERG_FIELD_ID);
-    Optional<Types.NestedField> icebergSchema = OrcToIcebergVisitor.visit(schema, schemaConverter);
-    schemaConverter.afterField("field" + typeIndex, schema);
-    options.add(visit(icebergSchema.get().type(), schema, visitor));
+    OrcToIcebergVisitor schemaConverter = new OrcToIcebergVisitorWithPseudoId();
+    schemaConverter.beforeField("field" + typeIndex, orcType);
+    Optional<Types.NestedField> icebergType = OrcToIcebergVisitor.visit(orcType, schemaConverter);
+    schemaConverter.afterField("field" + typeIndex, orcType);
+    options.add(visit(icebergType.get().type(), orcType, visitor));
   }
 
   public T record(Types.StructType iStruct, TypeDescription record, List<String> names, List<T> fields) {

orc/src/main/java/org/apache/iceberg/orc/OrcToIcebergVisitorWithPseudoId.java

@@ -0,0 +1,71 @@
+/*
+ * 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.iceberg.orc;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import org.apache.iceberg.types.Types;
+import org.apache.iceberg.types.Types.NestedField;
+import org.apache.orc.TypeDescription;
+
+public class OrcToIcebergVisitorWithPseudoId extends OrcToIcebergVisitor {
+
+  private static final String PSEUDO_ICEBERG_FIELD_ID = "-1";
+
+  @Override
+  public Optional<NestedField> union(TypeDescription union, List<Optional<NestedField>> options) {
+    union.setAttribute(org.apache.iceberg.orc.ORCSchemaUtil.ICEBERG_ID_ATTRIBUTE, PSEUDO_ICEBERG_FIELD_ID);
+    List<Optional<NestedField>> optionsCopy = new ArrayList<>(options);
+    optionsCopy.add(0, Optional.of(
+        Types.NestedField.of(Integer.parseInt(PSEUDO_ICEBERG_FIELD_ID), true,
+            ORCSchemaUtil.ICEBERG_UNION_TAG_FIELD_NAME, Types.IntegerType.get())));
+    return Optional.of(
+        Types.NestedField.of(Integer.parseInt(PSEUDO_ICEBERG_FIELD_ID), true, currentFieldName(), Types.StructType.of(
+            optionsCopy.stream().filter(Optional::isPresent).map(Optional::get).collect(Collectors.toList()))));
+  }
+
+  @Override
+  public Optional<NestedField> record(TypeDescription record, List<String> names,
+      List<Optional<NestedField>> fields) {
+    record.setAttribute(org.apache.iceberg.orc.ORCSchemaUtil.ICEBERG_ID_ATTRIBUTE, PSEUDO_ICEBERG_FIELD_ID);
+    return super.record(record, names, fields);
+  }
+
+  @Override
+  public Optional<NestedField> list(TypeDescription array, Optional<NestedField> element) {
+    array.setAttribute(org.apache.iceberg.orc.ORCSchemaUtil.ICEBERG_ID_ATTRIBUTE, PSEUDO_ICEBERG_FIELD_ID);
+    return super.list(array, element);
+  }
+
+  @Override
+  public Optional<NestedField> map(TypeDescription map, Optional<NestedField> key,
+      Optional<NestedField> value) {
+    map.setAttribute(org.apache.iceberg.orc.ORCSchemaUtil.ICEBERG_ID_ATTRIBUTE, PSEUDO_ICEBERG_FIELD_ID);
+    return super.map(map, key, value);
+  }
+
+  @Override
+  public Optional<NestedField> primitive(TypeDescription primitive) {
+    primitive.setAttribute(org.apache.iceberg.orc.ORCSchemaUtil.ICEBERG_ID_ATTRIBUTE, PSEUDO_ICEBERG_FIELD_ID);
+    return super.primitive(primitive);
+  }
+}

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

@@ -20,7 +20,7 @@
 package org.apache.iceberg.spark.data;
 
 import java.math.BigDecimal;
-import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.iceberg.orc.ORCSchemaUtil;
@@ -169,38 +169,23 @@ protected void set(InternalRow struct, int pos, Object value) {
 
   static class UnionReader implements OrcValueReader<Object> {
     private final OrcValueReader[] readers;
-    private final Type expectedIcebergSchema;
-    private int[] projectedFieldIdsToIdxInReturnedRow;
-    private boolean isTagFieldProjected;
-    private int numOfFieldsInReturnedRow;
+    private final Type expectedType;
+    private Map<Integer, Integer> idxInExpectedSchemaToIdxInReaders;
 
     private UnionReader(List<OrcValueReader<?>> readers, Type expected) {
-      this.readers = new OrcValueReader[readers.size()];
-      for (int i = 0; i < this.readers.length; i += 1) {
-        this.readers[i] = readers.get(i);
-      }
-      this.expectedIcebergSchema = expected;
+      this.readers = readers.toArray(new OrcValueReader[readers.size()]);
+      this.expectedType = expected;
 
       if (this.readers.length > 1) {
-        // 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 Integer.MIN_VALUE, it means the value of this field should not be stored
-        // in the returned row
-        this.projectedFieldIdsToIdxInReturnedRow = new int[readers.size()];
-        Arrays.fill(this.projectedFieldIdsToIdxInReturnedRow, Integer.MIN_VALUE);
-        this.numOfFieldsInReturnedRow = 0;
-        this.isTagFieldProjected = false;
-
-        for (Types.NestedField expectedStructField : expectedIcebergSchema.asStructType().fields()) {
-          String fieldName = expectedStructField.name();
-          if (fieldName.equals(ORCSchemaUtil.ICEBERG_UNION_TAG_FIELD_NAME)) {
-            this.isTagFieldProjected = true;
-            this.numOfFieldsInReturnedRow++;
-            continue;
+        idxInExpectedSchemaToIdxInReaders = new HashMap<>();
+        // Construct fieldIdxInExpectedSchemaToIdxInReaders
+        for (int i = 0; i < expectedType.asStructType().fields().size(); i += 1) {
+          String fieldName = expectedType.asStructType().fields().get(i).name();
+          if (!fieldName.equals(ORCSchemaUtil.ICEBERG_UNION_TAG_FIELD_NAME)) {
+            int idxInReader = Integer.parseInt(fieldName
+                    .substring(ORCSchemaUtil.ICEBERG_UNION_TYPE_FIELD_NAME_PREFIX_LENGTH));
+            this.idxInExpectedSchemaToIdxInReaders.put(i, idxInReader);
           }
-          int projectedFieldIndex = Integer.valueOf(fieldName
-                  .substring(ORCSchemaUtil.ICEBERG_UNION_TYPE_FIELD_NAME_PREFIX_LENGTH));
-          this.projectedFieldIdsToIdxInReturnedRow[projectedFieldIndex] = this.numOfFieldsInReturnedRow++;
         }
       }
     }
@@ -214,18 +199,21 @@ public Object nonNullRead(ColumnVector vector, int row) {
       if (readers.length == 1) {
         return value;
       } else {
-        InternalRow struct = new GenericInternalRow(numOfFieldsInReturnedRow);
+        InternalRow struct = new GenericInternalRow(this.expectedType.asStructType().fields().size());
         for (int i = 0; i < struct.numFields(); i += 1) {
           struct.setNullAt(i);
         }
-        if (this.isTagFieldProjected) {
-          struct.update(0, fieldIndex);
-        }
-
-        if (this.projectedFieldIdsToIdxInReturnedRow[fieldIndex] != Integer.MIN_VALUE) {
-          struct.update(this.projectedFieldIdsToIdxInReturnedRow[fieldIndex], value);
+        for (int i = 0; i < this.expectedType.asStructType().fields().size(); i += 1) {
+          String fieldName = expectedType.asStructType().fields().get(i).name();
+          if (fieldName.equals(ORCSchemaUtil.ICEBERG_UNION_TAG_FIELD_NAME)) {
+            struct.update(i, fieldIndex);
+          } else {
+            int idxInReader = this.idxInExpectedSchemaToIdxInReaders.get(i);
+            if (idxInReader == fieldIndex) {
+              struct.update(i, value);
+            }
+          }
         }
-
         return struct;
       }
     }