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);
  }