Use Id based mapping for nested fields in Iceberg Connector

lib/trino-orc/src/main/java/io/trino/orc/NameBasedFieldMapper.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed 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 io.trino.orc;
+
+import com.google.common.collect.Maps;
+import io.trino.orc.OrcReader.FieldMapper;
+
+import java.util.Locale;
+import java.util.Map;
+
+import static java.util.Objects.requireNonNull;
+
+public class NameBasedFieldMapper
+        implements FieldMapper
+{
+    private final Map<String, OrcColumn> nestedColumns;
+
+    private NameBasedFieldMapper(Map<String, OrcColumn> nestedColumns)
+    {
+        this.nestedColumns = requireNonNull(nestedColumns, "nestedColumns is null");
+    }
+
+    @Override
+    public OrcColumn get(String fieldName)
+    {
+        return nestedColumns.get(fieldName);
+    }
+
+    public static FieldMapper create(OrcColumn column)
+    {
+        requireNonNull(column, "column is null");
+        Map<String, OrcColumn> nestedColumns = Maps.uniqueIndex(
+                column.getNestedColumns(),
+                field -> field.getColumnName().toLowerCase(Locale.ENGLISH));
+
+        return new NameBasedFieldMapper(nestedColumns);
+    }
+}

lib/trino-orc/src/main/java/io/trino/orc/OrcReader.java

@@ -261,38 +261,15 @@ public OrcRecordReader createRecordReader(
         return createRecordReader(
                 readColumns,
                 readTypes,
+                Collections.nCopies(readColumns.size(), ProjectedLayout.fullyProjectedLayout()),
                 predicate,
                 0,
                 orcDataSource.getEstimatedSize(),
                 legacyFileTimeZone,
                 systemMemoryUsage,
                 initialBatchSize,
-                exceptionTransform);
-    }
-
-    public OrcRecordReader createRecordReader(
-            List<OrcColumn> readColumns,
-            List<Type> readTypes,
-            OrcPredicate predicate,
-            long offset,
-            long length,
-            DateTimeZone legacyFileTimeZone,
-            AggregatedMemoryContext systemMemoryUsage,
-            int initialBatchSize,
-            Function<Exception, RuntimeException> exceptionTransform)
-            throws OrcCorruptionException
-    {
-        return createRecordReader(
-                readColumns,
-                readTypes,
-                Collections.nCopies(readColumns.size(), ProjectedLayout.fullyProjectedLayout()),
-                predicate,
-                offset,
-                length,
-                legacyFileTimeZone,
-                systemMemoryUsage,
-                initialBatchSize,
-                exceptionTransform);
+                exceptionTransform,
+                NameBasedFieldMapper::create);
     }
 
     public OrcRecordReader createRecordReader(
@@ -305,7 +282,8 @@ public OrcRecordReader createRecordReader(
             DateTimeZone legacyFileTimeZone,
             AggregatedMemoryContext systemMemoryUsage,
             int initialBatchSize,
-            Function<Exception, RuntimeException> exceptionTransform)
+            Function<Exception, RuntimeException> exceptionTransform,
+            FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         return new OrcRecordReader(
@@ -331,7 +309,8 @@ public OrcRecordReader createRecordReader(
                 systemMemoryUsage,
                 writeValidation,
                 initialBatchSize,
-                exceptionTransform);
+                exceptionTransform,
+                fieldMapperFactory);
     }
 
     private static OrcDataSource wrapWithCacheIfTiny(OrcDataSource dataSource, DataSize maxCacheSize)
@@ -498,4 +477,15 @@ public static ProjectedLayout createProjectedLayout(OrcColumn root, List<List<St
             return new ProjectedLayout(Optional.of(fieldLayouts.build()));
         }
     }
+
+    public interface FieldMapperFactory
+    {
+        FieldMapper create(OrcColumn orcColumn);
+    }
+
+    // Used for mapping a nested field with the appropriate OrcColumn
+    public interface FieldMapper
+    {
+        OrcColumn get(String fieldName);
+    }
 }

lib/trino-orc/src/main/java/io/trino/orc/OrcRecordReader.java

@@ -24,6 +24,7 @@
 import io.airlift.units.DataSize;
 import io.trino.memory.context.AggregatedMemoryContext;
 import io.trino.memory.context.LocalMemoryContext;
+import io.trino.orc.OrcReader.FieldMapperFactory;
 import io.trino.orc.OrcWriteValidation.StatisticsValidation;
 import io.trino.orc.OrcWriteValidation.WriteChecksum;
 import io.trino.orc.OrcWriteValidation.WriteChecksumBuilder;
@@ -143,7 +144,8 @@ public OrcRecordReader(
             AggregatedMemoryContext systemMemoryUsage,
             Optional<OrcWriteValidation> writeValidation,
             int initialBatchSize,
-            Function<Exception, RuntimeException> exceptionTransform)
+            Function<Exception, RuntimeException> exceptionTransform,
+            FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         requireNonNull(readColumns, "readColumns is null");
@@ -239,7 +241,14 @@ public OrcRecordReader(
                 metadataReader,
                 writeValidation);
 
-        columnReaders = createColumnReaders(readColumns, readTypes, readLayouts, streamReadersSystemMemoryContext, blockFactory);
+        columnReaders = createColumnReaders(
+                readColumns,
+                readTypes,
+                readLayouts,
+                streamReadersSystemMemoryContext,
+                blockFactory,
+                fieldMapperFactory);
+
         currentBytesPerCell = new long[columnReaders.length];
         maxBytesPerCell = new long[columnReaders.length];
         nextBatchSize = initialBatchSize;
@@ -558,15 +567,16 @@ private static ColumnReader[] createColumnReaders(
             List<Type> readTypes,
             List<OrcReader.ProjectedLayout> readLayouts,
             AggregatedMemoryContext systemMemoryContext,
-            OrcBlockFactory blockFactory)
+            OrcBlockFactory blockFactory,
+            FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         ColumnReader[] columnReaders = new ColumnReader[columns.size()];
         for (int columnIndex = 0; columnIndex < columns.size(); columnIndex++) {
             Type readType = readTypes.get(columnIndex);
             OrcColumn column = columns.get(columnIndex);
             OrcReader.ProjectedLayout projectedLayout = readLayouts.get(columnIndex);
-            columnReaders[columnIndex] = createColumnReader(readType, column, projectedLayout, systemMemoryContext, blockFactory);
+            columnReaders[columnIndex] = createColumnReader(readType, column, projectedLayout, systemMemoryContext, blockFactory, fieldMapperFactory);
         }
         return columnReaders;
     }

lib/trino-orc/src/main/java/io/trino/orc/reader/ColumnReaders.java

@@ -18,6 +18,7 @@
 import io.trino.orc.OrcColumn;
 import io.trino.orc.OrcCorruptionException;
 import io.trino.orc.OrcReader;
+import io.trino.orc.OrcReader.FieldMapperFactory;
 import io.trino.spi.type.TimeType;
 import io.trino.spi.type.Type;
 
@@ -34,7 +35,8 @@ public static ColumnReader createColumnReader(
             OrcColumn column,
             OrcReader.ProjectedLayout projectedLayout,
             AggregatedMemoryContext systemMemoryContext,
-            OrcBlockFactory blockFactory)
+            OrcBlockFactory blockFactory,
+            FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         if (type instanceof TimeType) {
@@ -68,15 +70,15 @@ public static ColumnReader createColumnReader(
             case TIMESTAMP_INSTANT:
                 return new TimestampColumnReader(type, column, systemMemoryContext.newLocalMemoryContext(ColumnReaders.class.getSimpleName()));
             case LIST:
-                return new ListColumnReader(type, column, systemMemoryContext, blockFactory);
+                return new ListColumnReader(type, column, systemMemoryContext, blockFactory, fieldMapperFactory);
             case STRUCT:
-                return new StructColumnReader(type, column, projectedLayout, systemMemoryContext, blockFactory);
+                return new StructColumnReader(type, column, projectedLayout, systemMemoryContext, blockFactory, fieldMapperFactory);
             case MAP:
-                return new MapColumnReader(type, column, systemMemoryContext, blockFactory);
+                return new MapColumnReader(type, column, systemMemoryContext, blockFactory, fieldMapperFactory);
             case DECIMAL:
                 return new DecimalColumnReader(type, column, systemMemoryContext.newLocalMemoryContext(ColumnReaders.class.getSimpleName()));
             case UNION:
-                return new UnionColumnReader(type, column, systemMemoryContext, blockFactory);
+                return new UnionColumnReader(type, column, systemMemoryContext, blockFactory, fieldMapperFactory);
             default:
                 throw new IllegalArgumentException("Unsupported type: " + column.getColumnType());
         }

lib/trino-orc/src/main/java/io/trino/orc/reader/ListColumnReader.java

@@ -18,6 +18,7 @@
 import io.trino.orc.OrcBlockFactory;
 import io.trino.orc.OrcColumn;
 import io.trino.orc.OrcCorruptionException;
+import io.trino.orc.OrcReader.FieldMapperFactory;
 import io.trino.orc.metadata.ColumnEncoding;
 import io.trino.orc.metadata.ColumnMetadata;
 import io.trino.orc.stream.BooleanInputStream;
@@ -73,7 +74,7 @@ public class ListColumnReader
 
     private boolean rowGroupOpen;
 
-    public ListColumnReader(Type type, OrcColumn column, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory)
+    public ListColumnReader(Type type, OrcColumn column, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory, FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         requireNonNull(type, "type is null");
@@ -82,7 +83,13 @@ public ListColumnReader(Type type, OrcColumn column, AggregatedMemoryContext sys
 
         this.column = requireNonNull(column, "column is null");
         this.blockFactory = requireNonNull(blockFactory, "blockFactory is null");
-        this.elementColumnReader = createColumnReader(elementType, column.getNestedColumns().get(0), fullyProjectedLayout(), systemMemoryContext, blockFactory);
+        this.elementColumnReader = createColumnReader(
+                elementType,
+                column.getNestedColumns().get(0),
+                fullyProjectedLayout(),
+                systemMemoryContext,
+                blockFactory,
+                fieldMapperFactory);
     }
 
     @Override

lib/trino-orc/src/main/java/io/trino/orc/reader/MapColumnReader.java

@@ -18,6 +18,7 @@
 import io.trino.orc.OrcBlockFactory;
 import io.trino.orc.OrcColumn;
 import io.trino.orc.OrcCorruptionException;
+import io.trino.orc.OrcReader.FieldMapperFactory;
 import io.trino.orc.metadata.ColumnEncoding;
 import io.trino.orc.metadata.ColumnMetadata;
 import io.trino.orc.stream.BooleanInputStream;
@@ -77,7 +78,7 @@ public class MapColumnReader
 
     private boolean rowGroupOpen;
 
-    public MapColumnReader(Type type, OrcColumn column, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory)
+    public MapColumnReader(Type type, OrcColumn column, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory, FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         requireNonNull(type, "type is null");
@@ -86,8 +87,20 @@ public MapColumnReader(Type type, OrcColumn column, AggregatedMemoryContext syst
 
         this.column = requireNonNull(column, "column is null");
         this.blockFactory = requireNonNull(blockFactory, "blockFactory is null");
-        this.keyColumnReader = createColumnReader(this.type.getKeyType(), column.getNestedColumns().get(0), fullyProjectedLayout(), systemMemoryContext, blockFactory);
-        this.valueColumnReader = createColumnReader(this.type.getValueType(), column.getNestedColumns().get(1), fullyProjectedLayout(), systemMemoryContext, blockFactory);
+        this.keyColumnReader = createColumnReader(
+                this.type.getKeyType(),
+                column.getNestedColumns().get(0),
+                fullyProjectedLayout(),
+                systemMemoryContext,
+                blockFactory,
+                fieldMapperFactory);
+        this.valueColumnReader = createColumnReader(
+                this.type.getValueType(),
+                column.getNestedColumns().get(1),
+                fullyProjectedLayout(),
+                systemMemoryContext,
+                blockFactory,
+                fieldMapperFactory);
     }
 
     @Override

lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java

@@ -21,6 +21,8 @@
 import io.trino.orc.OrcColumn;
 import io.trino.orc.OrcCorruptionException;
 import io.trino.orc.OrcReader;
+import io.trino.orc.OrcReader.FieldMapper;
+import io.trino.orc.OrcReader.FieldMapperFactory;
 import io.trino.orc.metadata.ColumnEncoding;
 import io.trino.orc.metadata.ColumnMetadata;
 import io.trino.orc.stream.BooleanInputStream;
@@ -47,7 +49,6 @@
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Verify.verify;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static io.trino.orc.metadata.Stream.StreamKind.PRESENT;
 import static io.trino.orc.reader.ColumnReaders.createColumnReader;
 import static io.trino.orc.reader.ReaderUtils.verifyStreamType;
@@ -75,7 +76,13 @@ public class StructColumnReader
 
     private boolean rowGroupOpen;
 
-    StructColumnReader(Type type, OrcColumn column, OrcReader.ProjectedLayout readLayout, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory)
+    StructColumnReader(
+            Type type,
+            OrcColumn column,
+            OrcReader.ProjectedLayout readLayout,
+            AggregatedMemoryContext systemMemoryContext,
+            OrcBlockFactory blockFactory,
+            FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         requireNonNull(type, "type is null");
@@ -85,9 +92,7 @@ public class StructColumnReader
         this.column = requireNonNull(column, "column is null");
         this.blockFactory = requireNonNull(blockFactory, "blockFactory is null");
 
-        Map<String, OrcColumn> nestedColumns = column.getNestedColumns().stream()
-                .collect(toImmutableMap(stream -> stream.getColumnName().toLowerCase(Locale.ENGLISH), stream -> stream));
-
+        FieldMapper fieldMapper = fieldMapperFactory.create(column);
         ImmutableList.Builder<String> fieldNames = ImmutableList.builder();
         ImmutableMap.Builder<String, ColumnReader> structFields = ImmutableMap.builder();
         for (Field field : this.type.getFields()) {
@@ -96,12 +101,20 @@ public class StructColumnReader
                     .toLowerCase(Locale.ENGLISH);
             fieldNames.add(fieldName);
 
-            OrcColumn fieldStream = nestedColumns.get(fieldName);
+            OrcColumn fieldStream = fieldMapper.get(fieldName);
 
             if (fieldStream != null) {
                 OrcReader.ProjectedLayout fieldLayout = readLayout.getFieldLayout(fieldName);
                 if (fieldLayout != null) {
-                    structFields.put(fieldName, createColumnReader(field.getType(), fieldStream, fieldLayout, systemMemoryContext, blockFactory));
+                    structFields.put(
+                            fieldName,
+                            createColumnReader(
+                                    field.getType(),
+                                    fieldStream,
+                                    fieldLayout,
+                                    systemMemoryContext,
+                                    blockFactory,
+                                    fieldMapperFactory));
                 }
             }
         }

lib/trino-orc/src/main/java/io/trino/orc/reader/UnionColumnReader.java

@@ -19,6 +19,7 @@
 import io.trino.orc.OrcBlockFactory;
 import io.trino.orc.OrcColumn;
 import io.trino.orc.OrcCorruptionException;
+import io.trino.orc.OrcReader.FieldMapperFactory;
 import io.trino.orc.metadata.ColumnEncoding;
 import io.trino.orc.metadata.ColumnMetadata;
 import io.trino.orc.stream.BooleanInputStream;
@@ -81,7 +82,7 @@ public class UnionColumnReader
 
     private boolean rowGroupOpen;
 
-    UnionColumnReader(Type type, OrcColumn column, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory)
+    UnionColumnReader(Type type, OrcColumn column, AggregatedMemoryContext systemMemoryContext, OrcBlockFactory blockFactory, FieldMapperFactory fieldMapperFactory)
             throws OrcCorruptionException
     {
         requireNonNull(type, "type is null");
@@ -94,7 +95,13 @@ public class UnionColumnReader
         ImmutableList.Builder<ColumnReader> fieldReadersBuilder = ImmutableList.builder();
         List<OrcColumn> fields = column.getNestedColumns();
         for (int i = 0; i < fields.size(); i++) {
-            fieldReadersBuilder.add(createColumnReader(type.getTypeParameters().get(i + 1), fields.get(i), fullyProjectedLayout(), systemMemoryContext, blockFactory));
+            fieldReadersBuilder.add(createColumnReader(
+                    type.getTypeParameters().get(i + 1),
+                    fields.get(i),
+                    fullyProjectedLayout(),
+                    systemMemoryContext,
+                    blockFactory,
+                    fieldMapperFactory));
         }
         fieldReaders = fieldReadersBuilder.build();
     }