Generalize OrcReader's ProjectLayout

alexjo2144 · alexjo2144 · commit a2a9ffe2e0fb · 2021-11-10T12:41:04.000-05:00
Refactor needed for Iceberg projection pushdown. Iceberg
uses field ids to specify columns rather than column names.
diff --git a/lib/trino-orc/src/main/java/io/trino/orc/OrcReader.java b/lib/trino-orc/src/main/java/io/trino/orc/OrcReader.java
@@ -261,7 +261,7 @@ public OrcRecordReader createRecordReader(
         return createRecordReader(
                 readColumns,
                 readTypes,
-                Collections.nCopies(readColumns.size(), ProjectedLayout.fullyProjectedLayout()),
+                Collections.nCopies(readColumns.size(), fullyProjectedLayout()),
                 predicate,
                 0,
                 orcDataSource.getEstimatedSize(),
@@ -432,29 +432,40 @@ static void validateFile(
         }
     }
 
-    public static class ProjectedLayout
+    public interface ProjectedLayout
+    {
+        ProjectedLayout getFieldLayout(OrcColumn orcColumn);
+    }
+
+    /**
+     * Constructs a ProjectedLayout where all subfields must be read
+     */
+    public static ProjectedLayout fullyProjectedLayout()
+    {
+        return orcColumn -> fullyProjectedLayout();
+    }
+
+    public static class NameBasedProjectedLayout
+            implements ProjectedLayout
     {
         private final Optional<Map<String, ProjectedLayout>> fieldLayouts;
 
-        private ProjectedLayout(Optional<Map<String, ProjectedLayout>> fieldLayouts)
+        private NameBasedProjectedLayout(Optional<Map<String, ProjectedLayout>> fieldLayouts)
         {
             this.fieldLayouts = requireNonNull(fieldLayouts, "fieldLayouts is null");
         }
 
-        public ProjectedLayout getFieldLayout(String name)
+        @Override
+        public ProjectedLayout getFieldLayout(OrcColumn orcColumn)
         {
+            String name = orcColumn.getColumnName().toLowerCase(ENGLISH);
             if (fieldLayouts.isPresent()) {
                 return fieldLayouts.get().get(name);
             }
 
             return fullyProjectedLayout();
         }
 
-        public static ProjectedLayout fullyProjectedLayout()
-        {
-            return new ProjectedLayout(Optional.empty());
-        }
-
         public static ProjectedLayout createProjectedLayout(OrcColumn root, List<List<String>> dereferences)
         {
             if (dereferences.stream().map(List::size).anyMatch(Predicate.isEqual(0))) {
@@ -474,7 +485,7 @@ public static ProjectedLayout createProjectedLayout(OrcColumn root, List<List<St
                 }
             }
 
-            return new ProjectedLayout(Optional.of(fieldLayouts.build()));
+            return new NameBasedProjectedLayout(Optional.of(fieldLayouts.build()));
         }
     }
 
diff --git a/lib/trino-orc/src/main/java/io/trino/orc/reader/ListColumnReader.java b/lib/trino-orc/src/main/java/io/trino/orc/reader/ListColumnReader.java
@@ -39,7 +39,7 @@
 import java.util.Optional;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
-import static io.trino.orc.OrcReader.ProjectedLayout.fullyProjectedLayout;
+import static io.trino.orc.OrcReader.fullyProjectedLayout;
 import static io.trino.orc.metadata.Stream.StreamKind.LENGTH;
 import static io.trino.orc.metadata.Stream.StreamKind.PRESENT;
 import static io.trino.orc.reader.ColumnReaders.createColumnReader;
diff --git a/lib/trino-orc/src/main/java/io/trino/orc/reader/MapColumnReader.java b/lib/trino-orc/src/main/java/io/trino/orc/reader/MapColumnReader.java
@@ -40,7 +40,7 @@
 import java.util.Optional;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
-import static io.trino.orc.OrcReader.ProjectedLayout.fullyProjectedLayout;
+import static io.trino.orc.OrcReader.fullyProjectedLayout;
 import static io.trino.orc.metadata.Stream.StreamKind.LENGTH;
 import static io.trino.orc.metadata.Stream.StreamKind.PRESENT;
 import static io.trino.orc.reader.ColumnReaders.createColumnReader;
diff --git a/lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java b/lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java
@@ -104,7 +104,7 @@ public class StructColumnReader
             OrcColumn fieldStream = fieldMapper.get(fieldName);
 
             if (fieldStream != null) {
-                OrcReader.ProjectedLayout fieldLayout = readLayout.getFieldLayout(fieldStream.getColumnName().toLowerCase(ENGLISH));
+                OrcReader.ProjectedLayout fieldLayout = readLayout.getFieldLayout(fieldStream);
                 if (fieldLayout != null) {
                     structFields.put(
                             fieldName,
diff --git a/lib/trino-orc/src/main/java/io/trino/orc/reader/UnionColumnReader.java b/lib/trino-orc/src/main/java/io/trino/orc/reader/UnionColumnReader.java
@@ -49,7 +49,7 @@
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
-import static io.trino.orc.OrcReader.ProjectedLayout.fullyProjectedLayout;
+import static io.trino.orc.OrcReader.fullyProjectedLayout;
 import static io.trino.orc.metadata.Stream.StreamKind.DATA;
 import static io.trino.orc.metadata.Stream.StreamKind.PRESENT;
 import static io.trino.orc.reader.ColumnReaders.createColumnReader;
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcDeleteDeltaPageSource.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcDeleteDeltaPageSource.java
@@ -48,8 +48,8 @@
 import static com.google.common.collect.Maps.uniqueIndex;
 import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext;
 import static io.trino.orc.OrcReader.MAX_BATCH_SIZE;
-import static io.trino.orc.OrcReader.ProjectedLayout.fullyProjectedLayout;
 import static io.trino.orc.OrcReader.createOrcReader;
+import static io.trino.orc.OrcReader.fullyProjectedLayout;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_CANNOT_OPEN_SPLIT;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_MISSING_DATA;
 import static io.trino.plugin.hive.acid.AcidSchema.ACID_COLUMN_BUCKET;
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java
@@ -77,8 +77,8 @@
 import static com.google.common.collect.Maps.uniqueIndex;
 import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext;
 import static io.trino.orc.OrcReader.INITIAL_BATCH_SIZE;
-import static io.trino.orc.OrcReader.ProjectedLayout.createProjectedLayout;
-import static io.trino.orc.OrcReader.ProjectedLayout.fullyProjectedLayout;
+import static io.trino.orc.OrcReader.NameBasedProjectedLayout.createProjectedLayout;
+import static io.trino.orc.OrcReader.fullyProjectedLayout;
 import static io.trino.orc.metadata.OrcMetadataWriter.PRESTO_WRITER_ID;
 import static io.trino.orc.metadata.OrcMetadataWriter.TRINO_WRITER_ID;
 import static io.trino.orc.metadata.OrcType.OrcTypeKind.INT;
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProvider.java
@@ -93,7 +93,7 @@
 import static com.google.common.collect.Maps.uniqueIndex;
 import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext;
 import static io.trino.orc.OrcReader.INITIAL_BATCH_SIZE;
-import static io.trino.orc.OrcReader.ProjectedLayout.fullyProjectedLayout;
+import static io.trino.orc.OrcReader.fullyProjectedLayout;
 import static io.trino.parquet.ParquetTypeUtils.getColumnIO;
 import static io.trino.parquet.ParquetTypeUtils.getDescriptors;
 import static io.trino.parquet.ParquetTypeUtils.getParquetTypeByName;

Original file line number	Diff line number	Diff line change
`@@ -261,7 +261,7 @@ public OrcRecordReader createRecordReader(`
`261`	`261`	`return createRecordReader(`
`262`	`262`	`readColumns,`
`263`	`263`	`readTypes,`
`264`		`- Collections.nCopies(readColumns.size(), ProjectedLayout.fullyProjectedLayout()),`
	`264`	`+ Collections.nCopies(readColumns.size(), fullyProjectedLayout()),`
`265`	`265`	`predicate,`
`266`	`266`	`0,`
`267`	`267`	`orcDataSource.getEstimatedSize(),`
`@@ -432,29 +432,40 @@ static void validateFile(`
`432`	`432`	`}`
`433`	`433`	`}`
`434`	`434`
`435`		`- public static class ProjectedLayout`
	`435`	`+ public interface ProjectedLayout`
	`436`	`+ {`
	`437`	`+ ProjectedLayout getFieldLayout(OrcColumn orcColumn);`
	`438`	`+ }`
	`439`	`+`
	`440`	`+ /**`
	`441`	`+ * Constructs a ProjectedLayout where all subfields must be read`
	`442`	`+ */`
	`443`	`+ public static ProjectedLayout fullyProjectedLayout()`
	`444`	`+ {`
	`445`	`+ return orcColumn -> fullyProjectedLayout();`
	`446`	`+ }`
	`447`	`+`
	`448`	`+ public static class NameBasedProjectedLayout`
	`449`	`+ implements ProjectedLayout`
`436`	`450`	`{`
`437`	`451`	`private final Optional<Map<String, ProjectedLayout>> fieldLayouts;`
`438`	`452`
`439`		`- private ProjectedLayout(Optional<Map<String, ProjectedLayout>> fieldLayouts)`
	`453`	`+ private NameBasedProjectedLayout(Optional<Map<String, ProjectedLayout>> fieldLayouts)`
`440`	`454`	`{`
`441`	`455`	`this.fieldLayouts = requireNonNull(fieldLayouts, "fieldLayouts is null");`
`442`	`456`	`}`
`443`	`457`
`444`		`- public ProjectedLayout getFieldLayout(String name)`
	`458`	`+ @Override`
	`459`	`+ public ProjectedLayout getFieldLayout(OrcColumn orcColumn)`
`445`	`460`	`{`
	`461`	`+ String name = orcColumn.getColumnName().toLowerCase(ENGLISH);`
`446`	`462`	`if (fieldLayouts.isPresent()) {`
`447`	`463`	`return fieldLayouts.get().get(name);`
`448`	`464`	`}`
`449`	`465`
`450`	`466`	`return fullyProjectedLayout();`
`451`	`467`	`}`
`452`	`468`
`453`		`- public static ProjectedLayout fullyProjectedLayout()`
`454`		`- {`
`455`		`- return new ProjectedLayout(Optional.empty());`
`456`		`- }`
`457`		`-`
`458`	`469`	`public static ProjectedLayout createProjectedLayout(OrcColumn root, List<List<String>> dereferences)`
`459`	`470`	`{`
`460`	`471`	`if (dereferences.stream().map(List::size).anyMatch(Predicate.isEqual(0))) {`
`@@ -474,7 +485,7 @@ public static ProjectedLayout createProjectedLayout(OrcColumn root, List<List<St`
`474`	`485`	`}`
`475`	`486`	`}`
`476`	`487`
`477`		`- return new ProjectedLayout(Optional.of(fieldLayouts.build()));`
	`488`	`+ return new NameBasedProjectedLayout(Optional.of(fieldLayouts.build()));`
`478`	`489`	`}`
`479`	`490`	`}`
`480`	`491`