make FieldValueTypeInformation creators take a TypeDescriptor parameter

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/AutoValueSchema.java

@@ -62,7 +62,7 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor) {
               .collect(Collectors.toList());
       List<FieldValueTypeInformation> types = Lists.newArrayListWithCapacity(methods.size());
       for (int i = 0; i < methods.size(); ++i) {
-        types.add(FieldValueTypeInformation.forGetter(methods.get(i), i));
+        types.add(FieldValueTypeInformation.forGetter(typeDescriptor, methods.get(i), i));
       }
       types.sort(Comparator.comparing(FieldValueTypeInformation::getNumber));
       validateFieldNumbers(types);

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/FieldValueTypeInformation.java

@@ -27,6 +27,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Stream;
 import org.apache.beam.sdk.schemas.annotations.SchemaCaseFormat;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
@@ -126,17 +127,26 @@ public static FieldValueTypeInformation forOneOf(
   }
 
   public static FieldValueTypeInformation forField(Field field, int index) {
-    TypeDescriptor<?> type = TypeDescriptor.of(field.getGenericType());
+    return forField(null, field, index);
+  }
+
+  public static FieldValueTypeInformation forField(
+      @Nullable TypeDescriptor<?> typeDescriptor, Field field, int index) {
+    TypeDescriptor<?> type =
+        Optional.ofNullable(typeDescriptor)
+            .map(td -> (TypeDescriptor) td.resolveType(field.getGenericType()))
+            // fall back to previous behavior
+            .orElseGet(() -> TypeDescriptor.of(field.getGenericType()));
     return new AutoValue_FieldValueTypeInformation.Builder()
         .setName(getNameOverride(field.getName(), field))
         .setNumber(getNumberOverride(index, field))
         .setNullable(hasNullableAnnotation(field))
         .setType(type)
         .setRawType(type.getRawType())
         .setField(field)
-        .setElementType(getIterableComponentType(field))
-        .setMapKeyType(getMapKeyType(field))
-        .setMapValueType(getMapValueType(field))
+        .setElementType(getIterableComponentType(type))
+        .setMapKeyType(getMapKeyType(type))
+        .setMapValueType(getMapValueType(type))
         .setOneOfTypes(Collections.emptyMap())
         .setDescription(getFieldDescription(field))
         .build();
@@ -185,6 +195,11 @@ public static <T extends AnnotatedElement & Member> String getNameOverride(
   }
 
   public static FieldValueTypeInformation forGetter(Method method, int index) {
+    return forGetter(null, method, index);
+  }
+
+  public static FieldValueTypeInformation forGetter(
+      @Nullable TypeDescriptor<?> typeDescriptor, Method method, int index) {
     String name;
     if (method.getName().startsWith("get")) {
       name = ReflectUtils.stripPrefix(method.getName(), "get");
@@ -194,7 +209,12 @@ public static FieldValueTypeInformation forGetter(Method method, int index) {
       throw new RuntimeException("Getter has wrong prefix " + method.getName());
     }
 
-    TypeDescriptor<?> type = TypeDescriptor.of(method.getGenericReturnType());
+    TypeDescriptor<?> type =
+        Optional.ofNullable(typeDescriptor)
+            .map(td -> (TypeDescriptor) td.resolveType(method.getGenericReturnType()))
+            // fall back to previous behavior
+            .orElseGet(() -> TypeDescriptor.of(method.getGenericReturnType()));
+
     boolean nullable = hasNullableReturnType(method);
     return new AutoValue_FieldValueTypeInformation.Builder()
         .setName(getNameOverride(name, method))
@@ -253,18 +273,32 @@ private static boolean isNullableAnnotation(Annotation annotation) {
   }
 
   public static FieldValueTypeInformation forSetter(Method method) {
-    return forSetter(method, "set");
+    return forSetter(null, method);
   }
 
   public static FieldValueTypeInformation forSetter(Method method, String setterPrefix) {
+    return forSetter(null, method, setterPrefix);
+  }
+
+  public static FieldValueTypeInformation forSetter(
+      @Nullable TypeDescriptor<?> typeDescriptor, Method method) {
+    return forSetter(typeDescriptor, method, "set");
+  }
+
+  public static FieldValueTypeInformation forSetter(
+      @Nullable TypeDescriptor<?> typeDescriptor, Method method, String setterPrefix) {
     String name;
     if (method.getName().startsWith(setterPrefix)) {
       name = ReflectUtils.stripPrefix(method.getName(), setterPrefix);
     } else {
       throw new RuntimeException("Setter has wrong prefix " + method.getName());
     }
 
-    TypeDescriptor<?> type = TypeDescriptor.of(method.getGenericParameterTypes()[0]);
+    TypeDescriptor<?> type =
+        Optional.ofNullable(typeDescriptor)
+            .map(td -> (TypeDescriptor) td.resolveType(method.getGenericParameterTypes()[0]))
+            // fall back to previous behavior
+            .orElseGet(() -> TypeDescriptor.of(method.getGenericParameterTypes()[0]));
     boolean nullable = hasSingleNullableParameter(method);
     return new AutoValue_FieldValueTypeInformation.Builder()
         .setName(name)
@@ -283,10 +317,6 @@ public FieldValueTypeInformation withName(String name) {
     return toBuilder().setName(name).build();
   }
 
-  private static FieldValueTypeInformation getIterableComponentType(Field field) {
-    return getIterableComponentType(TypeDescriptor.of(field.getGenericType()));
-  }
-
   static @Nullable FieldValueTypeInformation getIterableComponentType(TypeDescriptor<?> valueType) {
     // TODO: Figure out nullable elements.
     TypeDescriptor<?> componentType = ReflectUtils.getIterableComponentType(valueType);
@@ -306,23 +336,13 @@ private static FieldValueTypeInformation getIterableComponentType(Field field) {
         .build();
   }
 
-  // If the Field is a map type, returns the key type, otherwise returns a null reference.
-
-  private static @Nullable FieldValueTypeInformation getMapKeyType(Field field) {
-    return getMapKeyType(TypeDescriptor.of(field.getGenericType()));
-  }
-
+  // If the type is a map type, returns the key type, otherwise returns a null reference.
   private static @Nullable FieldValueTypeInformation getMapKeyType(
       TypeDescriptor<?> typeDescriptor) {
     return getMapType(typeDescriptor, 0);
   }
 
-  // If the Field is a map type, returns the value type, otherwise returns a null reference.
-
-  private static @Nullable FieldValueTypeInformation getMapValueType(Field field) {
-    return getMapType(TypeDescriptor.of(field.getGenericType()), 1);
-  }
-
+  // If the type is a map type, returns the value type, otherwise returns a null reference.
   private static @Nullable FieldValueTypeInformation getMapValueType(
       TypeDescriptor<?> typeDescriptor) {
     return getMapType(typeDescriptor, 1);

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/JavaBeanSchema.java

@@ -68,7 +68,7 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor) {
               .collect(Collectors.toList());
       List<FieldValueTypeInformation> types = Lists.newArrayListWithCapacity(methods.size());
       for (int i = 0; i < methods.size(); ++i) {
-        types.add(FieldValueTypeInformation.forGetter(methods.get(i), i));
+        types.add(FieldValueTypeInformation.forGetter(typeDescriptor, methods.get(i), i));
       }
       types.sort(Comparator.comparing(FieldValueTypeInformation::getNumber));
       validateFieldNumbers(types);
@@ -114,7 +114,7 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor) {
       return ReflectUtils.getMethods(typeDescriptor.getRawType()).stream()
           .filter(ReflectUtils::isSetter)
           .filter(m -> !m.isAnnotationPresent(SchemaIgnore.class))
-          .map(FieldValueTypeInformation::forSetter)
+          .map(m -> FieldValueTypeInformation.forSetter(typeDescriptor, m))
           .map(
               t -> {
                 if (t.getMethod().getAnnotation(SchemaFieldNumber.class) != null) {

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/JavaFieldSchema.java

@@ -64,7 +64,7 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor) {
               .collect(Collectors.toList());
       List<FieldValueTypeInformation> types = Lists.newArrayListWithCapacity(fields.size());
       for (int i = 0; i < fields.size(); ++i) {
-        types.add(FieldValueTypeInformation.forField(fields.get(i), i));
+        types.add(FieldValueTypeInformation.forField(typeDescriptor, fields.get(i), i));
       }
       types.sort(Comparator.comparing(FieldValueTypeInformation::getNumber));
       validateFieldNumbers(types);

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/AutoValueUtils.java

@@ -199,7 +199,7 @@ private static boolean matchConstructor(
     Map<String, FieldValueTypeInformation> setterTypes =
         ReflectUtils.getMethods(builderClass).stream()
             .filter(ReflectUtils::isSetter)
-            .map(FieldValueTypeInformation::forSetter)
+            .map(m -> FieldValueTypeInformation.forSetter(TypeDescriptor.of(builderClass), m))
             .collect(Collectors.toMap(FieldValueTypeInformation::getName, Function.identity()));
 
     List<FieldValueTypeInformation> setterMethods =

sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/FieldValueTypeInformationTest.java

@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.schemas;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.Map;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.junit.Test;
+
+public class FieldValueTypeInformationTest {
+  public static class GenericClass<T> {
+    public T t;
+
+    public GenericClass(T t) {
+      this.t = t;
+    }
+
+    public T getT() {
+      return t;
+    }
+
+    public void setT(T t) {
+      this.t = t;
+    }
+  }
+
+  private final TypeDescriptor<GenericClass<Map<String, Integer>>> typeDescriptor =
+      new TypeDescriptor<GenericClass<Map<String, Integer>>>() {};
+  private final TypeDescriptor<Map<String, Integer>> expectedFieldTypeDescriptor =
+      new TypeDescriptor<Map<String, Integer>>() {};
+
+  @Test
+  public void testForGetter() throws Exception {
+    FieldValueTypeInformation actual =
+        FieldValueTypeInformation.forGetter(
+            typeDescriptor, GenericClass.class.getMethod("getT"), 0);
+    assertEquals(expectedFieldTypeDescriptor, actual.getType());
+  }
+
+  @Test
+  public void testForField() throws Exception {
+    FieldValueTypeInformation actual =
+        FieldValueTypeInformation.forField(typeDescriptor, GenericClass.class.getField("t"), 0);
+    assertEquals(expectedFieldTypeDescriptor, actual.getType());
+  }
+
+  @Test
+  public void testForSetter() throws Exception {
+    FieldValueTypeInformation actual =
+        FieldValueTypeInformation.forSetter(
+            typeDescriptor, GenericClass.class.getMethod("setT", Object.class));
+    assertEquals(expectedFieldTypeDescriptor, actual.getType());
+  }
+}

sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/schemas/utils/AvroUtils.java

@@ -821,7 +821,7 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor, Sch
         Method method = methods.get(i);
         if (ReflectUtils.isGetter(method)) {
           FieldValueTypeInformation fieldValueTypeInformation =
-              FieldValueTypeInformation.forGetter(method, i);
+              FieldValueTypeInformation.forGetter(typeDescriptor, method, i);
           String name = mapping.get(fieldValueTypeInformation.getName());
           if (name != null) {
             types.add(fieldValueTypeInformation.withName(name));
@@ -871,7 +871,8 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor) {
       for (int i = 0; i < classFields.size(); ++i) {
         java.lang.reflect.Field f = classFields.get(i);
         if (!f.isAnnotationPresent(AvroIgnore.class)) {
-          FieldValueTypeInformation typeInformation = FieldValueTypeInformation.forField(f, i);
+          FieldValueTypeInformation typeInformation =
+              FieldValueTypeInformation.forField(typeDescriptor, f, i);
           AvroName avroname = f.getAnnotation(AvroName.class);
           if (avroname != null) {
             typeInformation = typeInformation.withName(avroname.value());

sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoByteBuddyUtils.java

@@ -1018,34 +1018,41 @@ static Method getProtoGetter(Multimap<String, Method> methods, String name, Fiel
 
   public static @Nullable <ProtoBuilderT extends MessageLite.Builder>
       SchemaUserTypeCreator getBuilderCreator(
-          Class<?> protoClass, Schema schema, FieldValueTypeSupplier fieldValueTypeSupplier) {
-    Class<ProtoBuilderT> builderClass = getProtoGeneratedBuilder(protoClass);
+          TypeDescriptor<?> protoTypeDescriptor,
+          Schema schema,
+          FieldValueTypeSupplier fieldValueTypeSupplier) {
+    Class<ProtoBuilderT> builderClass = getProtoGeneratedBuilder(protoTypeDescriptor.getRawType());
     if (builderClass == null) {
       return null;
     }
     Multimap<String, Method> methods = ReflectUtils.getMethodsMap(builderClass);
     List<FieldValueSetter<ProtoBuilderT, Object>> setters =
         schema.getFields().stream()
-            .map(f -> getProtoFieldValueSetter(f, methods, builderClass))
+            .map(f -> getProtoFieldValueSetter(protoTypeDescriptor, f, methods, builderClass))
             .collect(Collectors.toList());
-    return createBuilderCreator(protoClass, builderClass, setters, schema);
+    return createBuilderCreator(protoTypeDescriptor.getRawType(), builderClass, setters, schema);
   }
 
   private static <ProtoBuilderT extends MessageLite.Builder>
       FieldValueSetter<ProtoBuilderT, Object> getProtoFieldValueSetter(
-          Field field, Multimap<String, Method> methods, Class<ProtoBuilderT> builderClass) {
+          TypeDescriptor<?> typeDescriptor,
+          Field field,
+          Multimap<String, Method> methods,
+          Class<ProtoBuilderT> builderClass) {
     if (field.getType().isLogicalType(OneOfType.IDENTIFIER)) {
       OneOfType oneOfType = field.getType().getLogicalType(OneOfType.class);
       TreeMap<Integer, FieldValueSetter<ProtoBuilderT, Object>> oneOfSetters = Maps.newTreeMap();
       for (Field oneOfField : oneOfType.getOneOfSchema().getFields()) {
-        FieldValueSetter setter = getProtoFieldValueSetter(oneOfField, methods, builderClass);
+        FieldValueSetter setter =
+            getProtoFieldValueSetter(typeDescriptor, oneOfField, methods, builderClass);
         oneOfSetters.put(getFieldNumber(oneOfField), setter);
       }
       return createOneOfSetter(field.getName(), oneOfSetters, builderClass);
     } else {
       Method method = getProtoSetter(methods, field.getName(), field.getType());
       return JavaBeanUtils.createSetter(
-          FieldValueTypeInformation.forSetter(method, protoSetterPrefix(field.getType())),
+          FieldValueTypeInformation.forSetter(
+              typeDescriptor, method, protoSetterPrefix(field.getType())),
           new ProtoTypeConversionsFactory());
     }
   }

sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoMessageSchema.java

@@ -72,7 +72,8 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor, Sch
             Method method = getProtoGetter(methods, oneOfField.getName(), oneOfField.getType());
             oneOfTypes.put(
                 oneOfField.getName(),
-                FieldValueTypeInformation.forGetter(method, i).withName(field.getName()));
+                FieldValueTypeInformation.forGetter(typeDescriptor, method, i)
+                    .withName(field.getName()));
           }
           // Add an entry that encapsulates information about all possible getters.
           types.add(
@@ -82,7 +83,9 @@ public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor, Sch
         } else {
           // This is a simple field. Add the getter.
           Method method = getProtoGetter(methods, field.getName(), field.getType());
-          types.add(FieldValueTypeInformation.forGetter(method, i).withName(field.getName()));
+          types.add(
+              FieldValueTypeInformation.forGetter(typeDescriptor, method, i)
+                  .withName(field.getName()));
         }
       }
       return types;
@@ -117,7 +120,7 @@ public SchemaUserTypeCreator schemaTypeCreator(
       TypeDescriptor<?> targetTypeDescriptor, Schema schema) {
     SchemaUserTypeCreator creator =
         ProtoByteBuddyUtils.getBuilderCreator(
-            targetTypeDescriptor.getRawType(), schema, new ProtoClassFieldValueTypeSupplier());
+            targetTypeDescriptor, schema, new ProtoClassFieldValueTypeSupplier());
     if (creator == null) {
       throw new RuntimeException("Cannot create creator for " + targetTypeDescriptor);
     }

sdks/java/io/thrift/src/main/java/org/apache/beam/sdk/io/thrift/ThriftSchema.java

@@ -242,10 +242,12 @@ private FieldValueTypeInformation fieldValueTypeInfo(Class<?> type, String field
       if (factoryMethods.size() > 1) {
         throw new IllegalStateException("Overloaded factory methods: " + factoryMethods);
       }
-      return FieldValueTypeInformation.forSetter(factoryMethods.get(0), "");
+      return FieldValueTypeInformation.forSetter(
+          TypeDescriptor.of(type), factoryMethods.get(0), "");
     } else {
       try {
-        return FieldValueTypeInformation.forField(type.getDeclaredField(fieldName), 0);
+        return FieldValueTypeInformation.forField(
+            TypeDescriptor.of(type), type.getDeclaredField(fieldName), 0);
       } catch (NoSuchFieldException e) {
         throw new IllegalArgumentException(e);
       }