How to use org.apache.gobblin.fork

/**
 * Returns a {@link org.apache.gobblin.fork.CopyableSchema} wrapper around the given {@link Schema}.
 * {@inheritDoc}
 * @see org.apache.gobblin.converter.Converter#convertSchema(java.lang.Object, org.apache.gobblin.configuration.WorkUnitState)
 */
@Override
public CopyableSchema convertSchema(Schema inputSchema, WorkUnitState workUnit) throws SchemaConversionException {
 return new CopyableSchema(inputSchema);
}

public RecordWithForkMap(ControlMessage<D> record, int activeBranchesForRecord) {
 this.record = record;
 this.forkMap = null;
 this.copiesLeft = activeBranchesForRecord;
 this.mustCopy = this.copiesLeft > 1;
 this.cloner = buildForkCloner();
}

 @Override
 public boolean test(RecordWithForkMap<D> dRecordWithForkMap) {
  return dRecordWithForkMap.sendToBranch(this.forkIdx);
 }
}

@Override
public GlobalMetadata<S> copy() throws CopyNotSupportedException {
 if (CopyHelper.isCopyable(schema)) {
  return new GlobalMetadata((S)CopyHelper.copy(schema));
 }
 throw new CopyNotSupportedException("Type is not copyable: " + schema.getClass().getName());
}

/**
 * Copy this object if needed.
 * @param thing : this object that needs to be copied
 * @return: a possibly copied instance
 * @throws CopyNotSupportedException if thing needs to be copied but cannot be
 */
public static Object copy(Object thing) throws CopyNotSupportedException {
 if (!isCopyable(thing)) {
  throw new CopyNotSupportedException(thing.getClass().getName() + " cannot be copied. See Copyable");
 }
 if (thing instanceof Copyable) {
  return ((Copyable) thing).copy();
 }
 // Support for a few primitive types out of the box
 if (thing instanceof byte[]) {
  byte[] copy = new byte[((byte[]) thing).length];
  System.arraycopy(thing, 0, copy, 0, ((byte[]) thing).length);
  return copy;
 }
 // Assume that everything other type is immutable, not checking this again
 return thing;
}

@Test
public void testCopyable()
  throws CopyNotSupportedException {
 Copyable c = mock(Copyable.class);
 Assert.assertTrue(CopyHelper.isCopyable(c));
 Object copy = new Object();
 when(c.copy()).thenReturn(copy);
 Assert.assertEquals(CopyHelper.copy(c), copy);
 Assert.assertEquals(CopyHelper.copy(c), copy);
}

@Test
public void testByteArray()
  throws CopyNotSupportedException {
 int length = RANDOM.nextInt(200);
 byte[] bytes = new byte[length];
 RANDOM.nextBytes(bytes);
 Assert.assertTrue(CopyHelper.isCopyable(bytes));
 byte[] copiedBytes = (byte[]) CopyHelper.copy(bytes);
 Assert.assertTrue(copiedBytes != bytes, "Copied bytes reference should be different for every copy after that");
 Assert.assertEquals(copiedBytes, bytes, "Copied bytes value should be the same");
}

@Override
public List<Boolean> forkSchema(WorkUnitState workUnitState, S input) {
 return this.embeddedForkOperator.forkSchema(workUnitState, input);
}

@Override
public List<Boolean> forkDataRecordImpl(WorkUnitState workUnitState, D input) {
 return this.embeddedForkOperator.forkDataRecord(workUnitState, input);
}

@Override
public int getBranches(WorkUnitState workUnitState) {
 return this.embeddedForkOperator.getBranches(workUnitState);
}

 @Override
 public GenericRecord copy()
   throws CopyNotSupportedException {
  if (!(this.record instanceof GenericData.Record)) {
   throw new CopyNotSupportedException(
     "The record to make copy is not an instance of " + GenericData.Record.class.getName());
  }
  // Make a deep copy of the original record
  return new GenericData.Record((GenericData.Record) this.record, true);
 }
}

@Override
public void init(WorkUnitState workUnitState) throws Exception {
 this.embeddedForkOperator.init(workUnitState);
 super.init(workUnitState, DecoratorUtils.resolveUnderlyingObject(this).getClass());
}

@Override
protected StreamEntity<D> buildClone() {
 try {
  return new RecordEnvelope<>((D) CopyHelper.copy(_record), this, false);
 } catch (CopyNotSupportedException cnse) {
  throw new UnsupportedOperationException(cnse);
 }
}

 /**
  * Returns a {@link org.apache.gobblin.fork.CopyableGenericRecord} wrapper around the given {@link GenericRecord}.
  * {@inheritDoc}
  * @see org.apache.gobblin.converter.Converter#convertRecord(java.lang.Object, java.lang.Object, org.apache.gobblin.configuration.WorkUnitState)
  */
 @Override
 public Iterable<CopyableGenericRecord> convertRecord(CopyableSchema outputSchema, GenericRecord inputRecord,
   WorkUnitState workUnit) throws DataConversionException {
  return new SingleRecordIterable<>(new CopyableGenericRecord(inputRecord));
 }
}

@Override
public List<Boolean> forkDataRecord(WorkUnitState workUnitState, D input) {
 this.records.clear();
 for (int i = 0; i < getBranches(workUnitState); i++) {
  this.records.add(Boolean.TRUE);
 }
 return this.records;
}

/**
 * Check if an object is copyable using the {@link #copy(Object)} method.
 * @param thing: the object that needs to be copied
 * @return: true if {@link CopyHelper} can copy this thing, false otherwise
 */
public static boolean isCopyable(Object thing) {
 if (
   (thing instanceof Copyable)
     || (thing instanceof byte[])
     || (isImmutableType(thing))
   ) {
  return true;
 }
 return false;
}

@Override
public List<Boolean> forkSchema(WorkUnitState workUnitState, String input) {
 return Collections.nCopies(getBranches(workUnitState), true);
}

@Override
public CopyableSchema convertSchema(String schema, WorkUnitState workUnit) {
 return new CopyableSchema(new Schema.Parser().parse(schema));
}

public RecordWithForkMap(RecordEnvelope<D> record, List<Boolean> forkMap) {
 this.record = record;
 this.forkMap = Lists.newArrayList(forkMap);
 this.mustCopy = mustCopy(forkMap);
 this.copiesLeft = this.forkMap.stream().filter(x -> x).count();
 this.cloner = buildForkCloner();
}

@Override
public List<Boolean> forkSchema(WorkUnitState workUnitState, S input) {
 this.schemas.clear();
 for (int i = 0; i < getBranches(workUnitState); i++) {
  this.schemas.add(Boolean.TRUE);
 }
 return this.schemas;
}