Merge pull request #54 from data-integrations/feature/update-record-format

Update to use a RecordFormat that doesn’t involve StreamEvent

Author: chtyim

kafka-plugins-0.10/src/main/java/co/cask/hydrator/plugin/batch/source/KafkaBatchSource.java

@@ -26,7 +26,6 @@
 import co.cask.cdap.api.data.format.StructuredRecord;
 import co.cask.cdap.api.data.schema.Schema;
 import co.cask.cdap.api.dataset.lib.KeyValue;
-import co.cask.cdap.api.flow.flowlet.StreamEvent;
 import co.cask.cdap.common.io.ByteBuffers;
 import co.cask.cdap.etl.api.Emitter;
 import co.cask.cdap.etl.api.PipelineConfigurer;
@@ -50,6 +49,7 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -68,7 +68,7 @@ public class KafkaBatchSource extends BatchSource<KafkaKey, KafkaMessage, Struct
   private final Kafka10BatchConfig config;
   private List<KafkaRequest> kafkaRequests;
   private Schema schema;
-  private RecordFormat<StreamEvent, StructuredRecord> recordFormat;
+  private RecordFormat<ByteBuffer, StructuredRecord> recordFormat;
   private String messageField;
   private FileContext fileContext;
   private Path offsetsFile;
@@ -231,7 +231,7 @@ public void transform(KeyValue<KafkaKey, KafkaMessage> input, Emitter<Structured
     if (config.getFormat() == null) {
       builder.set(messageField, ByteBuffers.getByteArray(input.getValue().getPayload()));
     } else {
-      StructuredRecord messageRecord = recordFormat.read(new StreamEvent(input.getValue().getPayload()));
+      StructuredRecord messageRecord = recordFormat.read(input.getValue().getPayload());
       for (Schema.Field field : messageRecord.getSchema().getFields()) {
         String fieldName = field.getName();
         builder.set(fieldName, messageRecord.get(fieldName));

kafka-plugins-0.10/src/main/java/co/cask/hydrator/plugin/source/KafkaStreamingSource.java

@@ -23,7 +23,6 @@
 import co.cask.cdap.api.data.format.RecordFormat;
 import co.cask.cdap.api.data.format.StructuredRecord;
 import co.cask.cdap.api.data.schema.Schema;
-import co.cask.cdap.api.flow.flowlet.StreamEvent;
 import co.cask.cdap.etl.api.PipelineConfigurer;
 import co.cask.cdap.etl.api.streaming.StreamingContext;
 import co.cask.cdap.etl.api.streaming.StreamingSource;
@@ -271,7 +270,7 @@ protected void addMessage(StructuredRecord.Builder builder, String messageField,
    * Everything here should be serializable, as Spark Streaming will serialize all functions.
    */
   private static class FormatFunction extends BaseFunction {
-    private transient RecordFormat<StreamEvent, StructuredRecord> recordFormat;
+    private transient RecordFormat<ByteBuffer, StructuredRecord> recordFormat;
 
     FormatFunction(long ts, KafkaConfig conf) {
       super(ts, conf);
@@ -283,11 +282,11 @@ protected void addMessage(StructuredRecord.Builder builder, String messageField,
       if (recordFormat == null) {
         Schema messageSchema = conf.getMessageSchema();
         FormatSpecification spec =
-          new FormatSpecification(conf.getFormat(), messageSchema, new HashMap<String, String>());
+          new FormatSpecification(conf.getFormat(), messageSchema, new HashMap<>());
         recordFormat = RecordFormats.createInitializedFormat(spec);
       }
 
-      StructuredRecord messageRecord = recordFormat.read(new StreamEvent(ByteBuffer.wrap(message)));
+      StructuredRecord messageRecord = recordFormat.read(ByteBuffer.wrap(message));
       for (Schema.Field field : messageRecord.getSchema().getFields()) {
         String fieldName = field.getName();
         builder.set(fieldName, messageRecord.get(fieldName));

kafka-plugins-0.8/src/main/java/co/cask/hydrator/plugin/batch/source/KafkaBatchSource.java

@@ -25,7 +25,6 @@
 import co.cask.cdap.api.data.format.StructuredRecord;
 import co.cask.cdap.api.data.schema.Schema;
 import co.cask.cdap.api.dataset.lib.KeyValue;
-import co.cask.cdap.api.flow.flowlet.StreamEvent;
 import co.cask.cdap.common.io.ByteBuffers;
 import co.cask.cdap.etl.api.Emitter;
 import co.cask.cdap.etl.api.PipelineConfigurer;
@@ -45,6 +44,7 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -64,7 +64,7 @@ public class KafkaBatchSource extends BatchSource<KafkaKey, KafkaMessage, Struct
   private Path offsetsFile;
   private List<KafkaRequest> kafkaRequests;
   private Schema schema;
-  private RecordFormat<StreamEvent, StructuredRecord> recordFormat;
+  private RecordFormat<ByteBuffer, StructuredRecord> recordFormat;
   private String messageField;
 
   public KafkaBatchSource(KafkaBatchConfig config) {
@@ -164,7 +164,7 @@ public void transform(KeyValue<KafkaKey, KafkaMessage> input, Emitter<Structured
     if (config.getFormat() == null) {
       builder.set(messageField, ByteBuffers.getByteArray(input.getValue().getPayload()));
     } else {
-      StructuredRecord messageRecord = recordFormat.read(new StreamEvent(input.getValue().getPayload()));
+      StructuredRecord messageRecord = recordFormat.read(input.getValue().getPayload());
       for (Schema.Field field : messageRecord.getSchema().getFields()) {
         String fieldName = field.getName();
         builder.set(fieldName, messageRecord.get(fieldName));

kafka-plugins-0.8/src/main/java/co/cask/hydrator/plugin/source/KafkaStreamingSource.java

@@ -23,7 +23,6 @@
 import co.cask.cdap.api.data.format.RecordFormat;
 import co.cask.cdap.api.data.format.StructuredRecord;
 import co.cask.cdap.api.data.schema.Schema;
-import co.cask.cdap.api.flow.flowlet.StreamEvent;
 import co.cask.cdap.etl.api.PipelineConfigurer;
 import co.cask.cdap.etl.api.streaming.StreamingContext;
 import co.cask.cdap.etl.api.streaming.StreamingSource;
@@ -271,7 +270,7 @@ protected void addMessage(StructuredRecord.Builder builder, String messageField,
    * Everything here should be serializable, as Spark Streaming will serialize all functions.
    */
   private static class FormatFunction extends BaseFunction {
-    private transient RecordFormat<StreamEvent, StructuredRecord> recordFormat;
+    private transient RecordFormat<ByteBuffer, StructuredRecord> recordFormat;
 
     FormatFunction(long ts, KafkaConfig conf) {
       super(ts, conf);
@@ -283,11 +282,11 @@ protected void addMessage(StructuredRecord.Builder builder, String messageField,
       if (recordFormat == null) {
         Schema messageSchema = conf.getMessageSchema();
         FormatSpecification spec =
-          new FormatSpecification(conf.getFormat(), messageSchema, new HashMap<String, String>());
+          new FormatSpecification(conf.getFormat(), messageSchema, new HashMap<>());
         recordFormat = RecordFormats.createInitializedFormat(spec);
       }
 
-      StructuredRecord messageRecord = recordFormat.read(new StreamEvent(ByteBuffer.wrap(message)));
+      StructuredRecord messageRecord = recordFormat.read(ByteBuffer.wrap(message));
       for (Schema.Field field : messageRecord.getSchema().getFields()) {
         String fieldName = field.getName();
         builder.set(fieldName, messageRecord.get(fieldName));