let writer pass latest schema to WorkUnitState
Closes from ZihanLi58/GOBBLIN-859
@@ -48,7 +48,6 @@ import org.apache.gobblin.dataset.PartitionDescriptor;
import org.apache.gobblin.instrumented.writer.InstrumentedDataWriterDecorator;
import org.apache.gobblin.instrumented.writer.InstrumentedPartitionedDataWriterDecorator;
import org.apache.gobblin.records.ControlMessageHandler;
-import org.apache.gobblin.source.extractor.CheckpointableWatermark;
import org.apache.gobblin.stream.ControlMessage;
import org.apache.gobblin.stream.MetadataUpdateControlMessage;
import org.apache.gobblin.stream.RecordEnvelope;
@@ -67,6 +66,7 @@ import org.apache.gobblin.writer.partitioner.WriterPartitioner;
@Slf4j
public class PartitionedDataWriter<S, D> extends WriterWrapper<D> implements FinalState, SpeculativeAttemptAwareConstruct, WatermarkAwareWriter<D> {
+ public static final String WRITER_LATEST_SCHEMA = "writer.latest.schema";
private static final GenericRecord NON_PARTITIONED_WRITER_KEY =
new GenericData.Record(SchemaBuilder.record("Dummy").fields().endRecord());
@@ -96,6 +96,9 @@ public class PartitionedDataWriter<S, D> extends WriterWrapper<D> implements Fin
this.closer = Closer.create();
this.writerBuilder = builder;
this.controlMessageHandler = new PartitionDataWriterMessageHandler();
+ if(builder.schema != null) {
+ this.state.setProp(WRITER_LATEST_SCHEMA, builder.getSchema());
+ }
this.partitionWriters = CacheBuilder.newBuilder().build(new CacheLoader<GenericRecord, DataWriter<D>>() {
@Override
public DataWriter<D> load(final GenericRecord key)
@@ -321,6 +324,8 @@ public class PartitionedDataWriter<S, D> extends WriterWrapper<D> implements Fin
if (message instanceof MetadataUpdateControlMessage) {
PartitionedDataWriter.this.writerBuilder.withSchema(((MetadataUpdateControlMessage) message)
.getGlobalMetadata().getSchema());
+ state.setProp(WRITER_LATEST_SCHEMA, ((MetadataUpdateControlMessage) message)
+ .getGlobalMetadata().getSchema());
}
for (DataWriter writer : PartitionedDataWriter.this.partitionWriters.asMap().values()) {
