Merge pull request #32081: make FieldValueTypeInformation creators take a TypeDescriptor parameter

Author: tilgalas

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

@@ -19,7 +19,6 @@
 
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
-import java.util.Comparator;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.schemas.annotations.SchemaIgnore;
@@ -32,13 +31,10 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** A {@link SchemaProvider} for AutoValue classes. */
-@SuppressWarnings({
-  "nullness", // TODO(https://github.com/apache/beam/issues/20497)
-  "rawtypes"
-})
 public class AutoValueSchema extends GetterBasedSchemaProviderV2 {
   /** {@link FieldValueTypeSupplier} that's based on AutoValue getters. */
   @VisibleForTesting
@@ -49,7 +45,11 @@ public static class AbstractGetterTypeSupplier implements FieldValueTypeSupplier
     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.
-      TypeDescriptor<?> targetTypeDescriptor = AutoValueUtils.getBaseAutoValueClass(typeDescriptor);
+      TypeDescriptor<?> targetTypeDescriptor =
+          Preconditions.checkNotNull(
+              AutoValueUtils.getBaseAutoValueClass(typeDescriptor),
+              "unable to determine base AutoValue class for type {}",
+              typeDescriptor);
 
       List<Method> methods =
           ReflectUtils.getMethods(targetTypeDescriptor.getRawType()).stream()
@@ -62,9 +62,9 @@ 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));
+      types.sort(JavaBeanUtils.comparingNullFirst(FieldValueTypeInformation::getNumber));
       validateFieldNumbers(types);
       return types;
     }
@@ -89,8 +89,8 @@ private static void validateFieldNumbers(List<FieldValueTypeInformation> types)
   }
 
   @Override
-  public List<FieldValueGetter> fieldValueGetters(
-      TypeDescriptor<?> targetTypeDescriptor, Schema schema) {
+  public <T> List<FieldValueGetter<@NonNull T, Object>> fieldValueGetters(
+      TypeDescriptor<T> targetTypeDescriptor, Schema schema) {
     return JavaBeanUtils.getGetters(
         targetTypeDescriptor,
         schema,

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

@@ -20,6 +20,9 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import org.apache.beam.sdk.values.TypeDescriptor;
+import org.checkerframework.checker.initialization.qual.NotOnlyInitialized;
+import org.checkerframework.checker.initialization.qual.UnknownInitialization;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -32,24 +35,25 @@
  * significant for larger schemas) on each lookup. This wrapper caches the value returned by the
  * inner factory, so the schema comparison only need happen on the first lookup.
  */
-@SuppressWarnings({
-  "nullness", // TODO(https://github.com/apache/beam/issues/20497)
-  "rawtypes"
-})
-public class CachingFactory<CreatedT> implements Factory<CreatedT> {
+public class CachingFactory<CreatedT extends @NonNull Object> implements Factory<CreatedT> {
   private transient @Nullable ConcurrentHashMap<TypeDescriptor<?>, CreatedT> cache = null;
 
-  private final Factory<CreatedT> innerFactory;
+  private final @NotOnlyInitialized Factory<CreatedT> innerFactory;
 
-  public CachingFactory(Factory<CreatedT> innerFactory) {
+  public CachingFactory(@UnknownInitialization Factory<CreatedT> innerFactory) {
     this.innerFactory = innerFactory;
   }
 
-  @Override
-  public CreatedT create(TypeDescriptor<?> typeDescriptor, Schema schema) {
+  private ConcurrentHashMap<TypeDescriptor<?>, CreatedT> getCache() {
     if (cache == null) {
       cache = new ConcurrentHashMap<>();
     }
+    return cache;
+  }
+
+  @Override
+  public CreatedT create(TypeDescriptor<?> typeDescriptor, Schema schema) {
+    ConcurrentHashMap<TypeDescriptor<?>, CreatedT> cache = getCache();
     CreatedT cached = cache.get(typeDescriptor);
     if (cached != null) {
       return cached;

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

@@ -19,6 +19,7 @@
 
 import java.io.Serializable;
 import org.apache.beam.sdk.annotations.Internal;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -29,7 +30,7 @@
  * <p>Implementations of this interface are generated at runtime to map object fields to Row fields.
  */
 @Internal
-public interface FieldValueGetter<ObjectT, ValueT> extends Serializable {
+public interface FieldValueGetter<ObjectT extends @NonNull Object, ValueT> extends Serializable {
   @Nullable
   ValueT get(ObjectT object);
 

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

@@ -27,7 +27,9 @@
 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.annotations.Internal;
 import org.apache.beam.sdk.schemas.annotations.SchemaCaseFormat;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
@@ -40,10 +42,7 @@
 
 /** Represents type information for a Java type that will be used to infer a Schema type. */
 @AutoValue
-@SuppressWarnings({
-  "nullness", // TODO(https://github.com/apache/beam/issues/20497)
-  "rawtypes"
-})
+@Internal
 public abstract class FieldValueTypeInformation implements Serializable {
   /** Optionally returns the field index. */
   public abstract @Nullable Integer getNumber();
@@ -125,18 +124,23 @@ public static FieldValueTypeInformation forOneOf(
         .build();
   }
 
-  public static FieldValueTypeInformation forField(Field field, int index) {
-    TypeDescriptor<?> type = TypeDescriptor.of(field.getGenericType());
+  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 +189,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 +203,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 +267,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 +311,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,34 +330,23 @@ 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);
   }
 
   // If the Field is a map type, returns the key or value type (0 is key type, 1 is value).
   // Otherwise returns a null reference.
-  @SuppressWarnings("unchecked")
   private static @Nullable FieldValueTypeInformation getMapType(
       TypeDescriptor<?> valueType, int index) {
-    TypeDescriptor mapType = ReflectUtils.getMapType(valueType, index);
+    TypeDescriptor<?> mapType = ReflectUtils.getMapType(valueType, index);
     if (mapType == null) {
       return null;
     }