fix tableRowFromMessage

Author: reuvenlax

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/SerializableBiFunctions.java

@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Useful {@link SerializableFunction} overrides. */
+public class SerializableBiFunctions {
+  /** Always returns the first argument. */
+  public static <FirstInputT, SecondInputT, OutputT>
+      SerializableBiFunction<FirstInputT, SecondInputT, FirstInputT> select1st(
+          SerializableBiFunction<@Nullable FirstInputT, @Nullable SecondInputT, OutputT>
+              biFunction) {
+    return (t, u) -> t;
+  }
+
+  /** Always returns the second argument. */
+  public static <FirstInputT, SecondInputT, OutputT>
+      SerializableBiFunction<FirstInputT, SecondInputT, SecondInputT> select2nd(
+          SerializableBiFunction<@Nullable FirstInputT, @Nullable SecondInputT, OutputT>
+              biFunction) {
+    return (t, u) -> u;
+  }
+
+  /** Convert to a unary function by fixing the first argument. */
+  public static <FirstInputT, SecondInputT, OutputT>
+      SerializableFunction<SecondInputT, OutputT> fix1st(
+          SerializableBiFunction<@Nullable FirstInputT, @Nullable SecondInputT, OutputT> biFunction,
+          @Nullable FirstInputT value) {
+    return u -> biFunction.apply(value, u);
+  }
+
+  /** Convert to a unary function by fixing the second argument. */
+  public static <FirstInputT, SecondInputT, OutputT>
+      SerializableFunction<FirstInputT, OutputT> fix2nd(
+          SerializableBiFunction<@Nullable FirstInputT, @Nullable SecondInputT, OutputT> biFunction,
+          @Nullable SecondInputT value) {
+    return t -> biFunction.apply(t, value);
+  }
+
+  /** Convert from a unary function by ignoring the first argument. */
+  public static <FirstInputT, SecondInputT, OutputT>
+      SerializableBiFunction<FirstInputT, SecondInputT, OutputT> ignore1st(
+          SerializableFunction<@Nullable SecondInputT, OutputT> function) {
+    return (t, u) -> function.apply(u);
+  }
+
+  /** Convert from a unary function by ignoring the second argument. */
+  public static <FirstInputT, SecondInputT, OutputT>
+      SerializableBiFunction<FirstInputT, SecondInputT, OutputT> ignore2nd(
+          SerializableFunction<@Nullable FirstInputT, OutputT> function) {
+    return (t, u) -> function.apply(t);
+  }
+}

sdks/java/extensions/protobuf/src/test/proto/proto3_schema_messages.proto

@@ -33,6 +33,51 @@ import "proto3_schema_options.proto";
 
 option java_package = "org.apache.beam.sdk.extensions.protobuf";
 
+message PrimitiveEncodedFields {
+  int64 encoded_timestamp = 1;
+  int32 encoded_date = 2;
+  bytes encoded_numeric = 3;
+  bytes encoded_bignumeric = 4;
+  int64 encoded_packed_datetime = 5;
+  int64 encoded_packed_time = 6;
+}
+
+message NestedEncodedFields {
+  PrimitiveEncodedFields nested = 1;
+  repeated PrimitiveEncodedFields nested_list = 2;
+}
+
+message PrimitiveUnEncodedFields {
+  string timestamp = 1;
+  string date = 2;
+  string numeric = 3;
+  string bignumeric = 4;
+  string datetime = 5;
+  string time = 6;
+}
+
+message NestedUnEncodedFields {
+  PrimitiveUnEncodedFields nested = 1;
+  repeated PrimitiveUnEncodedFields nested_list = 2;
+}
+
+message WrapperUnEncodedFields {
+  google.protobuf.FloatValue float = 1;
+  google.protobuf.DoubleValue double = 2;
+  google.protobuf.BoolValue bool = 3;
+  google.protobuf.Int32Value int32 = 4;
+  google.protobuf.Int64Value int64 = 5;
+  google.protobuf.UInt32Value uint32 = 6;
+  google.protobuf.UInt64Value uint64 = 7;
+  google.protobuf.BytesValue bytes = 8;
+  google.protobuf.Timestamp timestamp = 9;
+}
+
+message NestedWrapperUnEncodedFields {
+  WrapperUnEncodedFields nested = 1;
+  repeated WrapperUnEncodedFields nested_list = 2;
+}
+
 message Primitive {
     double primitive_double = 1;
     float primitive_float = 2;
@@ -287,4 +332,4 @@ message NoWrapPrimitive {
   optional bool bool = 13;
   optional string string = 14;
   optional bytes bytes = 15;
-}
+}

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java

@@ -182,6 +182,7 @@ public ByteString mergeNewFields(
   public TableRow toTableRow(ByteString protoBytes, Predicate<String> includeField) {
     try {
       return TableRowToStorageApiProto.tableRowFromMessage(
+          getSchemaInformation(),
           DynamicMessage.parseFrom(
               TableRowToStorageApiProto.wrapDescriptorProto(getDescriptor()), protoBytes),
           true,

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java

@@ -42,6 +42,7 @@
 import com.google.cloud.bigquery.storage.v1.AppendRowsRequest;
 import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest;
 import com.google.cloud.bigquery.storage.v1.DataFormat;
+import com.google.cloud.bigquery.storage.v1.ProtoSchemaConverter;
 import com.google.cloud.bigquery.storage.v1.ReadSession;
 import com.google.cloud.bigquery.storage.v1.ReadStream;
 import com.google.gson.JsonArray;
@@ -119,6 +120,8 @@
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.Reshuffle;
+import org.apache.beam.sdk.transforms.SerializableBiFunction;
+import org.apache.beam.sdk.transforms.SerializableBiFunctions;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.SerializableFunctions;
 import org.apache.beam.sdk.transforms.SimpleFunction;
@@ -2297,10 +2300,61 @@ public static <T extends Message> Write<T> writeProtos(Class<T> protoMessageClas
     if (DynamicMessage.class.equals(protoMessageClass)) {
       throw new IllegalArgumentException("DynamicMessage is not supported.");
     }
-    return BigQueryIO.<T>write()
-        .withFormatFunction(
-            m -> TableRowToStorageApiProto.tableRowFromMessage(m, false, Predicates.alwaysTrue()))
-        .withWriteProtosClass(protoMessageClass);
+    try {
+      return BigQueryIO.<T>write()
+          .toBuilder()
+          .setFormatFunction(FormatProto.fromClass(protoMessageClass))
+          .build()
+          .withWriteProtosClass(protoMessageClass);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private static class FormatProto<T extends Message>
+      implements SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, T, TableRow> {
+    transient TableRowToStorageApiProto.SchemaInformation inferredSchemaInformation;
+    final Class<T> protoMessageClass;
+
+    FormatProto(Class<T> protoMessageClass) {
+      this.protoMessageClass = protoMessageClass;
+    }
+
+    TableRowToStorageApiProto.SchemaInformation inferSchemaInformation() {
+      try {
+        if (inferredSchemaInformation == null) {
+          Descriptors.Descriptor descriptor =
+              (Descriptors.Descriptor)
+                  org.apache.beam.sdk.util.Preconditions.checkStateNotNull(
+                          protoMessageClass.getMethod("getDescriptor"))
+                      .invoke(null);
+          Descriptors.Descriptor convertedDescriptor =
+              TableRowToStorageApiProto.wrapDescriptorProto(
+                  ProtoSchemaConverter.convert(descriptor).getProtoDescriptor());
+          TableSchema tableSchema =
+              TableRowToStorageApiProto.protoSchemaToTableSchema(
+                  TableRowToStorageApiProto.tableSchemaFromDescriptor(convertedDescriptor));
+          this.inferredSchemaInformation =
+              TableRowToStorageApiProto.SchemaInformation.fromTableSchema(tableSchema);
+        }
+        return inferredSchemaInformation;
+      } catch (Exception e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    static <T extends Message> FormatProto<T> fromClass(Class<T> protoMessageClass)
+        throws Exception {
+      return new FormatProto<>(protoMessageClass);
+    }
+
+    @Override
+    public TableRow apply(TableRowToStorageApiProto.SchemaInformation schemaInformation, T input) {
+      TableRowToStorageApiProto.SchemaInformation localSchemaInformation =
+          schemaInformation != null ? schemaInformation : inferSchemaInformation();
+      return TableRowToStorageApiProto.tableRowFromMessage(
+          localSchemaInformation, input, false, Predicates.alwaysTrue());
+    }
   }
 
   /** Implementation of {@link #write}. */
@@ -2354,9 +2408,13 @@ public enum Method {
     abstract @Nullable SerializableFunction<ValueInSingleWindow<T>, TableDestination>
         getTableFunction();
 
-    abstract @Nullable SerializableFunction<T, TableRow> getFormatFunction();
+    abstract @Nullable SerializableBiFunction<
+            TableRowToStorageApiProto.SchemaInformation, T, TableRow>
+        getFormatFunction();
 
-    abstract @Nullable SerializableFunction<T, TableRow> getFormatRecordOnFailureFunction();
+    abstract @Nullable SerializableBiFunction<
+            TableRowToStorageApiProto.SchemaInformation, T, TableRow>
+        getFormatRecordOnFailureFunction();
 
     abstract RowWriterFactory.@Nullable AvroRowWriterFactory<T, ?, ?> getAvroRowWriterFactory();
 
@@ -2467,10 +2525,13 @@ abstract static class Builder<T> {
       abstract Builder<T> setTableFunction(
           SerializableFunction<ValueInSingleWindow<T>, TableDestination> tableFunction);
 
-      abstract Builder<T> setFormatFunction(SerializableFunction<T, TableRow> formatFunction);
+      abstract Builder<T> setFormatFunction(
+          SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, T, TableRow>
+              formatFunction);
 
       abstract Builder<T> setFormatRecordOnFailureFunction(
-          SerializableFunction<T, TableRow> formatFunction);
+          SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, T, TableRow>
+              formatFunction);
 
       abstract Builder<T> setAvroRowWriterFactory(
           RowWriterFactory.AvroRowWriterFactory<T, ?, ?> avroRowWriterFactory);
@@ -2718,7 +2779,9 @@ public Write<T> to(DynamicDestinations<T, ?> dynamicDestinations) {
 
     /** Formats the user's type into a {@link TableRow} to be written to BigQuery. */
     public Write<T> withFormatFunction(SerializableFunction<T, TableRow> formatFunction) {
-      return toBuilder().setFormatFunction(formatFunction).build();
+      return toBuilder()
+          .setFormatFunction(SerializableBiFunctions.ignore1st(formatFunction))
+          .build();
     }
 
     /**
@@ -2733,7 +2796,9 @@ public Write<T> withFormatFunction(SerializableFunction<T, TableRow> formatFunct
      */
     public Write<T> withFormatRecordOnFailureFunction(
         SerializableFunction<T, TableRow> formatFunction) {
-      return toBuilder().setFormatRecordOnFailureFunction(formatFunction).build();
+      return toBuilder()
+          .setFormatRecordOnFailureFunction(SerializableBiFunctions.ignore1st(formatFunction))
+          .build();
     }
 
     /**
@@ -3599,9 +3664,10 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) {
     private <DestinationT> WriteResult expandTyped(
         PCollection<T> input, DynamicDestinations<T, DestinationT> dynamicDestinations) {
       boolean optimizeWrites = getOptimizeWrites();
-      SerializableFunction<T, TableRow> formatFunction = getFormatFunction();
-      SerializableFunction<T, TableRow> formatRecordOnFailureFunction =
-          getFormatRecordOnFailureFunction();
+      SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, T, TableRow>
+          formatFunction = getFormatFunction();
+      SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, T, TableRow>
+          formatRecordOnFailureFunction = getFormatRecordOnFailureFunction();
       RowWriterFactory.AvroRowWriterFactory<T, ?, DestinationT> avroRowWriterFactory =
           (RowWriterFactory.AvroRowWriterFactory<T, ?, DestinationT>) getAvroRowWriterFactory();
 
@@ -3623,7 +3689,8 @@ private <DestinationT> WriteResult expandTyped(
           // If no format function set, then we will automatically convert the input type to a
           // TableRow.
           // TODO: it would be trivial to convert to avro records here instead.
-          formatFunction = BigQueryUtils.toTableRow(input.getToRowFunction());
+          formatFunction =
+              SerializableBiFunctions.ignore1st(BigQueryUtils.toTableRow(input.getToRowFunction()));
         }
         // Infer the TableSchema from the input Beam schema.
         // TODO: If the user provided a schema, we should use that. There are things that can be
@@ -3769,8 +3836,9 @@ private <DestinationT> WriteResult continueExpandTyped(
                     getCreateDisposition(),
                     dynamicDestinations,
                     elementCoder,
-                    tableRowWriterFactory.getToRowFn(),
-                    tableRowWriterFactory.getToFailsafeRowFn())
+                    SerializableBiFunctions.fix1st(tableRowWriterFactory.getToRowFn(), null),
+                    SerializableBiFunctions.fix1st(
+                        tableRowWriterFactory.getToFailsafeRowFn(), null))
                 .withInsertRetryPolicy(retryPolicy)
                 .withTestServices(getBigQueryServices())
                 .withExtendedErrorInfo(getExtendedErrorInfo())

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java

@@ -55,6 +55,7 @@
 import org.apache.beam.sdk.schemas.Schema.FieldType;
 import org.apache.beam.sdk.schemas.logicaltypes.NanosDuration;
 import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableBiFunction;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
 import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter;
@@ -641,14 +642,15 @@ public Write<?> fromConfigRow(Row configRow, PipelineOptions options) {
         if (formatFunctionBytes != null) {
           builder =
               builder.setFormatFunction(
-                  (SerializableFunction<?, TableRow>) fromByteArray(formatFunctionBytes));
+                  (SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, ?, TableRow>)
+                      fromByteArray(formatFunctionBytes));
         }
         byte[] formatRecordOnFailureFunctionBytes =
             configRow.getBytes("format_record_on_failure_function");
         if (formatRecordOnFailureFunctionBytes != null) {
           builder =
               builder.setFormatRecordOnFailureFunction(
-                  (SerializableFunction<?, TableRow>)
+                  (SerializableBiFunction<TableRowToStorageApiProto.SchemaInformation, ?, TableRow>)
                       fromByteArray(formatRecordOnFailureFunctionBytes));
         }
         byte[] avroRowWriterFactoryBytes = configRow.getBytes("avro_row_writer_factory");