Fix nullness in BigQueryIO

.github/trigger_files/beam_PostCommit_Java.json

@@ -1,4 +1,5 @@
 {
+  "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO",
   "comment": "Modify this file in a trivial way to cause this test suite to run",
   "modification": 6
 } 

.github/trigger_files/beam_PostCommit_Java_BigQueryEarlyRollout.json

@@ -0,0 +1,3 @@
+{
+  "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO"
+}

.github/trigger_files/beam_PostCommit_Java_DataflowV1.json

@@ -1,4 +1,5 @@
 {
+  "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO",
   "https://github.com/apache/beam/pull/39330": "Fix DataflowV1 test failure by fixing getSimpleName access",
   "https://github.com/apache/beam/pull/34902": "Introducing OutputBuilder",
   "https://github.com/apache/beam/pull/35177": "Introducing WindowedValueReceiver to runners",

.github/trigger_files/beam_PostCommit_Java_DataflowV2.json

@@ -1,4 +1,5 @@
 {
+  "https://github.com/apache/beam/pull/39893": "Fix nullness in BigQueryIO",
   "modification": 9,
   "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface"
 }

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

@@ -317,7 +317,6 @@ public static DynamicMessage messageFromGenericRecord(
     return builder.build();
   }
 
-  @SuppressWarnings("nullness")
   private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Schema.Field field) {
     @Nullable Schema schema = field.schema();
 
@@ -360,12 +359,14 @@ private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Sch
         if (valueType == null) {
           throw new RuntimeException("Unexpected null element type!");
         }
-        TableFieldSchema keyFieldSchema =
-            fieldDescriptorFromAvroField(
-                new Schema.Field("key", keyType, "key of the map entry", null));
-        TableFieldSchema valueFieldSchema =
-            fieldDescriptorFromAvroField(
-                new Schema.Field("value", valueType, "value of the map entry", null));
+        // The Avro Field constructor accepts a null default value, but Avro is not annotated.
+        @SuppressWarnings("nullness")
+        Schema.Field keyField = new Schema.Field("key", keyType, "key of the map entry", null);
+        @SuppressWarnings("nullness")
+        Schema.Field valueField =
+            new Schema.Field("value", valueType, "value of the map entry", null);
+        TableFieldSchema keyFieldSchema = fieldDescriptorFromAvroField(keyField);
+        TableFieldSchema valueFieldSchema = fieldDescriptorFromAvroField(valueField);
         builder =
             builder
                 .setType(TableFieldSchema.Type.STRUCT)
@@ -382,9 +383,10 @@ private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Sch
         Preconditions.checkState(
             elementType.getType() != Schema.Type.UNION,
             "Multiple non-null union types are not supported.");
-        TableFieldSchema unionFieldSchema =
-            fieldDescriptorFromAvroField(
-                new Schema.Field(field.name(), elementType, field.doc(), null));
+        // The Avro Field constructor accepts a null default value, but Avro is not annotated.
+        @SuppressWarnings("nullness")
+        Schema.Field unionField = new Schema.Field(field.name(), elementType, field.doc(), null);
+        TableFieldSchema unionFieldSchema = fieldDescriptorFromAvroField(unionField);
         builder =
             builder
                 .setType(unionFieldSchema.getType())

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

@@ -29,9 +29,6 @@ class AvroRowWriter<AvroT, T> extends BigQueryRowWriter<T> {
   private final Schema schema;
   private final SerializableFunction<AvroWriteRequest<T>, AvroT> toAvroRecord;
 
-  @SuppressWarnings({
-    "nullness" // calling superclass method in constructor flagged as error; TODO: fix
-  })
   AvroRowWriter(
       String basename,
       Schema schema,
@@ -42,8 +39,12 @@ class AvroRowWriter<AvroT, T> extends BigQueryRowWriter<T> {
 
     this.schema = schema;
     this.toAvroRecord = toAvroRecord;
-    this.writer =
+    // getOutputStream() is established by the superclass constructor, which the checker
+    // cannot see through the partially-initialized receiver.
+    @SuppressWarnings("nullness")
+    DataFileWriter<AvroT> initializedWriter =
         new DataFileWriter<>(writerFactory.apply(schema)).create(schema, getOutputStream());
+    this.writer = initializedWriter;
   }
 
   @Override

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

@@ -162,12 +162,12 @@ class BatchLoads<DestinationT, ElementT>
   private long maxBytesPerPartition;
   private int numFileShards;
   private @Nullable Duration triggeringFrequency;
-  private ValueProvider<String> customGcsTempLocation;
+  private @Nullable ValueProvider<String> customGcsTempLocation;
   private @Nullable ValueProvider<String> loadJobProjectId;
   private final Coder<ElementT> elementCoder;
   private final RowWriterFactory<ElementT, DestinationT> rowWriterFactory;
   private final @Nullable String kmsKey;
-  private final String tempDataset;
+  private final @Nullable String tempDataset;
   private final BadRecordRouter badRecordRouter;
   private final ErrorHandler<BadRecord, ?> badRecordErrorHandler;
   private Coder<TableDestination> tableDestinationCoder;
@@ -181,15 +181,15 @@ class BatchLoads<DestinationT, ElementT>
       boolean singletonTable,
       DynamicDestinations<?, DestinationT> dynamicDestinations,
       Coder<DestinationT> destinationCoder,
-      ValueProvider<String> customGcsTempLocation,
+      @Nullable ValueProvider<String> customGcsTempLocation,
       @Nullable ValueProvider<String> loadJobProjectId,
       boolean ignoreUnknownValues,
       Coder<ElementT> elementCoder,
       RowWriterFactory<ElementT, DestinationT> rowWriterFactory,
       @Nullable String kmsKey,
       boolean clusteringEnabled,
       boolean useAvroLogicalTypes,
-      String tempDataset,
+      @Nullable String tempDataset,
       BadRecordRouter badRecordRouter,
       ErrorHandler<BadRecord, ?> badRecordErrorHandler) {
     bigQueryServices = new BigQueryServicesImpl();
@@ -249,7 +249,7 @@ public void setMaxNumWritersPerBundle(int maxNumWritersPerBundle) {
     this.maxNumWritersPerBundle = maxNumWritersPerBundle;
   }
 
-  public void setTriggeringFrequency(Duration triggeringFrequency) {
+  public void setTriggeringFrequency(@Nullable Duration triggeringFrequency) {
     this.triggeringFrequency = triggeringFrequency;
   }
 
@@ -285,6 +285,7 @@ public void validate(@Nullable PipelineOptions maybeOptions) {
     PipelineOptions options = Preconditions.checkArgumentNotNull(maybeOptions);
     // We will use a BigQuery load job -- validate the temp location.
     String tempLocation;
+    ValueProvider<String> customGcsTempLocation = this.customGcsTempLocation;
     if (customGcsTempLocation == null) {
       tempLocation = options.getTempLocation();
     } else {
@@ -424,7 +425,7 @@ private WriteResult expandTriggered(PCollection<KV<DestinationT, ElementT>> inpu
             .apply("ExtractTempTables", Values.create())
             .apply(
                 ParDo.of(
-                        new UpdateSchemaDestination<DestinationT>(
+                        new UpdateSchemaDestination<>(
                             bigQueryServices,
                             zeroLoadJobIdPrefixView,
                             loadJobProjectId,
@@ -530,7 +531,7 @@ public WriteResult expandUntriggered(PCollection<KV<DestinationT, ElementT>> inp
             .apply("ReifyRenameInput", new ReifyAsIterable<>())
             .apply(
                 ParDo.of(
-                        new UpdateSchemaDestination<DestinationT>(
+                        new UpdateSchemaDestination<>(
                             bigQueryServices,
                             zeroLoadJobIdPrefixView,
                             loadJobProjectId,
@@ -592,6 +593,8 @@ private PCollectionView<String> createTempFilePrefixView(
                       @ProcessElement
                       public void getTempFilePrefix(ProcessContext c) {
                         String tempLocationRoot;
+                        ValueProvider<String> customGcsTempLocation =
+                            BatchLoads.this.customGcsTempLocation;
                         if (customGcsTempLocation != null && customGcsTempLocation.get() != null) {
                           tempLocationRoot = customGcsTempLocation.get();
                         } else {

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

@@ -546,9 +546,6 @@ private static boolean hasNamespaceCollision(List<TableFieldSchema> fieldSchemas
     return false;
   }
 
-  @SuppressWarnings({
-    "nullness" // Avro library not annotated
-  })
   private static Field convertField(
       TableFieldSchema bigQueryField, Boolean useAvroLogicalTypes, @Nullable String namespace) {
     String fieldName = bigQueryField.getName();
@@ -569,11 +566,15 @@ private static Field convertField(
     } else if (!"REQUIRED".equals(bqMode)) {
       throw new IllegalArgumentException(String.format("Unknown BigQuery Field Mode: %s", bqMode));
     }
-    return new Field(
-        fieldName,
-        fieldSchema,
-        bigQueryField.getDescription(),
-        (Object) null /* Cast to avoid deprecated JsonNode constructor. */);
+    // The Avro Field constructor accepts a null default value, but Avro is not annotated.
+    @SuppressWarnings("nullness")
+    Field field =
+        new Field(
+            fieldName,
+            fieldSchema,
+            bigQueryField.getDescription(),
+            (Object) null /* Cast to avoid deprecated JsonNode constructor. */);
+    return field;
   }
 
   static TableSchema fromGenericAvroSchema(Schema schema) {

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

@@ -21,6 +21,7 @@
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.schemas.AutoValueSchema;
 import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.checkerframework.dataflow.qual.Pure;
 
 /**
  * Configuration for reading from BigQuery.
@@ -31,9 +32,6 @@
  * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam
  * repository.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 @DefaultSchema(AutoValueSchema.class)
 @AutoValue
 public abstract class BigQueryExportReadSchemaTransformConfiguration {
@@ -44,21 +42,25 @@ public static Builder builder() {
   }
 
   /** Configures the BigQuery read job with the SQL query. */
+  @Pure
   @Nullable
   public abstract String getQuery();
 
   /**
    * Specifies a table for a BigQuery read job. See {@link BigQueryIO.TypedRead#from(String)} for
    * more details on the expected format.
    */
+  @Pure
   @Nullable
   public abstract String getTableSpec();
 
   /** BigQuery geographic location where the query job will be executed. */
+  @Pure
   @Nullable
   public abstract String getQueryLocation();
 
   /** Enables BigQuery's Standard SQL dialect when reading from a query. */
+  @Pure
   @Nullable
   public abstract Boolean getUseStandardSql();
 

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

@@ -32,7 +32,7 @@
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TypeDescriptor;
 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.Strings;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * An implementation of {@link TypedSchemaTransformProvider} for BigQuery read jobs configured using
@@ -42,9 +42,6 @@
  * provide no backwards compatibility guarantees, and it should not be implemented outside the Beam
  * repository.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 @Internal
 @AutoService(SchemaTransformProvider.class)
 public class BigQueryExportReadSchemaTransformProvider
@@ -96,7 +93,7 @@ public List<String> outputCollectionNames() {
    */
   protected static class BigQueryExportSchemaTransform extends SchemaTransform {
     /** An instance of {@link BigQueryServices} used for testing. */
-    private BigQueryServices testBigQueryServices = null;
+    private @Nullable BigQueryServices testBigQueryServices = null;
 
     private final BigQueryExportReadSchemaTransformConfiguration configuration;
 
@@ -135,19 +132,19 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
     BigQueryIO.TypedRead<TableRow> toTypedRead() {
       BigQueryIO.TypedRead<TableRow> read = BigQueryIO.readTableRowsWithSchema();
 
-      if (!Strings.isNullOrEmpty(configuration.getQuery())) {
+      if (configuration.getQuery() != null && !configuration.getQuery().isEmpty()) {
         read = read.fromQuery(configuration.getQuery());
       }
 
-      if (!Strings.isNullOrEmpty(configuration.getTableSpec())) {
+      if (configuration.getTableSpec() != null && !configuration.getTableSpec().isEmpty()) {
         read = read.from(configuration.getTableSpec());
       }
 
       if (configuration.getUseStandardSql() != null && configuration.getUseStandardSql()) {
         read = read.usingStandardSql();
       }
 
-      if (!Strings.isNullOrEmpty(configuration.getQueryLocation())) {
+      if (configuration.getQueryLocation() != null && !configuration.getQueryLocation().isEmpty()) {
         read = read.withQueryLocation(configuration.getQueryLocation());
       }
 

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.io.gcp.bigquery;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
@@ -27,6 +28,8 @@
 import com.google.api.services.bigquery.model.Dataset;
 import com.google.api.services.bigquery.model.ErrorProto;
 import com.google.api.services.bigquery.model.Job;
+import com.google.api.services.bigquery.model.JobConfiguration;
+import com.google.api.services.bigquery.model.JobConfigurationLoad;
 import com.google.api.services.bigquery.model.JobReference;
 import com.google.api.services.bigquery.model.JobStatus;
 import com.google.api.services.bigquery.model.Table;
@@ -471,9 +474,6 @@ static <K, V> List<V> getOrCreateMapListValue(Map<K, List<V>> map, K key) {
    *
    * <p>If the project id is omitted, the default project id is used.
    */
-  @SuppressWarnings({
-    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-  })
   public static TableReference parseTableSpec(String tableSpec) {
     Matcher match = BigQueryIO.TABLE_SPEC.matcher(tableSpec);
     if (!match.matches()) {
@@ -545,7 +545,9 @@ public static TableReference parseTableSpec(String tableSpec) {
     }
 
     TableReference ref = new TableReference();
-    ref.setProjectId(project);
+    // The project id is optional; the API client accepts a null project id but is not annotated.
+    @SuppressWarnings("nullness")
+    TableReference unused = ref.setProjectId(project);
     return ref.setDatasetId(dataset).setTableId(table);
   }
 
@@ -561,9 +563,6 @@ private static IllegalArgumentException invalidTableSpec(String tableSpec) {
             tableSpec));
   }
 
-  @SuppressWarnings({
-    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-  })
   public static TableReference parseTableUrn(String tableUrn) {
     Matcher match = BigQueryIO.TABLE_URN_SPEC.matcher(tableUrn);
     if (!match.matches()) {
@@ -573,10 +572,10 @@ public static TableReference parseTableUrn(String tableUrn) {
               + tableUrn);
     }
 
-    TableReference ref = new TableReference();
-    ref.setProjectId(match.group("PROJECT"));
-
-    return ref.setDatasetId(match.group("DATASET")).setTableId(match.group("TABLE"));
+    return new TableReference()
+        .setProjectId(checkStateNotNull(match.group("PROJECT")))
+        .setDatasetId(checkStateNotNull(match.group("DATASET")))
+        .setTableId(checkStateNotNull(match.group("TABLE")));
   }
 
   /** Strip off any partition decorator information from a tablespec. */
@@ -585,19 +584,22 @@ public static String stripPartitionDecorator(String tableSpec) {
     return (index == -1) ? tableSpec : tableSpec.substring(0, index);
   }
 
-  @SuppressWarnings({
-    "nullness" // The BigQuery API library is documented to accept nulls but is not annotated
-  })
   static String jobToPrettyString(@Nullable Job job) throws IOException {
-    if (job != null && job.getConfiguration().getLoad() != null) {
+    if (job == null) {
+      return "null";
+    }
+    JobConfiguration configuration = job.getConfiguration();
+    if (configuration != null && configuration.getLoad() != null) {
       // Removing schema and sourceUris from error messages for load jobs since these fields can be
       // quite long and error message might not be displayed properly in runner specific logs.
       job = job.clone();
-      job.getConfiguration().getLoad().setSchema(null);
-      job.getConfiguration().getLoad().setSourceUris(null);
+      JobConfigurationLoad load = checkStateNotNull(job.getConfiguration()).getLoad();
+      // The BigQuery API library is documented to accept nulls here but is not annotated.
+      @SuppressWarnings("nullness")
+      JobConfigurationLoad unused = load.setSchema(null).setSourceUris(null);
     }
 
-    return job == null ? "null" : job.toPrettyString();
+    return job.toPrettyString();
   }
 
   static String statusToPrettyString(@Nullable JobStatus status) throws IOException {