Nested predicate push down to Parquet Reader

presto-hive/src/main/java/com/facebook/presto/hive/HiveColumnHandle.java

@@ -15,11 +15,13 @@
 
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
+import com.facebook.presto.spi.NestedField;
 import com.facebook.presto.spi.type.TypeManager;
 import com.facebook.presto.spi.type.TypeSignature;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 
@@ -30,6 +32,7 @@
 import static com.facebook.presto.hive.HiveType.HIVE_STRING;
 import static com.facebook.presto.spi.type.BigintType.BIGINT;
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Arrays.asList;
 import static java.util.Objects.requireNonNull;
 
 public class HiveColumnHandle
@@ -60,6 +63,7 @@ public enum ColumnType
     private final int hiveColumnIndex;
     private final ColumnType columnType;
     private final Optional<String> comment;
+    private final Optional<NestedField> nestedField;
 
     @JsonCreator
     public HiveColumnHandle(
@@ -68,7 +72,8 @@ public HiveColumnHandle(
             @JsonProperty("typeSignature") TypeSignature typeSignature,
             @JsonProperty("hiveColumnIndex") int hiveColumnIndex,
             @JsonProperty("columnType") ColumnType columnType,
-            @JsonProperty("comment") Optional<String> comment)
+            @JsonProperty("comment") Optional<String> comment,
+            @JsonProperty("nestedField") Optional<NestedField> nestedField)
     {
         this.name = requireNonNull(name, "name is null");
         checkArgument(hiveColumnIndex >= 0 || columnType == PARTITION_KEY || columnType == SYNTHESIZED, "hiveColumnIndex is negative");
@@ -77,6 +82,7 @@ public HiveColumnHandle(
         this.typeName = requireNonNull(typeSignature, "type is null");
         this.columnType = requireNonNull(columnType, "columnType is null");
         this.comment = requireNonNull(comment, "comment is null");
+        this.nestedField = requireNonNull(nestedField, "nestedField is null");
     }
 
     @JsonProperty
@@ -112,12 +118,23 @@ public ColumnMetadata getColumnMetadata(TypeManager typeManager)
         return new ColumnMetadata(name, typeManager.getType(typeName), null, isHidden());
     }
 
+    public List<String> getNameList()
+    {
+        return asList(name.split("\\."));
+    }
+
     @JsonProperty
     public Optional<String> getComment()
     {
         return comment;
     }
 
+    @JsonProperty
+    public Optional<NestedField> getNestedField()
+    {
+        return nestedField;
+    }
+
     @JsonProperty
     public TypeSignature getTypeSignature()
     {
@@ -133,7 +150,7 @@ public ColumnType getColumnType()
     @Override
     public int hashCode()
     {
-        return Objects.hash(name, hiveColumnIndex, hiveType, columnType, comment);
+        return Objects.hash(name, hiveColumnIndex, hiveType, columnType, comment, nestedField);
     }
 
     @Override
@@ -150,13 +167,14 @@ public boolean equals(Object obj)
                 Objects.equals(this.hiveColumnIndex, other.hiveColumnIndex) &&
                 Objects.equals(this.hiveType, other.hiveType) &&
                 Objects.equals(this.columnType, other.columnType) &&
+                Objects.equals(this.nestedField, other.nestedField) &&
                 Objects.equals(this.comment, other.comment);
     }
 
     @Override
     public String toString()
     {
-        return name + ":" + hiveType + ":" + hiveColumnIndex + ":" + columnType;
+        return name + ":" + hiveType + ":" + hiveColumnIndex + ":" + columnType + ":" + nestedField;
     }
 
     public static HiveColumnHandle updateRowIdHandle()
@@ -167,12 +185,12 @@ public static HiveColumnHandle updateRowIdHandle()
         // plan-time support for row-by-row delete so that planning doesn't fail. This is why we need
         // rowid handle. Note that in Hive connector, rowid handle is not implemented beyond plan-time.
 
-        return new HiveColumnHandle(UPDATE_ROW_ID_COLUMN_NAME, HIVE_LONG, BIGINT.getTypeSignature(), -1, SYNTHESIZED, Optional.empty());
+        return new HiveColumnHandle(UPDATE_ROW_ID_COLUMN_NAME, HIVE_LONG, BIGINT.getTypeSignature(), -1, SYNTHESIZED, Optional.empty(), Optional.empty());
     }
 
     public static HiveColumnHandle pathColumnHandle()
     {
-        return new HiveColumnHandle(PATH_COLUMN_NAME, PATH_HIVE_TYPE, PATH_TYPE_SIGNATURE, PATH_COLUMN_INDEX, SYNTHESIZED, Optional.empty());
+        return new HiveColumnHandle(PATH_COLUMN_NAME, PATH_HIVE_TYPE, PATH_TYPE_SIGNATURE, PATH_COLUMN_INDEX, SYNTHESIZED, Optional.empty(), Optional.empty());
     }
 
     /**
@@ -182,7 +200,7 @@ public static HiveColumnHandle pathColumnHandle()
      */
     public static HiveColumnHandle bucketColumnHandle()
     {
-        return new HiveColumnHandle(BUCKET_COLUMN_NAME, BUCKET_HIVE_TYPE, BUCKET_TYPE_SIGNATURE, BUCKET_COLUMN_INDEX, SYNTHESIZED, Optional.empty());
+        return new HiveColumnHandle(BUCKET_COLUMN_NAME, BUCKET_HIVE_TYPE, BUCKET_TYPE_SIGNATURE, BUCKET_COLUMN_INDEX, SYNTHESIZED, Optional.empty(), Optional.empty());
     }
 
     public static boolean isPathColumnHandle(HiveColumnHandle column)

presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java

@@ -45,6 +45,7 @@
 import com.facebook.presto.spi.Constraint;
 import com.facebook.presto.spi.DiscretePredicates;
 import com.facebook.presto.spi.InMemoryRecordSet;
+import com.facebook.presto.spi.NestedField;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.RecordCursor;
 import com.facebook.presto.spi.SchemaTableName;
@@ -143,6 +144,7 @@
 import static com.facebook.presto.hive.HiveSessionProperties.isSortedWritingEnabled;
 import static com.facebook.presto.hive.HiveSessionProperties.isStatisticsEnabled;
 import static com.facebook.presto.hive.HiveSessionProperties.isWritingStagingFilesEnabled;
+import static com.facebook.presto.hive.HiveStorageFormat.PARQUET;
 import static com.facebook.presto.hive.HiveTableProperties.AVRO_SCHEMA_URL;
 import static com.facebook.presto.hive.HiveTableProperties.BUCKETED_BY_PROPERTY;
 import static com.facebook.presto.hive.HiveTableProperties.BUCKET_COUNT_PROPERTY;
@@ -166,6 +168,7 @@
 import static com.facebook.presto.hive.HiveUtil.decodeViewData;
 import static com.facebook.presto.hive.HiveUtil.encodeViewData;
 import static com.facebook.presto.hive.HiveUtil.getPartitionKeyColumnHandles;
+import static com.facebook.presto.hive.HiveUtil.getRegularColumnHandles;
 import static com.facebook.presto.hive.HiveUtil.hiveColumnHandles;
 import static com.facebook.presto.hive.HiveUtil.schemaTableName;
 import static com.facebook.presto.hive.HiveUtil.toPartitionValues;
@@ -544,6 +547,32 @@ public Map<String, ColumnHandle> getColumnHandles(ConnectorSession session, Conn
         return columnHandles.build();
     }
 
+    @Override
+    public Map<NestedField, ColumnHandle> getNestedColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle, Collection<NestedField> nestedFields)
+    {
+        SchemaTableName tableName = schemaTableName(tableHandle);
+        Optional<Table> table = metastore.getTable(tableName.getSchemaName(), tableName.getTableName());
+        if (!table.isPresent()) {
+            throw new TableNotFoundException(tableName);
+        }
+
+        if (extractHiveStorageFormat(table.get()).equals(PARQUET)) {
+            List<HiveColumnHandle> regularColumnHandles = getRegularColumnHandles(table.get());
+            Map<String, HiveColumnHandle> regularHiveColumnHandles = regularColumnHandles.stream()
+                    .collect(Collectors.toMap(HiveColumnHandle::getName, identity()));
+            ImmutableMap.Builder<NestedField, ColumnHandle> nestedColumnHandles = ImmutableMap.builder();
+            for (NestedField field : nestedFields) {
+                HiveColumnHandle hiveColumnHandle = regularHiveColumnHandles.get(field.getBase());
+                Optional<HiveType> type = hiveColumnHandle.getHiveType().getFieldType(field);
+                if (hiveColumnHandle != null) {
+                    nestedColumnHandles.put(field, new HiveColumnHandle(field.getName(), type.get(), type.get().getTypeSignature(), hiveColumnHandle.getHiveColumnIndex(), hiveColumnHandle.getColumnType(), hiveColumnHandle.getComment(), Optional.of(field)));
+                }
+            }
+            return nestedColumnHandles.build();
+        }
+        return ImmutableMap.of();
+    }
+
     @SuppressWarnings("TryWithIdenticalCatches")
     @Override
     public Map<SchemaTableName, List<ColumnMetadata>> listTableColumns(ConnectorSession session, SchemaTablePrefix prefix)
@@ -2174,7 +2203,8 @@ else if (column.isHidden()) {
                     column.getType().getTypeSignature(),
                     ordinal,
                     columnType,
-                    Optional.ofNullable(column.getComment())));
+                    Optional.ofNullable(column.getComment()),
+                    Optional.empty()));
             ordinal++;
         }
 

presto-hive/src/main/java/com/facebook/presto/hive/HivePageSourceProvider.java

@@ -319,8 +319,14 @@ public static List<ColumnMapping> buildColumnMappings(
             for (HiveColumnHandle column : columns) {
                 Optional<HiveType> coercionFrom = Optional.ofNullable(columnCoercions.get(column.getHiveColumnIndex()));
                 if (column.getColumnType() == REGULAR) {
-                    checkArgument(regularColumnIndices.add(column.getHiveColumnIndex()), "duplicate hiveColumnIndex in columns list");
-                    columnMappings.add(regular(column, regularIndex, coercionFrom));
+                    if (column.getNestedField().isPresent()) {
+                        Optional<HiveType> hiveType = coercionFrom.flatMap(type -> type.getFieldType(column.getNestedField().get()));
+                        columnMappings.add(regular(column, regularIndex, hiveType));
+                    }
+                    else {
+                        checkArgument(regularColumnIndices.add(column.getHiveColumnIndex()), "duplicate hiveColumnIndex in columns list");
+                        columnMappings.add(regular(column, regularIndex, coercionFrom));
+                    }
                     regularIndex++;
                 }
                 else {
@@ -365,7 +371,8 @@ public static List<HiveColumnHandle> toColumnHandles(List<ColumnMapping> regular
                                 columnMapping.getCoercionFrom().get().getTypeSignature(),
                                 columnHandle.getHiveColumnIndex(),
                                 columnHandle.getColumnType(),
-                                Optional.empty());
+                                Optional.empty(),
+                                columnHandle.getNestedField());
                     })
                     .collect(toList());
         }

presto-hive/src/main/java/com/facebook/presto/hive/HiveType.java

@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.hive;
 
+import com.facebook.presto.spi.NestedField;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.type.NamedTypeSignature;
 import com.facebook.presto.spi.type.RowFieldName;
@@ -156,6 +157,24 @@ public boolean isSupportedType()
         return isSupportedType(getTypeInfo());
     }
 
+    public Optional<HiveType> getFieldType(NestedField nestedField)
+    {
+        TypeInfo typeInfo = getTypeInfo();
+        for (String field : nestedField.getRemaining()) {
+            if (!(typeInfo instanceof StructTypeInfo)) {
+                throw new IllegalArgumentException("Invalid type: " + typeInfo + ". expecting RowType");
+            }
+            StructTypeInfo structTypeInfo = (StructTypeInfo) typeInfo;
+            try {
+                typeInfo = structTypeInfo.getStructFieldTypeInfo(field);
+            }
+            catch (RuntimeException e) {
+                return Optional.empty();
+            }
+        }
+        return Optional.of(toHiveType(typeInfo));
+    }
+
     public static boolean isSupportedType(TypeInfo typeInfo)
     {
         switch (typeInfo.getCategory()) {

presto-hive/src/main/java/com/facebook/presto/hive/HiveUtil.java

@@ -830,7 +830,7 @@ public static List<HiveColumnHandle> getRegularColumnHandles(Table table)
             // ignore unsupported types rather than failing
             HiveType hiveType = field.getType();
             if (hiveType.isSupportedType()) {
-                columns.add(new HiveColumnHandle(field.getName(), hiveType, hiveType.getTypeSignature(), hiveColumnIndex, REGULAR, field.getComment()));
+                columns.add(new HiveColumnHandle(field.getName(), hiveType, hiveType.getTypeSignature(), hiveColumnIndex, REGULAR, field.getComment(), Optional.empty()));
             }
             hiveColumnIndex++;
         }
@@ -848,7 +848,7 @@ public static List<HiveColumnHandle> getPartitionKeyColumnHandles(Table table)
             if (!hiveType.isSupportedType()) {
                 throw new PrestoException(NOT_SUPPORTED, format("Unsupported Hive type %s found in partition keys of table %s.%s", hiveType, table.getDatabaseName(), table.getTableName()));
             }
-            columns.add(new HiveColumnHandle(field.getName(), hiveType, hiveType.getTypeSignature(), -1, PARTITION_KEY, field.getComment()));
+            columns.add(new HiveColumnHandle(field.getName(), hiveType, hiveType.getTypeSignature(), -1, PARTITION_KEY, field.getComment(), Optional.empty()));
         }
 
         return columns.build();

presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcPageSourceFactory.java

@@ -263,7 +263,7 @@ private static List<HiveColumnHandle> getPhysicalHiveColumnHandles(List<HiveColu
                 physicalOrdinal = nextMissingColumnIndex;
                 nextMissingColumnIndex++;
             }
-            physicalColumns.add(new HiveColumnHandle(column.getName(), column.getHiveType(), column.getTypeSignature(), physicalOrdinal, column.getColumnType(), column.getComment()));
+            physicalColumns.add(new HiveColumnHandle(column.getName(), column.getHiveType(), column.getTypeSignature(), physicalOrdinal, column.getColumnType(), column.getComment(), column.getNestedField()));
         }
         return physicalColumns.build();
     }