From 235a3eff6f57b6a6400079364e0fbc10c20c001f Mon Sep 17 00:00:00 2001 From: Kadin Rabo Date: Fri, 3 Apr 2026 11:55:56 -0400 Subject: [PATCH 1/2] feat(core): add typeVariationReference to UserDefined type The Substrait proto defines type_variation_reference on UserDefined but substrait-java ignores it. Add the field to the POJO, thread it through both proto converters, and add builder/DSL overloads. Closes #567 --- .../io/substrait/dsl/SubstraitBuilder.java | 17 ++++++++++ .../src/main/java/io/substrait/type/Type.java | 15 ++++++++- .../java/io/substrait/type/TypeCreator.java | 9 ++++++ .../type/proto/BaseProtoConverter.java | 3 +- .../substrait/type/proto/BaseProtoTypes.java | 6 ++-- .../type/proto/ProtoTypeConverter.java | 1 + .../type/proto/TypeProtoConverter.java | 13 ++++++-- .../extension/TypeExtensionTest.java | 32 +++++++++++++++++++ 8 files changed, 89 insertions(+), 7 deletions(-) diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 0f58da6af..2bb8cf6ce 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -1510,6 +1510,23 @@ public Type.UserDefined userDefinedType(String urn, String typeName) { return Type.UserDefined.builder().urn(urn).name(typeName).nullable(false).build(); } + /** + * Creates a user-defined type with the specified URN, type name, and type variation. + * + * @param urn the URN of the extension containing the type + * @param typeName the name of the user-defined type + * @param typeVariationReference the type variation reference + * @return a new {@link Type.UserDefined} + */ + public Type.UserDefined userDefinedType(String urn, String typeName, int typeVariationReference) { + return Type.UserDefined.builder() + .urn(urn) + .name(typeName) + .nullable(false) + .typeVariationReference(typeVariationReference) + .build(); + } + // Misc /** diff --git a/core/src/main/java/io/substrait/type/Type.java b/core/src/main/java/io/substrait/type/Type.java index e82f8c1cf..3821cb085 100644 --- a/core/src/main/java/io/substrait/type/Type.java +++ b/core/src/main/java/io/substrait/type/Type.java @@ -441,8 +441,21 @@ public java.util.List typeParameters() { return java.util.Collections.emptyList(); } + /** + * Returns the type variation reference for this user-defined type. + * + *

Type variations allow different physical representations or semantics for the same logical + * type. The reference value maps to an {@code ExtensionTypeVariation} declaration in the plan. + * + * @return the type variation reference, or {@code 0} if using the default variation + */ + // Note: Cannot use @Value.Default here due to an Immutables bug with @Value.Enclosing + // and multiple @Value.Default fields in the same class (generates broken InitShim code). + // The default value (0) is set in the builder() factory method below. + public abstract int typeVariationReference(); + public static ImmutableType.UserDefined.Builder builder() { - return ImmutableType.UserDefined.builder(); + return ImmutableType.UserDefined.builder().typeVariationReference(0); } @Override diff --git a/core/src/main/java/io/substrait/type/TypeCreator.java b/core/src/main/java/io/substrait/type/TypeCreator.java index 6a897417e..0331c9e76 100644 --- a/core/src/main/java/io/substrait/type/TypeCreator.java +++ b/core/src/main/java/io/substrait/type/TypeCreator.java @@ -115,6 +115,15 @@ public Type userDefined(String urn, String name) { return Type.UserDefined.builder().nullable(nullable).urn(urn).name(name).build(); } + public Type userDefined(String urn, String name, int typeVariationReference) { + return Type.UserDefined.builder() + .nullable(nullable) + .urn(urn) + .name(name) + .typeVariationReference(typeVariationReference) + .build(); + } + public static TypeCreator of(boolean nullability) { return nullability ? NULLABLE : REQUIRED; } diff --git a/core/src/main/java/io/substrait/type/proto/BaseProtoConverter.java b/core/src/main/java/io/substrait/type/proto/BaseProtoConverter.java index 3909d4033..898fe87d0 100644 --- a/core/src/main/java/io/substrait/type/proto/BaseProtoConverter.java +++ b/core/src/main/java/io/substrait/type/proto/BaseProtoConverter.java @@ -175,6 +175,7 @@ public final T visit(final Type.Map expr) { public final T visit(final Type.UserDefined expr) { int ref = extensionCollector.getTypeReference(SimpleExtension.TypeAnchor.of(expr.urn(), expr.name())); - return typeContainer(expr).userDefined(ref, expr.typeParameters()); + return typeContainer(expr) + .userDefined(ref, expr.typeVariationReference(), expr.typeParameters()); } } diff --git a/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java b/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java index 47842382f..dfbd1fca5 100644 --- a/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java +++ b/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java @@ -131,10 +131,12 @@ public final T struct(T... types) { public abstract T map(T key, T value); - public abstract T userDefined(int ref); + public abstract T userDefined(int ref, int typeVariationReference); public abstract T userDefined( - int ref, java.util.List typeParameters); + int ref, + int typeVariationReference, + java.util.List typeParameters); protected abstract T wrap(Object o); diff --git a/core/src/main/java/io/substrait/type/proto/ProtoTypeConverter.java b/core/src/main/java/io/substrait/type/proto/ProtoTypeConverter.java index 24231aebc..00b2532b7 100644 --- a/core/src/main/java/io/substrait/type/proto/ProtoTypeConverter.java +++ b/core/src/main/java/io/substrait/type/proto/ProtoTypeConverter.java @@ -107,6 +107,7 @@ public Type from(io.substrait.proto.Type type) { userDefined.getTypeParametersList().stream() .map(this::from) .collect(java.util.stream.Collectors.toList())) + .typeVariationReference(userDefined.getTypeVariationReference()) .build(); } case USER_DEFINED_TYPE_REFERENCE: diff --git a/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java b/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java index 102c295a2..88035ada1 100644 --- a/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java +++ b/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java @@ -181,17 +181,24 @@ public Type map(Type key, Type value) { } @Override - public Type userDefined(int ref) { + public Type userDefined(int ref, int typeVariationReference) { return wrap( - Type.UserDefined.newBuilder().setTypeReference(ref).setNullability(nullability).build()); + Type.UserDefined.newBuilder() + .setTypeReference(ref) + .setTypeVariationReference(typeVariationReference) + .setNullability(nullability) + .build()); } @Override public Type userDefined( - int ref, java.util.List typeParameters) { + int ref, + int typeVariationReference, + java.util.List typeParameters) { return wrap( Type.UserDefined.newBuilder() .setTypeReference(ref) + .setTypeVariationReference(typeVariationReference) .setNullability(nullability) .addAllTypeParameters( typeParameters.stream() diff --git a/core/src/test/java/io/substrait/extension/TypeExtensionTest.java b/core/src/test/java/io/substrait/extension/TypeExtensionTest.java index 3a3c2ec63..efefb2870 100644 --- a/core/src/test/java/io/substrait/extension/TypeExtensionTest.java +++ b/core/src/test/java/io/substrait/extension/TypeExtensionTest.java @@ -78,6 +78,38 @@ void roundtripCustomType() { io.substrait.proto.Plan protoPlan = planProtoConverter.toProto(plan); Plan planReturned = protoPlanConverter.from(protoPlan); assertEquals(plan, planReturned); + + // verify default typeVariationReference is 0 + assertEquals(0, ((Type.UserDefined) customType1).typeVariationReference()); + } + + @Test + void roundtripCustomTypeWithVariationReference() { + Type customTypeWithVariation = sb.userDefinedType(URN, "customType1", 42); + + List tableName = Stream.of("example").collect(Collectors.toList()); + List columnNames = Stream.of("custom_type_column").collect(Collectors.toList()); + List types = Stream.of(customTypeWithVariation).collect(Collectors.toList()); + + Plan plan = sb.plan(sb.root(sb.namedScan(tableName, columnNames, types))); + + io.substrait.proto.Plan protoPlan = planProtoConverter.toProto(plan); + + // verify the type variation reference is set in the proto + io.substrait.proto.Type protoType = + protoPlan + .getRelations(0) + .getRoot() + .getInput() + .getRead() + .getBaseSchema() + .getStruct() + .getTypes(0); + assertEquals(42, protoType.getUserDefined().getTypeVariationReference()); + + // verify full roundtrip + Plan planReturned = protoPlanConverter.from(protoPlan); + assertEquals(plan, planReturned); } @Test From daa5f32969e2845ded4940b3225745ea9ef8d73c Mon Sep 17 00:00:00 2001 From: Kadin Rabo Date: Wed, 8 Apr 2026 10:15:17 -0400 Subject: [PATCH 2/2] address review: UserDefined as interface, simplify overloads - Change UserDefined from abstract class to interface, enabling clean @Value.Default for typeVariationReference (fixes InitShim bug) - Delegate 2-arg overloads to 3-arg in TypeCreator and SubstraitBuilder - Make BaseProtoTypes.userDefined(ref, tvr) a concrete default method - Remove duplicate 2-arg override in TypeProtoConverter - Simplify test per reviewer feedback --- .../io/substrait/dsl/SubstraitBuilder.java | 2 +- .../src/main/java/io/substrait/type/Type.java | 22 +++++++++---------- .../java/io/substrait/type/TypeCreator.java | 2 +- .../substrait/type/proto/BaseProtoTypes.java | 4 +++- .../type/proto/TypeProtoConverter.java | 10 --------- .../extension/TypeExtensionTest.java | 14 ------------ 6 files changed, 16 insertions(+), 38 deletions(-) diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 2bb8cf6ce..54d4f63a3 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -1507,7 +1507,7 @@ public Expression.WindowFunctionInvocation windowFn( * @return a new {@link Type.UserDefined} */ public Type.UserDefined userDefinedType(String urn, String typeName) { - return Type.UserDefined.builder().urn(urn).name(typeName).nullable(false).build(); + return userDefinedType(urn, typeName, 0); } /** diff --git a/core/src/main/java/io/substrait/type/Type.java b/core/src/main/java/io/substrait/type/Type.java index 3821cb085..801e83880 100644 --- a/core/src/main/java/io/substrait/type/Type.java +++ b/core/src/main/java/io/substrait/type/Type.java @@ -422,11 +422,11 @@ public R accept(final TypeVisitor typeVisitor) th } @Value.Immutable - abstract class UserDefined implements Type { + interface UserDefined extends Type { - public abstract String urn(); + String urn(); - public abstract String name(); + String name(); /** * Returns the type parameters for this user-defined type. @@ -437,7 +437,7 @@ abstract class UserDefined implements Type { * @return a list of type parameters, or an empty list if this type is not parameterized */ @Value.Default - public java.util.List typeParameters() { + default java.util.List typeParameters() { return java.util.Collections.emptyList(); } @@ -449,17 +449,17 @@ public java.util.List typeParameters() { * * @return the type variation reference, or {@code 0} if using the default variation */ - // Note: Cannot use @Value.Default here due to an Immutables bug with @Value.Enclosing - // and multiple @Value.Default fields in the same class (generates broken InitShim code). - // The default value (0) is set in the builder() factory method below. - public abstract int typeVariationReference(); + @Value.Default + default int typeVariationReference() { + return 0; + } - public static ImmutableType.UserDefined.Builder builder() { - return ImmutableType.UserDefined.builder().typeVariationReference(0); + static ImmutableType.UserDefined.Builder builder() { + return ImmutableType.UserDefined.builder(); } @Override - public R accept(TypeVisitor typeVisitor) throws E { + default R accept(TypeVisitor typeVisitor) throws E { return typeVisitor.visit(this); } } diff --git a/core/src/main/java/io/substrait/type/TypeCreator.java b/core/src/main/java/io/substrait/type/TypeCreator.java index 0331c9e76..0362223bd 100644 --- a/core/src/main/java/io/substrait/type/TypeCreator.java +++ b/core/src/main/java/io/substrait/type/TypeCreator.java @@ -112,7 +112,7 @@ public Type.Map map(Type key, Type value) { } public Type userDefined(String urn, String name) { - return Type.UserDefined.builder().nullable(nullable).urn(urn).name(name).build(); + return userDefined(urn, name, 0); } public Type userDefined(String urn, String name, int typeVariationReference) { diff --git a/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java b/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java index dfbd1fca5..c6490c9b1 100644 --- a/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java +++ b/core/src/main/java/io/substrait/type/proto/BaseProtoTypes.java @@ -131,7 +131,9 @@ public final T struct(T... types) { public abstract T map(T key, T value); - public abstract T userDefined(int ref, int typeVariationReference); + public T userDefined(int ref, int typeVariationReference) { + return userDefined(ref, typeVariationReference, java.util.Collections.emptyList()); + } public abstract T userDefined( int ref, diff --git a/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java b/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java index 88035ada1..aa7b4d6a6 100644 --- a/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java +++ b/core/src/main/java/io/substrait/type/proto/TypeProtoConverter.java @@ -180,16 +180,6 @@ public Type map(Type key, Type value) { Type.Map.newBuilder().setKey(key).setValue(value).setNullability(nullability).build()); } - @Override - public Type userDefined(int ref, int typeVariationReference) { - return wrap( - Type.UserDefined.newBuilder() - .setTypeReference(ref) - .setTypeVariationReference(typeVariationReference) - .setNullability(nullability) - .build()); - } - @Override public Type userDefined( int ref, diff --git a/core/src/test/java/io/substrait/extension/TypeExtensionTest.java b/core/src/test/java/io/substrait/extension/TypeExtensionTest.java index efefb2870..c5d709a26 100644 --- a/core/src/test/java/io/substrait/extension/TypeExtensionTest.java +++ b/core/src/test/java/io/substrait/extension/TypeExtensionTest.java @@ -94,20 +94,6 @@ void roundtripCustomTypeWithVariationReference() { Plan plan = sb.plan(sb.root(sb.namedScan(tableName, columnNames, types))); io.substrait.proto.Plan protoPlan = planProtoConverter.toProto(plan); - - // verify the type variation reference is set in the proto - io.substrait.proto.Type protoType = - protoPlan - .getRelations(0) - .getRoot() - .getInput() - .getRead() - .getBaseSchema() - .getStruct() - .getTypes(0); - assertEquals(42, protoType.getUserDefined().getTypeVariationReference()); - - // verify full roundtrip Plan planReturned = protoPlanConverter.from(protoPlan); assertEquals(plan, planReturned); }