[ 
https://issues.apache.org/jira/browse/BEAM-4454?focusedWorklogId=171429&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-171429
 ]

ASF GitHub Bot logged work on BEAM-4454:
----------------------------------------

                Author: ASF GitHub Bot
            Created on: 03/Dec/18 00:24
            Start Date: 03/Dec/18 00:24
    Worklog Time Spent: 10m 
      Work Description: reuvenlax commented on a change in pull request #7181: 
[BEAM-4454] Add more AVRO utilities to convert between Beam and Avro.
URL: https://github.com/apache/beam/pull/7181#discussion_r238122681
 
 

 ##########
 File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/AvroUtils.java
 ##########
 @@ -20,87 +20,210 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
-import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import org.apache.avro.Schema.Type;
 import org.apache.avro.generic.GenericEnumSymbol;
 import org.apache.avro.generic.GenericFixed;
 import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.generic.GenericRecordBuilder;
+import org.apache.avro.reflect.ReflectData;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
 import org.apache.beam.sdk.values.Row;
 
 /** Utils to convert AVRO records to Beam rows. */
 @Experimental(Experimental.Kind.SCHEMAS)
 public class AvroUtils {
+  // Unwrap an AVRO schema into the base type an whether it is nullable.
+  static class TypeWithNullability {
+    public final org.apache.avro.Schema type;
+    public final boolean nullable;
+
+    TypeWithNullability(org.apache.avro.Schema avroSchema) {
+      if (avroSchema.getType() == org.apache.avro.Schema.Type.UNION) {
+        List<org.apache.avro.Schema> types = avroSchema.getTypes();
+
+        // optional fields in AVRO have form of:
+        // {"name": "foo", "type": ["null", "something"]}
+
+        // don't need recursion because nested unions aren't supported in AVRO
+        List<org.apache.avro.Schema> nonNullTypes =
+            types
+                .stream()
+                .filter(x -> x.getType() != org.apache.avro.Schema.Type.NULL)
+                .collect(Collectors.toList());
+
+        if (nonNullTypes.size() == types.size() || nonNullTypes.isEmpty()) {
+          // union without `null` or all 'null' union, keep as is.
+          type = avroSchema;
+          nullable = false;
+        } else if (nonNullTypes.size() > 1) {
+          type = org.apache.avro.Schema.createUnion(nonNullTypes);
+          nullable = true;
+        } else {
+          // One non-null type.
+          type = nonNullTypes.get(0);
+          nullable = true;
+        }
+      } else {
+        type = avroSchema;
+        nullable = false;
+      }
+    }
+  }
+
   private AvroUtils() {}
 
   /**
    * Converts AVRO schema to Beam row schema.
    *
    * @param schema schema of type RECORD
    */
-  public static Schema toSchema(@Nonnull org.apache.avro.Schema schema) {
+  public static Schema toBeamSchema(org.apache.avro.Schema schema) {
     Schema.Builder builder = Schema.builder();
 
     for (org.apache.avro.Schema.Field field : schema.getFields()) {
-      org.apache.avro.Schema unwrapped = unwrapNullableSchema(field.schema());
+      TypeWithNullability nullableType = new 
TypeWithNullability(field.schema());
+      builder.addField(field.name(), toFieldType(nullableType));
+    }
+
+    return builder.build();
+  }
+
+  /** Converts a Beam Schema into an AVRO schema. */
+  public static org.apache.avro.Schema toAvroSchema(Schema beamSchema) {
+    List<org.apache.avro.Schema.Field> fields = Lists.newArrayList();
+    for (Schema.Field field : beamSchema.getFields()) {
+      org.apache.avro.Schema fieldSchema = getFieldSchema(field.getType());
+      org.apache.avro.Schema.Field recordField =
+          new org.apache.avro.Schema.Field(
+              field.getName(), fieldSchema, field.getDescription(), null);
+      fields.add(recordField);
+    }
+    org.apache.avro.Schema avroSchema = 
org.apache.avro.Schema.createRecord(fields);
+    return avroSchema;
+  }
+
+  /**
+   * Strict conversion from AVRO to Beam, strict because it doesn't do 
widening or narrowing during
+   * conversion. If Schema is not provided, one is inferred from the AVRO 
schema.
+   */
+  public static Row toBeamRowStrict(GenericRecord record, @Nullable Schema 
schema) {
+    if (schema == null) {
+      schema = toBeamSchema(record.getSchema());
+    }
+
+    Row.Builder builder = Row.withSchema(schema);
+    org.apache.avro.Schema avroSchema = record.getSchema();
+
+    for (Schema.Field field : schema.getFields()) {
+      Object value = record.get(field.getName());
+      org.apache.avro.Schema fieldAvroSchema = 
avroSchema.getField(field.getName()).schema();
 
-      if (!unwrapped.equals(field.schema())) {
-        builder.addNullableField(field.name(), toFieldType(unwrapped));
+      if (value == null) {
+        builder.addValue(null);
       } else {
-        builder.addField(field.name(), toFieldType(unwrapped));
+        builder.addValue(convertAvroFieldStrict(value, fieldAvroSchema, 
field.getType()));
       }
     }
 
     return builder.build();
   }
 
+  /**
+   * Convert from a Beam Row to an AVRO GenericRecord. If a Schema is not 
provided, one is inferred
+   * from the Beam schema on the orw.
+   */
+  public static GenericRecord toGenericRecord(
+      Row row, @Nullable org.apache.avro.Schema avroSchema) {
+    Schema beamSchema = row.getSchema();
+    // Use the provided AVRO schema if present, otherwise infer an AVRO schema 
from the row
+    // schema.
+    if (avroSchema != null && avroSchema.getFields().size() != 
beamSchema.getFieldCount()) {
+      throw new IllegalArgumentException("AVRO schema doesn't match row 
schema");
 
 Review comment:
   Done

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Issue Time Tracking
-------------------

    Worklog Id:     (was: 171429)
    Time Spent: 2.5h  (was: 2h 20m)

> Provide automatic schema registration for AVROs
> -----------------------------------------------
>
>                 Key: BEAM-4454
>                 URL: https://issues.apache.org/jira/browse/BEAM-4454
>             Project: Beam
>          Issue Type: Sub-task
>          Components: sdk-java-core
>            Reporter: Reuven Lax
>            Assignee: Reuven Lax
>            Priority: Major
>          Time Spent: 2.5h
>  Remaining Estimate: 0h
>
> Need to make sure this is a compatible change



--
This message was sent by Atlassian JIRA
(v7.6.3#76005)

Reply via email to