Implement dereference pushdown for the Iceberg connector

Author: alexjo2144

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveApplyProjectionUtil.java

@@ -27,7 +27,7 @@
 
 import static java.util.Objects.requireNonNull;
 
-final class HiveApplyProjectionUtil
+public final class HiveApplyProjectionUtil
 {
     private HiveApplyProjectionUtil() {}
 

plugin/trino-hive/src/main/java/io/trino/plugin/hive/ReaderColumns.java

@@ -26,13 +26,13 @@
  * - the projected columns required by a connector level pagesource and
  * -  the columns supplied by format-specific page source
  * <p>
- * Currently used in {@link HivePageSource}.
+ * Currently used in {@link HivePageSource} and {@code io.trino.plugin.iceberg.IcebergPageSource}.
  */
 public class ReaderColumns
 {
     // columns to be read by the reader (ordered)
     private final List<ColumnHandle> readerColumns;
-    // indices for mapping expected hive column handles to the reader's column handles
+    // indices for mapping expected column handles to the reader's column handles
     private final List<Integer> readerBlockIndices;
 
     public ReaderColumns(List<? extends ColumnHandle> readerColumns, List<Integer> readerBlockIndices)

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ColumnIdentity.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.iceberg;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -40,6 +41,7 @@ public class ColumnIdentity
     private final TypeCategory typeCategory;
     // Underlying ImmutableMap is used to maintain the column ordering
     private final Map<Integer, ColumnIdentity> children;
+    private final Map<Integer, Integer> childFieldIdToIndex;
 
     @JsonCreator
     public ColumnIdentity(
@@ -56,10 +58,14 @@ public ColumnIdentity(
                 children.isEmpty() == (typeCategory == PRIMITIVE),
                 "Children should be empty if and only if column type is primitive");
         ImmutableMap.Builder<Integer, ColumnIdentity> childrenBuilder = ImmutableMap.builder();
-        for (ColumnIdentity child : children) {
+        ImmutableMap.Builder<Integer, Integer> childFieldIdToIndex = ImmutableMap.builder();
+        for (int i = 0; i < children.size(); i++) {
+            ColumnIdentity child = children.get(i);
             childrenBuilder.put(child.getId(), child);
+            childFieldIdToIndex.put(child.getId(), i);
         }
         this.children = childrenBuilder.build();
+        this.childFieldIdToIndex = childFieldIdToIndex.build();
     }
 
     @JsonProperty
@@ -86,6 +92,20 @@ public List<ColumnIdentity> getChildren()
         return ImmutableList.copyOf(children.values());
     }
 
+    @JsonIgnore
+    public ColumnIdentity getChildByFieldId(int fieldId)
+    {
+        checkArgument(children.containsKey(fieldId), "ColumnIdentity %s does not contain child with field id %s", this, fieldId);
+        return children.get(fieldId);
+    }
+
+    @JsonIgnore
+    public int getChildIndexByFieldId(int fieldId)
+    {
+        checkArgument(childFieldIdToIndex.containsKey(fieldId), "ColumnIdentity %s does not contain child with field id %s", this, fieldId);
+        return childFieldIdToIndex.get(fieldId);
+    }
+
     @Override
     public boolean equals(Object o)
     {

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ExpressionConverter.java

@@ -82,7 +82,7 @@ public static Expression toIcebergExpression(TupleDomain<IcebergColumnHandle> tu
         for (Map.Entry<IcebergColumnHandle, Domain> entry : domainMap.entrySet()) {
             IcebergColumnHandle columnHandle = entry.getKey();
             Domain domain = entry.getValue();
-            expression = and(expression, toIcebergExpression(columnHandle.getName(), columnHandle.getType(), domain));
+            expression = and(expression, toIcebergExpression(columnHandle.getQualifiedName(), columnHandle.getType(), domain));
         }
         return expression;
     }

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergColumnHandle.java

@@ -16,40 +16,52 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.type.Type;
-import io.trino.spi.type.TypeManager;
-import org.apache.iceberg.types.Types;
 
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 
-import static io.trino.plugin.iceberg.ColumnIdentity.createColumnIdentity;
-import static io.trino.plugin.iceberg.ColumnIdentity.primitiveColumnIdentity;
-import static io.trino.plugin.iceberg.TypeConverter.toTrinoType;
 import static java.util.Objects.requireNonNull;
 
 public class IcebergColumnHandle
         implements ColumnHandle
 {
-    private final ColumnIdentity columnIdentity;
+    private final ColumnIdentity baseColumnIdentity;
+    private final Type baseType;
+    // The list of field ids to indicate the projected part of the top-level column represented by baseColumnIdentity
+    private final List<Integer> path;
     private final Type type;
     private final Optional<String> comment;
+    // Cache of ColumnIdentity#getId to ensure quick access, even with dereferences
+    private final int id;
 
     @JsonCreator
     public IcebergColumnHandle(
-            @JsonProperty("columnIdentity") ColumnIdentity columnIdentity,
+            @JsonProperty("baseColumnIdentity") ColumnIdentity baseColumnIdentity,
+            @JsonProperty("baseType") Type baseType,
+            @JsonProperty("path") List<Integer> path,
             @JsonProperty("type") Type type,
             @JsonProperty("comment") Optional<String> comment)
     {
-        this.columnIdentity = requireNonNull(columnIdentity, "columnIdentity is null");
+        this.baseColumnIdentity = requireNonNull(baseColumnIdentity, "baseColumnIdentity is null");
+        this.baseType = requireNonNull(baseType, "baseType is null");
+        this.path = ImmutableList.copyOf(requireNonNull(path, "path is null"));
         this.type = requireNonNull(type, "type is null");
         this.comment = requireNonNull(comment, "comment is null");
+        this.id = path.isEmpty() ? baseColumnIdentity.getId() : Iterables.getLast(path);
     }
 
-    @JsonProperty
+    @JsonIgnore
     public ColumnIdentity getColumnIdentity()
     {
+        ColumnIdentity columnIdentity = baseColumnIdentity;
+        for (int fieldId : path) {
+            columnIdentity = columnIdentity.getChildByFieldId(fieldId);
+        }
         return columnIdentity;
     }
 
@@ -59,6 +71,24 @@ public Type getType()
         return type;
     }
 
+    @JsonProperty
+    public ColumnIdentity getBaseColumnIdentity()
+    {
+        return baseColumnIdentity;
+    }
+
+    @JsonProperty
+    public Type getBaseType()
+    {
+        return baseType;
+    }
+
+    @JsonIgnore
+    public IcebergColumnHandle getBaseColumn()
+    {
+        return new IcebergColumnHandle(getBaseColumnIdentity(), getBaseType(), ImmutableList.of(), getBaseType(), Optional.empty());
+    }
+
     @JsonProperty
     public Optional<String> getComment()
     {
@@ -68,19 +98,50 @@ public Optional<String> getComment()
     @JsonIgnore
     public int getId()
     {
-        return columnIdentity.getId();
+        return id;
     }
 
+    /**
+     * For nested columns, this is the unqualified name of the last field in the path
+     */
     @JsonIgnore
     public String getName()
     {
-        return columnIdentity.getName();
+        return getColumnIdentity().getName();
+    }
+
+    @JsonProperty
+    public List<Integer> getPath()
+    {
+        return path;
+    }
+
+    /**
+     * The dot separated path components used to address this column, including all dereferences and the column name.
+     */
+    @JsonIgnore
+    public String getQualifiedName()
+    {
+        ImmutableList.Builder<String> pathNames = ImmutableList.builder();
+        ColumnIdentity columnIdentity = baseColumnIdentity;
+        pathNames.add(columnIdentity.getName());
+        for (int fieldId : path) {
+            columnIdentity = columnIdentity.getChildByFieldId(fieldId);
+            pathNames.add(columnIdentity.getName());
+        }
+        // Iceberg tables are guaranteed not to have ambiguous column names so joining them like this must uniquely identify a single column.
+        return String.join(".", pathNames.build());
+    }
+
+    public boolean isBaseColumn()
+    {
+        return path.isEmpty();
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(columnIdentity, type, comment);
+        return Objects.hash(baseColumnIdentity, baseType, path, type, comment);
     }
 
     @Override
@@ -93,7 +154,9 @@ public boolean equals(Object obj)
             return false;
         }
         IcebergColumnHandle other = (IcebergColumnHandle) obj;
-        return Objects.equals(this.columnIdentity, other.columnIdentity) &&
+        return Objects.equals(this.baseColumnIdentity, other.baseColumnIdentity) &&
+                Objects.equals(this.baseType, other.baseType) &&
+                Objects.equals(this.path, other.path) &&
                 Objects.equals(this.type, other.type) &&
                 Objects.equals(this.comment, other.comment);
     }
@@ -103,17 +166,4 @@ public String toString()
     {
         return getId() + ":" + getName() + ":" + type.getDisplayName();
     }
-
-    public static IcebergColumnHandle primitiveIcebergColumnHandle(int id, String name, Type type, Optional<String> comment)
-    {
-        return new IcebergColumnHandle(primitiveColumnIdentity(id, name), type, comment);
-    }
-
-    public static IcebergColumnHandle create(Types.NestedField column, TypeManager typeManager)
-    {
-        return new IcebergColumnHandle(
-                createColumnIdentity(column),
-                toTrinoType(column.type(), typeManager),
-                Optional.ofNullable(column.doc()));
-    }
 }

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConfig.java

@@ -37,6 +37,7 @@ public class IcebergConfig
     private CatalogType catalogType = HIVE_METASTORE;
     private Duration dynamicFilteringWaitTimeout = new Duration(0, SECONDS);
     private boolean tableStatisticsEnabled = true;
+    private boolean projectionPushdownEnabled = true;
 
     public CatalogType getCatalogType()
     {
@@ -153,4 +154,17 @@ public boolean isTableStatisticsEnabled()
     {
         return tableStatisticsEnabled;
     }
+
+    public boolean isProjectionPushdownEnabled()
+    {
+        return projectionPushdownEnabled;
+    }
+
+    @Config("iceberg.projection-pushdown-enabled")
+    @ConfigDescription("Read only required fields from a struct")
+    public IcebergConfig setProjectionPushdownEnabled(boolean projectionPushdownEnabled)
+    {
+        this.projectionPushdownEnabled = projectionPushdownEnabled;
+        return this;
+    }
 }