[BEAM-4453] Add convenience methods for pojo and javabean schema registration.

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

@@ -28,6 +28,7 @@
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.function.Function;
+import javax.annotation.Nullable;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
 import org.apache.beam.sdk.transforms.SerializableFunction;
@@ -67,8 +68,40 @@ private static class SchemaEntry<T> {
   private final Map<TypeDescriptor, SchemaEntry> entries = Maps.newHashMap();
   private final ArrayDeque<SchemaProvider> providers;
 
+  class PerTypeRegisteredProvider implements SchemaProvider {
+    private Map<TypeDescriptor, SchemaProvider> providers = Maps.newHashMap();
+
+    void registerProvider(TypeDescriptor typeDescriptor, SchemaProvider schemaProvider) {
+      providers.put(typeDescriptor, schemaProvider);
+    }
+
+    @Nullable
+    @Override
+    public <T> Schema schemaFor(TypeDescriptor<T> typeDescriptor) {
+      SchemaProvider schemaProvider = providers.get(typeDescriptor);
+      return (schemaProvider != null) ? schemaProvider.schemaFor(typeDescriptor) : null;
+    }
+
+    @Nullable
+    @Override
+    public <T> SerializableFunction<T, Row> toRowFunction(TypeDescriptor<T> typeDescriptor) {
+      SchemaProvider schemaProvider = providers.get(typeDescriptor);
+      return (schemaProvider != null) ? schemaProvider.toRowFunction(typeDescriptor) : null;
+    }
+
+    @Nullable
+    @Override
+    public <T> SerializableFunction<Row, T> fromRowFunction(TypeDescriptor<T> typeDescriptor) {
+      SchemaProvider schemaProvider = providers.get(typeDescriptor);
+      return (schemaProvider != null) ? schemaProvider.fromRowFunction(typeDescriptor) : null;
+    }
+  }
+
+  private PerTypeRegisteredProvider perTypeRegisteredProviders = new PerTypeRegisteredProvider();
+
   private SchemaRegistry() {
     providers = new ArrayDeque<>(REGISTERED_SCHEMA_PROVIDERS);
+    providers.addFirst(perTypeRegisteredProviders);
   }
 
   public static SchemaRegistry createDefault() {
@@ -104,6 +137,57 @@ public void registerSchemaProvider(SchemaProvider schemaProvider) {
     providers.addFirst(schemaProvider);
   }
 
+  /** Register a {@link SchemaProvider} to be used for a specific type. * */
+  public <T> void registerSchemaProvider(Class<T> clazz, SchemaProvider schemaProvider) {
+    registerSchemaProvider(TypeDescriptor.of(clazz), schemaProvider);
+  }
+
+  /** Register a {@link SchemaProvider} to be used for a specific type. * */
+  public <T> void registerSchemaProvider(
+      TypeDescriptor<T> typeDescriptor, SchemaProvider schemaProvider) {
+    perTypeRegisteredProviders.registerProvider(typeDescriptor, schemaProvider);
+  }
+
+  /**
+   * Register a POJO type for automatic schema inference.
+   *
+   * <p>Currently schema field names will match field names in the POJO, and all fields must be
+   * mutable (i.e. no final fields).
+   */
+  public <T> void registerPOJO(Class<T> clazz) {
+    registerPOJO(TypeDescriptor.of(clazz));
+  }
+
+  /**
+   * Register a POJO type for automatic schema inference.
+   *
+   * <p>Currently schema field names will match field names in the POJO, and all fields must be
+   * mutable (i.e. no final fields).
+   */
+  public <T> void registerPOJO(TypeDescriptor<T> typeDescriptor) {
+    registerSchemaProvider(typeDescriptor, new JavaFieldSchema());
+  }
+
+  /**
+   * Register a JavaBean type for automatic schema inference.
+   *
+   * <p>Currently schema field names will match getter names in the bean, and all getters must have
+   * matching setters.
+   */
+  public <T> void registerJavaBean(Class<T> clazz) {
+    registerJavaBean(TypeDescriptor.of(clazz));
+  }
+
+  /**
+   * Register a JavaBean type for automatic schema inference.
+   *
+   * <p>Currently schema field names will match getter names in the bean, and all getters must have
+   * matching setters.
+   */
+  public <T> void registerJavaBean(TypeDescriptor<T> typeDescriptor) {
+    registerSchemaProvider(typeDescriptor, new JavaBeanSchema());
+  }
+
   /**
    * Get a schema for a given {@link Class} type. If no schema exists, throws {@link
    * NoSuchSchemaException}.
@@ -151,13 +235,13 @@ public <T> SerializableFunction<T, Row> getToRowFunction(TypeDescriptor<T> typeD
     return getProviderResult((SchemaProvider p) -> p.toRowFunction(typeDescriptor));
   }
 
-  /** Rerieve the function that converts a {@link Row} object to an object of the specified type. */
+  /** Retrieve the function that converts a {@link Row} object to the specified type. */
   public <T> SerializableFunction<Row, T> getFromRowFunction(Class<T> clazz)
       throws NoSuchSchemaException {
     return getFromRowFunction(TypeDescriptor.of(clazz));
   }
 
-  /** Rerieve the function that converts a {@link Row} object to an object of the specified type. */
+  /** Retrieve the function that converts a {@link Row} object to the specified type. */
   public <T> SerializableFunction<Row, T> getFromRowFunction(TypeDescriptor<T> typeDescriptor)
       throws NoSuchSchemaException {
     SchemaEntry entry = entries.get(typeDescriptor);

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

@@ -18,11 +18,16 @@
 
 package org.apache.beam.sdk.schemas;
 
+import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.SIMPLE_BEAN_SCHEMA;
+import static org.apache.beam.sdk.schemas.utils.TestPOJOs.SIMPLE_POJO_SCHEMA;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import com.google.auto.service.AutoService;
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import org.apache.beam.sdk.schemas.utils.TestJavaBeans.SimpleBean;
+import org.apache.beam.sdk.schemas.utils.TestPOJOs.SimplePOJO;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TypeDescriptor;
@@ -202,4 +207,20 @@ public void testDefaultSchemaProvider() throws NoSuchSchemaException {
     SchemaRegistry registry = SchemaRegistry.createDefault();
     assertEquals(EMPTY_SCHEMA, registry.getSchema(TestDefaultSchemaClass.class));
   }
+
+  @Test
+  public void testRegisterPojo() throws NoSuchSchemaException {
+    SchemaRegistry registry = SchemaRegistry.createDefault();
+    registry.registerPOJO(SimplePOJO.class);
+    Schema schema = registry.getSchema(SimplePOJO.class);
+    assertTrue(SIMPLE_POJO_SCHEMA.equivalent(schema));
+  }
+
+  @Test
+  public void testRegisterJavaBean() throws NoSuchSchemaException {
+    SchemaRegistry registry = SchemaRegistry.createDefault();
+    registry.registerJavaBean(SimpleBean.class);
+    Schema schema = registry.getSchema(SimpleBean.class);
+    assertTrue(SIMPLE_BEAN_SCHEMA.equivalent(schema));
+  }
 }