Return same variant instance and update test

apache · Feb 13, 2025 · 93586e3 · 93586e3
1 parent f53769b
commit 93586e3
Show file tree

Hide file tree

Showing 17 changed files with 47 additions and 45 deletions.
diff --git a/api/src/main/java/org/apache/iceberg/Accessors.java b/api/src/main/java/org/apache/iceberg/Accessors.java
@@ -233,7 +233,7 @@ public Map<Integer, Accessor<StructLike>> struct(
     }
 
     @Override
-    public Map<Integer, Accessor<StructLike>> variant() {
+    public Map<Integer, Accessor<StructLike>> variant(Types.VariantType variant) {
       return null;
     }
 

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

diff --git a/api/src/main/java/org/apache/iceberg/types/FindTypeVisitor.java b/api/src/main/java/org/apache/iceberg/types/FindTypeVisitor.java
@@ -77,9 +77,9 @@ public Type map(Types.MapType map, Type keyResult, Type valueResult) {
   }
 
   @Override
-  public Type variant() {
-    if (predicate.test(Types.VariantType.get())) {
-      return Types.VariantType.get();
+  public Type variant(Types.VariantType variant) {
+    if (predicate.test(variant)) {
+      return variant;
     }
 
     return null;

diff --git a/api/src/main/java/org/apache/iceberg/types/GetProjectedIds.java b/api/src/main/java/org/apache/iceberg/types/GetProjectedIds.java
@@ -76,7 +76,7 @@ public Set<Integer> map(Types.MapType map, Set<Integer> keyResult, Set<Integer>
   }
 
   @Override
-  public Set<Integer> variant() {
+  public Set<Integer> variant(Types.VariantType variant) {
     return null;
   }
 }
diff --git a/api/src/main/java/org/apache/iceberg/types/IndexById.java b/api/src/main/java/org/apache/iceberg/types/IndexById.java
@@ -66,7 +66,7 @@ public Map<Integer, Types.NestedField> map(
   }
 
   @Override
-  public Map<Integer, Types.NestedField> variant() {
+  public Map<Integer, Types.NestedField> variant(Types.VariantType variant) {
     return null;
   }
 }
diff --git a/api/src/main/java/org/apache/iceberg/types/IndexByName.java b/api/src/main/java/org/apache/iceberg/types/IndexByName.java
@@ -177,7 +177,7 @@ public Map<String, Integer> map(
   }
 
   @Override
-  public Map<String, Integer> variant() {
+  public Map<String, Integer> variant(Types.VariantType variant) {
     return nameToId;
   }
 

diff --git a/api/src/main/java/org/apache/iceberg/types/IndexParents.java b/api/src/main/java/org/apache/iceberg/types/IndexParents.java
@@ -77,7 +77,7 @@ public Map<Integer, Integer> map(
   }
 
   @Override
-  public Map<Integer, Integer> variant() {
+  public Map<Integer, Integer> variant(Types.VariantType variant) {
     return idToParent;
   }
 

diff --git a/api/src/main/java/org/apache/iceberg/types/PruneColumns.java b/api/src/main/java/org/apache/iceberg/types/PruneColumns.java
@@ -160,7 +160,7 @@ public Type map(Types.MapType map, Type ignored, Type valueResult) {
   }
 
   @Override
-  public Type variant() {
+  public Type variant(Types.VariantType variant) {
     return null;
   }
 

diff --git a/api/src/main/java/org/apache/iceberg/types/ReassignIds.java b/api/src/main/java/org/apache/iceberg/types/ReassignIds.java
@@ -158,8 +158,8 @@ public Type map(Types.MapType map, Supplier<Type> keyTypeFuture, Supplier<Type>
   }
 
   @Override
-  public Type variant() {
-    return Types.VariantType.get();
+  public Type variant(Types.VariantType variant) {
+    return variant;
   }
 
   @Override

diff --git a/api/src/main/java/org/apache/iceberg/types/Type.java b/api/src/main/java/org/apache/iceberg/types/Type.java
@@ -82,6 +82,10 @@ default Types.MapType asMapType() {
     throw new IllegalArgumentException("Not a map type: " + this);
   }
 
+  default Types.VariantType asVariantType() {
+    throw new IllegalArgumentException("Not a variant type: " + this);
+  }
+
   default boolean isNestedType() {
     return false;
   }

diff --git a/api/src/main/java/org/apache/iceberg/types/TypeUtil.java b/api/src/main/java/org/apache/iceberg/types/TypeUtil.java
@@ -616,7 +616,7 @@ public T map(Types.MapType map, T keyResult, T valueResult) {
       return null;
     }
 
-    public T variant() {
+    public T variant(Types.VariantType variant) {
       throw new UnsupportedOperationException("Unsupported type: variant");
     }
 
@@ -684,7 +684,7 @@ public static <T> T visit(Type type, SchemaVisitor<T> visitor) {
         return visitor.map(map, keyResult, valueResult);
 
       case VARIANT:
-        return visitor.variant();
+        return visitor.variant(type.asVariantType());
 
       default:
         return visitor.primitive(type.asPrimitiveType());
@@ -712,8 +712,8 @@ public T map(Types.MapType map, Supplier<T> keyResult, Supplier<T> valueResult)
       return null;
     }
 
-    public T variant() {
-      return null;
+    public T variant(Types.VariantType variant) {
+      throw new UnsupportedOperationException("Unsupported type: variant");
     }
 
     public T primitive(Type.PrimitiveType primitive) {
@@ -793,7 +793,7 @@ public static <T> T visit(Type type, CustomOrderSchemaVisitor<T> visitor) {
             new VisitFuture<>(map.valueType(), visitor));
 
       case VARIANT:
-        return visitor.variant();
+        return visitor.variant(type.asVariantType());
 
       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
@@ -85,10 +85,10 @@ public static Type fromTypeName(String typeString) {
   public static PrimitiveType fromPrimitiveString(String typeString) {
     Type type = fromTypeName(typeString);
     if (type.isPrimitiveType()) {
-      return (PrimitiveType) type;
+      return type.asPrimitiveType();
     }
 
-    throw new IllegalArgumentException("Cannot parse type string to primitive: " + typeString);
+    throw new IllegalArgumentException("Cannot parse type string: variant is not a primitive type");
   }
 
   public static class BooleanType extends PrimitiveType {
@@ -445,6 +445,11 @@ public boolean isVariantType() {
       return true;
     }
 
+    @Override
+    public VariantType asVariantType() {
+      return this;
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) {

diff --git a/api/src/test/java/org/apache/iceberg/types/TestTypeUtil.java b/api/src/test/java/org/apache/iceberg/types/TestTypeUtil.java
@@ -665,11 +665,8 @@ private static Stream<Arguments> testTypes() {
   public void testAssignFreshIdsWithType(Type testType) {
     Schema schema =
         new Schema(required(0, "v", testType), required(1, "A", Types.IntegerType.get()));
-    Schema sourceSchema =
-        new Schema(required(1, "v", testType), required(2, "A", Types.IntegerType.get()));
 
-    Schema assignedSchema =
-        TypeUtil.assignFreshIds(sourceSchema, new AtomicInteger(10)::incrementAndGet);
+    Schema assignedSchema = TypeUtil.assignFreshIds(schema, new AtomicInteger(10)::incrementAndGet);
     Schema expectedSchema =
         new Schema(required(11, "v", testType), required(12, "A", Types.IntegerType.get()));
     assertThat(assignedSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
@@ -683,7 +680,7 @@ public void testReassignIdsWithType(Type testType) {
     Schema sourceSchema =
         new Schema(required(1, "v", testType), required(2, "A", Types.IntegerType.get()));
 
-    final Schema reassignedSchema = TypeUtil.reassignIds(schema, sourceSchema);
+    Schema reassignedSchema = TypeUtil.reassignIds(schema, sourceSchema);
     assertThat(reassignedSchema.asStruct()).isEqualTo(sourceSchema.asStruct());
   }
 
@@ -692,36 +689,32 @@ public void testReassignIdsWithType(Type testType) {
   public void testIndexByIdWithType(Type testType) {
     Schema schema =
         new Schema(required(0, "v", testType), required(1, "A", Types.IntegerType.get()));
-    Schema sourceSchema =
-        new Schema(required(1, "v", testType), required(2, "A", Types.IntegerType.get()));
 
-    Map<Integer, Types.NestedField> indexByIds = TypeUtil.indexById(sourceSchema.asStruct());
-    assertThat(indexByIds.get(1).type()).isEqualTo(testType);
+    Map<Integer, Types.NestedField> indexByIds = TypeUtil.indexById(schema.asStruct());
+    assertThat(indexByIds.get(0).type()).isEqualTo(testType);
   }
 
   @ParameterizedTest
   @MethodSource("testTypes")
   public void testIndexNameByIdWithType(Type testType) {
     Schema schema =
         new Schema(required(0, "v", testType), required(1, "A", Types.IntegerType.get()));
-    Schema sourceSchema =
-        new Schema(required(1, "v", testType), required(2, "A", Types.IntegerType.get()));
 
-    Map<Integer, String> indexNameByIds = TypeUtil.indexNameById(sourceSchema.asStruct());
-    assertThat(indexNameByIds.get(1)).isEqualTo("v");
+    Map<Integer, String> indexNameByIds = TypeUtil.indexNameById(schema.asStruct());
+    assertThat(indexNameByIds.get(0)).isEqualTo("v");
   }
 
   @ParameterizedTest
   @MethodSource("testTypes")
   public void testProjectWithType(Type testType) {
-    Schema sourceSchema =
-        new Schema(required(1, "v", testType), required(2, "A", Types.IntegerType.get()));
+    Schema schema =
+        new Schema(required(0, "v", testType), required(1, "A", Types.IntegerType.get()));
 
-    Schema expectedSchema = new Schema(Lists.newArrayList(required(1, "v", testType)));
-    Schema projectedSchema = TypeUtil.project(sourceSchema, Sets.newHashSet(1));
+    Schema expectedSchema = new Schema(Lists.newArrayList(required(0, "v", testType)));
+    Schema projectedSchema = TypeUtil.project(schema, Sets.newHashSet(0));
     assertThat(projectedSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
 
-    Set<Integer> projectedIds = TypeUtil.getProjectedIds(sourceSchema);
-    assertThat(Set.of(1, 2)).isEqualTo(projectedIds);
+    Set<Integer> projectedIds = TypeUtil.getProjectedIds(schema);
+    assertThat(Set.of(0, 1)).isEqualTo(projectedIds);
   }
 }
diff --git a/api/src/test/java/org/apache/iceberg/types/TestTypes.java b/api/src/test/java/org/apache/iceberg/types/TestTypes.java
@@ -69,10 +69,10 @@ public void fromPrimitiveString() {
 
     assertThatExceptionOfType(IllegalArgumentException.class)
         .isThrownBy(() -> Types.fromPrimitiveString("variant"))
-        .withMessage("Cannot parse type string to primitive: variant");
+        .withMessage("Cannot parse type string: variant is not a primitive type");
     assertThatExceptionOfType(IllegalArgumentException.class)
         .isThrownBy(() -> Types.fromPrimitiveString("Variant"))
-        .withMessage("Cannot parse type string to primitive: Variant");
+        .withMessage("Cannot parse type string: variant is not a primitive type");
 
     assertThatExceptionOfType(IllegalArgumentException.class)
         .isThrownBy(() -> Types.fromPrimitiveString("abcdefghij"))

diff --git a/core/src/main/java/org/apache/iceberg/avro/TypeToSchema.java b/core/src/main/java/org/apache/iceberg/avro/TypeToSchema.java
@@ -188,7 +188,7 @@ public Schema map(Types.MapType map, Schema keySchema, Schema valueSchema) {
   }
 
   @Override
-  public Schema variant() {
+  public Schema variant(Types.VariantType variant) {
     String recordName = "r" + fieldIds.peek();
     return Schema.createRecord(
         recordName,

diff --git a/core/src/main/java/org/apache/iceberg/schema/SchemaWithPartnerVisitor.java b/core/src/main/java/org/apache/iceberg/schema/SchemaWithPartnerVisitor.java
@@ -107,7 +107,7 @@ public static <P, T> T visit(
 
         return visitor.map(map, partner, keyResult, valueResult);
       case VARIANT:
-        return visitor.variant(partner);
+        return visitor.variant(type.asVariantType(), partner);
       default:
         return visitor.primitive(type.asPrimitiveType(), partner);
     }
@@ -161,7 +161,7 @@ public R map(Types.MapType map, P partner, R keyResult, R valueResult) {
     return null;
   }
 
-  public R variant(P partner) {
+  public R variant(Types.VariantType variant, P partner) {
     throw new UnsupportedOperationException("Unsupported type: variant");
   }
 

diff --git a/core/src/main/java/org/apache/iceberg/schema/UnionByNameVisitor.java b/core/src/main/java/org/apache/iceberg/schema/UnionByNameVisitor.java
@@ -143,7 +143,7 @@ public Boolean map(
   }
 
   @Override
-  public Boolean variant(Integer partnerId) {
+  public Boolean variant(Types.VariantType variant, Integer partnerId) {
     return partnerId == null;
   }