Core: add variant type support

apache · Dec 19, 2024 · 68d690e · 68d690e
1 parent 88a2596
commit 68d690e
Show file tree

Hide file tree

Showing 14 changed files with 217 additions and 86 deletions.
diff --git a/api/src/main/java/org/apache/iceberg/transforms/Identity.java b/api/src/main/java/org/apache/iceberg/transforms/Identity.java
@@ -19,14 +19,17 @@
 package org.apache.iceberg.transforms;
 
 import java.io.ObjectStreamException;
+import java.util.Set;
 import org.apache.iceberg.expressions.BoundPredicate;
 import org.apache.iceberg.expressions.Expressions;
 import org.apache.iceberg.expressions.UnboundPredicate;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.types.Type;
+import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.SerializableFunction;
 
 class Identity<T> implements Transform<T, T> {
+  private static final Set<Type> UNSUPPORTED_TYPES = Set.of(Types.VariantType.get());
   private static final Identity<?> INSTANCE = new Identity<>();
 
   private final Type type;
@@ -39,7 +42,7 @@ class Identity<T> implements Transform<T, T> {
   @Deprecated
   public static <I> Identity<I> get(Type type) {
     Preconditions.checkArgument(
-        type.typeId() != Type.TypeID.VARIANT, "Unsupported type for identity: %s", type);
+        !UNSUPPORTED_TYPES.contains(type), "Unsupported type for identity: %s", type);
 
     return new Identity<>(type);
   }

diff --git a/api/src/main/java/org/apache/iceberg/types/AssignFreshIds.java b/api/src/main/java/org/apache/iceberg/types/AssignFreshIds.java
@@ -124,6 +124,11 @@ public Type map(Types.MapType map, Supplier<Type> keyFuture, Supplier<Type> valu
     }
   }
 
+  @Override
+  public Type variant() {
+    return Types.VariantType.get();
+  }
+
   @Override
   public Type primitive(Type.PrimitiveType primitive) {
     return primitive;

diff --git a/api/src/main/java/org/apache/iceberg/types/GetProjectedIds.java b/api/src/main/java/org/apache/iceberg/types/GetProjectedIds.java
@@ -47,7 +47,9 @@ public Set<Integer> struct(Types.StructType struct, List<Set<Integer>> fieldResu
 
   @Override
   public Set<Integer> field(Types.NestedField field, Set<Integer> fieldResult) {
-    if ((includeStructIds && field.type().isStructType()) || field.type().isPrimitiveType()) {
+    if ((includeStructIds && field.type().isStructType())
+        || field.type().isPrimitiveType()
+        || field.type().isVariantType()) {
       fieldIds.add(field.fieldId());
     }
     return fieldIds;

diff --git a/api/src/main/java/org/apache/iceberg/types/ReassignIds.java b/api/src/main/java/org/apache/iceberg/types/ReassignIds.java
@@ -161,6 +161,11 @@ public Type map(Types.MapType map, Supplier<Type> keyTypeFuture, Supplier<Type>
     }
   }
 
+  @Override
+  public Type variant() {
+    return Types.VariantType.get();
+  }
+
   @Override
   public Type primitive(Type.PrimitiveType primitive) {
     return primitive; // nothing to reassign

diff --git a/api/src/main/java/org/apache/iceberg/types/Type.java b/api/src/main/java/org/apache/iceberg/types/Type.java
@@ -97,6 +97,10 @@ default boolean isMapType() {
     return false;
   }
 
+  default boolean isVariantType() {
+    return false;
+  }
+
   default NestedType asNestedType() {
     throw new IllegalArgumentException("Not a nested type: " + this);
   }

diff --git a/api/src/main/java/org/apache/iceberg/types/TypeUtil.java b/api/src/main/java/org/apache/iceberg/types/TypeUtil.java
@@ -709,6 +709,10 @@ public T map(Types.MapType map, Supplier<T> keyResult, Supplier<T> valueResult)
       return null;
     }
 
+    public T variant() {
+      return null;
+    }
+
     public T primitive(Type.PrimitiveType primitive) {
       return null;
     }
@@ -785,6 +789,9 @@ public static <T> T visit(Type type, CustomOrderSchemaVisitor<T> visitor) {
             new VisitFuture<>(map.keyType(), visitor),
             new VisitFuture<>(map.valueType(), visitor));
 
+      case VARIANT:
+        return visitor.variant();
+
       default:
         return visitor.primitive(type.asPrimitiveType());
     }

diff --git a/api/src/main/java/org/apache/iceberg/types/Types.java b/api/src/main/java/org/apache/iceberg/types/Types.java
@@ -429,6 +429,11 @@ public String toString() {
       return "variant";
     }
 
+    @Override
+    public boolean isVariantType() {
+      return true;
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) {