Merge pull request #1272 from data-integrations/cherry_fix_emptybqsink

[6.9.2][Cherry-pick]  PLUGIN-1445 : BQ sink in case empty task: Use a supplier and initialize only when there are records to write

Author: sahusanket

src/main/java/io/cdap/plugin/gcp/bigquery/sink/AvroOutputFormat.java

@@ -28,14 +28,19 @@
 import org.apache.hadoop.io.NullWritable;
 import org.apache.hadoop.mapreduce.RecordWriter;
 import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.io.UncheckedIOException;
+import java.util.function.Supplier;
 
 /**
  * avro output format.
  */
 public class AvroOutputFormat extends AvroKeyOutputFormat<GenericRecord> {
+
   public AvroOutputFormat() {
     super();
   }
@@ -65,21 +70,34 @@ public RecordWriter<AvroKey<GenericRecord>, NullWritable> getRecordWriter(TaskAt
     }
 
     GenericData dataModel = AvroSerialization.createDataModel(conf);
-    return create(writerSchema, dataModel, getCompressionCodec(context),
-                  getAvroFileOutputStream(context), getSyncInterval(context));
+
+    return create(writerSchema, dataModel, getCompressionCodec(context), getOutputStreamSupplier(context),
+                  getSyncInterval(context));
+  }
+
+  //Creating a supplier for outputstream , which will get triggered only if there is a record to written.
+  //This is to avoid creation of output stream for empty tasks ( no writes )
+  private Supplier<OutputStream> getOutputStreamSupplier(TaskAttemptContext context) {
+    return  () -> {
+      try {
+        return getAvroFileOutputStream(context);
+      } catch (IOException e) {
+        throw new UncheckedIOException(e);
+      }
+    };
   }
 
   /**
    * Creates a new record writer instance.
    *
    * @param writerSchema The writer schema for the records to write.
    * @param compressionCodec The compression type for the writer file.
-   * @param outputStream The target output stream for the records.
+   * @param outputStreamSupplier A supplier to give the target output stream for the records.
    * @param syncInterval The sync interval for the writer file.
    */
   private RecordWriter<AvroKey<GenericRecord>, NullWritable> create(
     Schema writerSchema, GenericData dataModel, CodecFactory compressionCodec,
-    OutputStream outputStream, int syncInterval) throws IOException {
-    return new AvroRecordWriter(writerSchema, dataModel, compressionCodec, outputStream, syncInterval);
+    Supplier<OutputStream> outputStreamSupplier, int syncInterval) {
+    return new AvroRecordWriter(writerSchema, dataModel, compressionCodec, outputStreamSupplier, syncInterval);
   }
 }

src/main/java/io/cdap/plugin/gcp/bigquery/sink/AvroRecordWriter.java

@@ -30,6 +30,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.function.Supplier;
 
 /**
  * avro record writer
@@ -40,36 +41,36 @@ public class AvroRecordWriter extends RecordWriter<AvroKey<GenericRecord>, NullW
   private Schema prevSchema;
   private GenericData dataModel;
   private CodecFactory compressionCodec;
-  private OutputStream outputStream;
+  private Supplier<OutputStream> outputStreamSupplier;
   private int syncInterval;
 
   /**
    * Constructor.
    *
    * @param writerSchema The writer schema for the records in the Avro container file.
    * @param compressionCodec A compression codec factory for the Avro container file.
-   * @param outputStream The output stream to write the Avro container file to.
+   * @param outputStreamSupplier A supplier to give the output stream to write the Avro container file to.
    * @param syncInterval The sync interval for the Avro container file.
    * @throws IOException If the record writer cannot be opened.
    */
   public AvroRecordWriter(Schema writerSchema, GenericData dataModel, CodecFactory compressionCodec,
-                          OutputStream outputStream, int syncInterval) throws IOException {
+                          Supplier<OutputStream>  outputStreamSupplier, int syncInterval) {
     this.dataModel = dataModel;
     this.compressionCodec = compressionCodec;
-    this.outputStream = outputStream;
     this.syncInterval = syncInterval;
+    this.outputStreamSupplier = outputStreamSupplier;
   }
   /**
    * Constructor.
    *
    * @param writerSchema The writer schema for the records in the Avro container file.
    * @param compressionCodec A compression codec factory for the Avro container file.
-   * @param outputStream The output stream to write the Avro container file to.
+   * @param outputStreamSupplier The output stream to write the Avro container file to.
    * @throws IOException If the record writer cannot be opened.
    */
   public AvroRecordWriter(Schema writerSchema, GenericData dataModel,
-                          CodecFactory compressionCodec, OutputStream outputStream) throws IOException {
-    this(writerSchema, dataModel, compressionCodec, outputStream,
+                          CodecFactory compressionCodec, Supplier<OutputStream>  outputStreamSupplier) {
+    this(writerSchema, dataModel, compressionCodec, outputStreamSupplier,
          DataFileConstants.DEFAULT_SYNC_INTERVAL);
   }
 
@@ -78,6 +79,7 @@ public AvroRecordWriter(Schema writerSchema, GenericData dataModel,
   public void write(AvroKey<GenericRecord> record, NullWritable ignore) throws IOException {
     // Create an Avro container file and a writer to it.
     Schema writerSchema = record.datum().getSchema();
+
     if (mAvroFileWriter == null) {
       createFileWriter(writerSchema);
     }
@@ -94,7 +96,7 @@ private void createFileWriter(Schema writerSchema) throws IOException {
     mAvroFileWriter = new DataFileWriter<GenericRecord>(dataModel.createDatumWriter(writerSchema));
     mAvroFileWriter.setCodec(compressionCodec);
     mAvroFileWriter.setSyncInterval(syncInterval);
-    mAvroFileWriter.create(writerSchema, outputStream);
+    mAvroFileWriter.create(writerSchema, outputStreamSupplier.get());
     prevSchema = writerSchema;
   }
 

src/test/java/io/cdap/plugin/gcp/bigquery/sink/AvroRecordWriterTest.java

@@ -0,0 +1,97 @@
+package io.cdap.plugin.gcp.bigquery.sink;
+
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.plugin.format.avro.StructuredToAvroTransformer;
+import org.apache.avro.file.CodecFactory;
+import org.apache.avro.file.DataFileReader;
+import org.apache.avro.file.SeekableByteArrayInput;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.mapred.AvroKey;
+import org.apache.hadoop.io.NullWritable;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.time.LocalDate;
+import java.time.LocalTime;
+import java.util.function.Supplier;
+
+
+public class AvroRecordWriterTest {
+
+  @Test
+  public void testSuccessfulWrites() throws IOException {
+    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+    Supplier<OutputStream> outputStreamSupplier = () -> byteArrayOutputStream;
+
+    AvroRecordWriter avroRecordWriter = new AvroRecordWriter(
+      null,
+      GenericData.get(),
+      CodecFactory.nullCodec(),
+      outputStreamSupplier,
+      33);
+
+    //Create a Generic record
+    Schema schema = Schema.recordOf("record",
+                                    Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
+                                    Schema.Field.of("name", Schema.of(Schema.Type.STRING)),
+                                    Schema.Field.of("price", Schema.of(Schema.Type.DOUBLE)),
+                                    Schema.Field.of("dt", Schema.nullableOf(Schema.of(Schema.LogicalType.DATE))),
+                                    Schema.Field.of("time", Schema.of(Schema.LogicalType.TIME_MICROS)),
+                                    Schema.Field.of("timestamp",
+                                                    Schema.nullableOf(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS))));
+
+    StructuredRecord sourceRecord = StructuredRecord.builder(schema)
+      .set("id", 1L)
+      .set("name", "alice")
+      .set("price", 1.2d)
+      .setDate("dt", LocalDate.of(2018, 11, 11))
+      .setTime("time", LocalTime.of(11, 11, 11))
+      .set("timestamp", 1464181635000000L).build();
+    GenericRecord sourceGenericRecord = new StructuredToAvroTransformer(schema).transform(sourceRecord);
+    AvroKey<GenericRecord> record = new AvroKey(sourceGenericRecord);
+
+    //Write using avroRecordWriter
+    avroRecordWriter.write(record, Mockito.mock(NullWritable.class));
+    avroRecordWriter.close(null);
+
+    //Now read the avro record from the Output stream
+    GenericDatumReader genericDatumReader = new GenericDatumReader<GenericRecord>(sourceGenericRecord.getSchema());
+    SeekableByteArrayInput inputStream = new SeekableByteArrayInput(byteArrayOutputStream.toByteArray());
+    DataFileReader dataFileReader = new DataFileReader<GenericRecord>(inputStream, genericDatumReader);
+
+    //Check if it matches
+    Assert.assertTrue(dataFileReader.iterator().hasNext());
+    //Record read must match the sourceGenericRecord
+    Assert.assertEquals(dataFileReader.iterator().next(), sourceGenericRecord);
+  }
+
+
+  /**
+   * If there are no records, then the output stream supplier should not be called
+   * If called then will throw exception
+   *
+   * @throws IOException
+   */
+  @Test
+  public void testZeroWrites() throws IOException {
+    Supplier<OutputStream> outputStreamSupplier = () -> {
+      throw new IllegalStateException("This supplier should not be called");
+    };
+
+    AvroRecordWriter avroRecordWriter = new AvroRecordWriter(
+      null,
+      GenericData.get(),
+      CodecFactory.nullCodec(),
+      outputStreamSupplier,
+      33);
+
+    avroRecordWriter.close(null);
+  }
+}