[iceberg] branch master updated: Add select and selectNot to TypeUtil for structs (#960)

[blue]

This is an automated email from the ASF dual-hosted git repository.

blue pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new 339d83b  Add select and selectNot to TypeUtil for structs (#960)
339d83b is described below

commit 339d83b0c492c66151b4b9324eeb488ed6a0ea02
Author: Ryan Blue <b...@apache.org>
AuthorDate: Wed Jul 8 18:06:20 2020 -0700

    Add select and selectNot to TypeUtil for structs (#960)
---
 .../java/org/apache/iceberg/types/TypeUtil.java    | 41 +++++++++++++++++-----
 1 file changed, 32 insertions(+), 9 deletions(-)

diff --git a/api/src/main/java/org/apache/iceberg/types/TypeUtil.java 
b/api/src/main/java/org/apache/iceberg/types/TypeUtil.java
index 36f7113..51db577 100644
--- a/api/src/main/java/org/apache/iceberg/types/TypeUtil.java
+++ b/api/src/main/java/org/apache/iceberg/types/TypeUtil.java
@@ -41,35 +41,58 @@ public class TypeUtil {
 
   public static Schema select(Schema schema, Set<Integer> fieldIds) {
     Preconditions.checkNotNull(schema, "Schema cannot be null");
-    Preconditions.checkNotNull(fieldIds, "Field ids cannot be null");
 
-    Type result = visit(schema, new PruneColumns(fieldIds));
+    Types.StructType result = select(schema.asStruct(), fieldIds);
     if (schema.asStruct() == result) {
       return schema;
     } else if (result != null) {
       if (schema.getAliases() != null) {
-        return new Schema(result.asNestedType().fields(), schema.getAliases());
+        return new Schema(result.fields(), schema.getAliases());
       } else {
-        return new Schema(result.asNestedType().fields());
+        return new Schema(result.fields());
       }
     }
 
     return new Schema(ImmutableList.of(), schema.getAliases());
   }
 
+  public static Types.StructType select(Types.StructType struct, Set<Integer> 
fieldIds) {
+    Preconditions.checkNotNull(struct, "Struct cannot be null");
+    Preconditions.checkNotNull(fieldIds, "Field ids cannot be null");
+
+    Type result = visit(struct, new PruneColumns(fieldIds));
+    if (struct == result) {
+      return struct;
+    } else if (result != null) {
+      return result.asStructType();
+    }
+
+    return Types.StructType.of();
+  }
+
   public static Set<Integer> getProjectedIds(Schema schema) {
-    return visit(schema, new GetProjectedIds());
+    return ImmutableSet.copyOf(getIdsInternal(schema.asStruct()));
   }
 
-  public static Set<Integer> getProjectedIds(Type schema) {
-    if (schema.isPrimitiveType()) {
+  public static Set<Integer> getProjectedIds(Type type) {
+    if (type.isPrimitiveType()) {
       return ImmutableSet.of();
     }
-    return ImmutableSet.copyOf(visit(schema, new GetProjectedIds()));
+    return ImmutableSet.copyOf(getIdsInternal(type));
+  }
+
+  private static Set<Integer> getIdsInternal(Type type) {
+    return visit(type, new GetProjectedIds());
+  }
+
+  public static Types.StructType selectNot(Types.StructType struct, 
Set<Integer> fieldIds) {
+    Set<Integer> projectedIds = getIdsInternal(struct);
+    projectedIds.removeAll(fieldIds);
+    return select(struct, projectedIds);
   }
 
   public static Schema selectNot(Schema schema, Set<Integer> fieldIds) {
-    Set<Integer> projectedIds = getProjectedIds(schema);
+    Set<Integer> projectedIds = getIdsInternal(schema.asStruct());
     projectedIds.removeAll(fieldIds);
     return select(schema, projectedIds);
   }