Add support for schema providers for generic classes

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

@@ -46,13 +46,13 @@ public static class AbstractGetterTypeSupplier implements FieldValueTypeSupplier
     public static final AbstractGetterTypeSupplier INSTANCE = new AbstractGetterTypeSupplier();
 
     @Override
-    public List<FieldValueTypeInformation> get(Class<?> clazz) {
+    public List<FieldValueTypeInformation> get(TypeDescriptor<?> typeDescriptor) {
 
       // If the generated class is passed in, we want to look at the base class to find the getters.
-      Class<?> targetClass = AutoValueUtils.getBaseAutoValueClass(clazz);
+      TypeDescriptor<?> targetTypeDescriptor = AutoValueUtils.getBaseAutoValueClass(typeDescriptor);
 
       List<Method> methods =
-          ReflectUtils.getMethods(targetClass).stream()
+          ReflectUtils.getMethods(targetTypeDescriptor.getRawType()).stream()
               .filter(ReflectUtils::isGetter)
               // All AutoValue getters are marked abstract.
               .filter(m -> Modifier.isAbstract(m.getModifiers()))
@@ -62,7 +62,7 @@ public List<FieldValueTypeInformation> get(Class<?> clazz) {
               .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);
@@ -89,27 +89,30 @@ private static void validateFieldNumbers(List<FieldValueTypeInformation> types)
   }
 
   @Override
-  public List<FieldValueGetter> fieldValueGetters(Class<?> targetClass, Schema schema) {
+  public List<FieldValueGetter> fieldValueGetters(
+      TypeDescriptor<?> targetTypeDescriptor, Schema schema) {
     return JavaBeanUtils.getGetters(
-        targetClass,
+        targetTypeDescriptor,
         schema,
         AbstractGetterTypeSupplier.INSTANCE,
         new DefaultTypeConversionsFactory());
   }
 
   @Override
   public List<FieldValueTypeInformation> fieldValueTypeInformations(
-      Class<?> targetClass, Schema schema) {
-    return JavaBeanUtils.getFieldTypes(targetClass, schema, AbstractGetterTypeSupplier.INSTANCE);
+      TypeDescriptor<?> targetTypeDescriptor, Schema schema) {
+    return JavaBeanUtils.getFieldTypes(
+        targetTypeDescriptor, schema, AbstractGetterTypeSupplier.INSTANCE);
   }
 
   @Override
-  public SchemaUserTypeCreator schemaTypeCreator(Class<?> targetClass, Schema schema) {
+  public SchemaUserTypeCreator schemaTypeCreator(
+      TypeDescriptor<?> targetTypeDescriptor, Schema schema) {
     // If a static method is marked with @SchemaCreate, use that.
-    Method annotated = ReflectUtils.getAnnotatedCreateMethod(targetClass);
+    Method annotated = ReflectUtils.getAnnotatedCreateMethod(targetTypeDescriptor.getRawType());
     if (annotated != null) {
       return JavaBeanUtils.getStaticCreator(
-          targetClass,
+          targetTypeDescriptor,
           annotated,
           schema,
           AbstractGetterTypeSupplier.INSTANCE,
@@ -119,7 +122,8 @@ public SchemaUserTypeCreator schemaTypeCreator(Class<?> targetClass, Schema sche
     // Try to find a generated builder class. If one exists, use that to generate a
     // SchemaTypeCreator for creating AutoValue objects.
     SchemaUserTypeCreator creatorFactory =
-        AutoValueUtils.getBuilderCreator(targetClass, schema, AbstractGetterTypeSupplier.INSTANCE);
+        AutoValueUtils.getBuilderCreator(
+            targetTypeDescriptor, schema, AbstractGetterTypeSupplier.INSTANCE);
     if (creatorFactory != null) {
       return creatorFactory;
     }
@@ -128,9 +132,10 @@ public SchemaUserTypeCreator schemaTypeCreator(Class<?> targetClass, Schema sche
     // class. Use that for creating AutoValue objects.
     creatorFactory =
         AutoValueUtils.getConstructorCreator(
-            targetClass, schema, AbstractGetterTypeSupplier.INSTANCE);
+            targetTypeDescriptor, schema, AbstractGetterTypeSupplier.INSTANCE);
     if (creatorFactory == null) {
-      throw new RuntimeException("Could not find a way to create AutoValue class " + targetClass);
+      throw new RuntimeException(
+          "Could not find a way to create AutoValue class " + targetTypeDescriptor);
     }
 
     return creatorFactory;
@@ -139,6 +144,6 @@ public SchemaUserTypeCreator schemaTypeCreator(Class<?> targetClass, Schema sche
   @Override
   public <T> @Nullable Schema schemaFor(TypeDescriptor<T> typeDescriptor) {
     return JavaBeanUtils.schemaFromJavaBeanClass(
-        typeDescriptor.getRawType(), AbstractGetterTypeSupplier.INSTANCE);
+        typeDescriptor, AbstractGetterTypeSupplier.INSTANCE);
   }
 }

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

@@ -19,6 +19,7 @@
 
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -36,7 +37,7 @@
   "rawtypes"
 })
 public class CachingFactory<CreatedT> implements Factory<CreatedT> {
-  private transient @Nullable ConcurrentHashMap<Class, CreatedT> cache = null;
+  private transient @Nullable ConcurrentHashMap<TypeDescriptor<?>, CreatedT> cache = null;
 
   private final Factory<CreatedT> innerFactory;
 
@@ -45,16 +46,16 @@ public CachingFactory(Factory<CreatedT> innerFactory) {
   }
 
   @Override
-  public CreatedT create(Class<?> clazz, Schema schema) {
+  public CreatedT create(TypeDescriptor<?> typeDescriptor, Schema schema) {
     if (cache == null) {
       cache = new ConcurrentHashMap<>();
     }
-    CreatedT cached = cache.get(clazz);
+    CreatedT cached = cache.get(typeDescriptor);
     if (cached != null) {
       return cached;
     }
-    cached = innerFactory.create(clazz, schema);
-    cache.put(clazz, cached);
+    cached = innerFactory.create(typeDescriptor, schema);
+    cache.put(typeDescriptor, cached);
     return cached;
   }
 

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

@@ -19,9 +19,10 @@
 
 import java.io.Serializable;
 import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.values.TypeDescriptor;
 
 /** A Factory interface for schema-related objects for a specific Java type. */
 @Internal
 public interface Factory<T> extends Serializable {
-  T create(Class<?> clazz, Schema schema);
+  T create(TypeDescriptor<?> typeDescriptor, Schema schema);
 }

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

@@ -28,6 +28,7 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.AutoValue_FieldValueTypeInformation.Builder;
 import org.apache.beam.sdk.schemas.annotations.SchemaCaseFormat;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
@@ -78,7 +79,7 @@ public abstract class FieldValueTypeInformation implements Serializable {
   /** If the field has a description, returns the description for the field. */
   public abstract @Nullable String getDescription();
 
-  abstract Builder toBuilder();
+  public abstract Builder toBuilder();
 
   @AutoValue.Builder
   public abstract static class Builder {
@@ -106,7 +107,7 @@ public abstract static class Builder {
 
     public abstract Builder setDescription(@Nullable String fieldDescription);
 
-    abstract FieldValueTypeInformation build();
+    public abstract FieldValueTypeInformation build();
   }
 
   public static FieldValueTypeInformation forOneOf(
@@ -125,18 +126,19 @@ public static FieldValueTypeInformation forOneOf(
         .build();
   }
 
-  public static FieldValueTypeInformation forField(Field field, int index) {
-    TypeDescriptor<?> type = TypeDescriptor.of(field.getGenericType());
+  public static FieldValueTypeInformation forField(
+      TypeDescriptor typeDescriptor, Field field, int index) {
+    TypeDescriptor<?> type = typeDescriptor.resolveType(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();
@@ -184,7 +186,8 @@ public static <T extends AnnotatedElement & Member> String getNameOverride(
     return fieldDescription.value();
   }
 
-  public static FieldValueTypeInformation forGetter(Method method, int index) {
+  public static FieldValueTypeInformation forGetter(
+      TypeDescriptor typeDescriptor, Method method, int index) {
     String name;
     if (method.getName().startsWith("get")) {
       name = ReflectUtils.stripPrefix(method.getName(), "get");
@@ -193,8 +196,7 @@ public static FieldValueTypeInformation forGetter(Method method, int index) {
     } else {
       throw new RuntimeException("Getter has wrong prefix " + method.getName());
     }
-
-    TypeDescriptor<?> type = TypeDescriptor.of(method.getGenericReturnType());
+    TypeDescriptor<?> type = typeDescriptor.resolveType(method.getGenericReturnType());
     boolean nullable = hasNullableReturnType(method);
     return new AutoValue_FieldValueTypeInformation.Builder()
         .setName(getNameOverride(name, method))
@@ -252,19 +254,20 @@ private static boolean isNullableAnnotation(Annotation annotation) {
     return annotation.annotationType().getSimpleName().equals("Nullable");
   }
 
-  public static FieldValueTypeInformation forSetter(Method method) {
-    return forSetter(method, "set");
+  public static FieldValueTypeInformation forSetter(TypeDescriptor typeDescriptor, Method method) {
+    return forSetter(typeDescriptor, method, "set");
   }
 
-  public static FieldValueTypeInformation forSetter(Method method, String setterPrefix) {
+  public static FieldValueTypeInformation forSetter(
+      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 = typeDescriptor.resolveType(method.getGenericParameterTypes()[0]);
     boolean nullable = hasSingleNullableParameter(method);
     return new AutoValue_FieldValueTypeInformation.Builder()
         .setName(name)
@@ -279,15 +282,27 @@ public static FieldValueTypeInformation forSetter(Method method, String setterPr
         .build();
   }
 
+  public static FieldValueTypeInformation.Builder builder() {
+    return new AutoValue_FieldValueTypeInformation.Builder();
+  }
+
   public FieldValueTypeInformation withName(String name) {
     return toBuilder().setName(name).build();
   }
 
-  private static FieldValueTypeInformation getIterableComponentType(Field field) {
-    return getIterableComponentType(TypeDescriptor.of(field.getGenericType()));
+  public FieldValueTypeInformation withTypesFrom(FieldValueTypeInformation other) {
+    return toBuilder()
+        .setType(other.getType())
+        .setRawType(other.getRawType())
+        .setElementType(other.getElementType())
+        .setMapKeyType(other.getMapKeyType())
+        .setMapValueType(other.getMapValueType())
+        .setOneOfTypes(other.getOneOfTypes())
+        .build();
   }
 
-  static @Nullable FieldValueTypeInformation getIterableComponentType(TypeDescriptor<?> valueType) {
+  public static @Nullable FieldValueTypeInformation getIterableComponentType(
+      TypeDescriptor<?> valueType) {
     // TODO: Figure out nullable elements.
     TypeDescriptor<?> componentType = ReflectUtils.getIterableComponentType(valueType);
     if (componentType == null) {
@@ -308,22 +323,14 @@ private static FieldValueTypeInformation getIterableComponentType(Field field) {
 
   // 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()));
-  }
-
-  private static @Nullable FieldValueTypeInformation getMapKeyType(
+  public 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);
-  }
-
-  private static @Nullable FieldValueTypeInformation getMapValueType(
+  public static @Nullable FieldValueTypeInformation getMapValueType(
       TypeDescriptor<?> typeDescriptor) {
     return getMapType(typeDescriptor, 1);
   }