org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker Java Examples

@Test
public void testInvokeProcessElementOutputDisallowedAfterFailedTryClaim() throws Exception {
  DoFn<Void, String> brokenFn =
      new DoFn<Void, String>() {
        @ProcessElement
        public void process(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
          assertFalse(tracker.tryClaim(6L));
          c.output("foo");
        }

        @GetInitialRestriction
        public OffsetRange getInitialRestriction(@Element Void element) {
          throw new UnsupportedOperationException("Should not be called in this test");
        }
      };
  e.expectMessage("Output is not allowed after a failed tryClaim()");
  runTest(brokenFn, new OffsetRange(0, 5));
}
 

@Test
public void testDoFnWithReturn() throws Exception {
  class MockFn extends DoFn<String, String> {
    @DoFn.ProcessElement
    public ProcessContinuation processElement(
        ProcessContext c, RestrictionTracker<SomeRestriction, Void> tracker) throws Exception {
      return null;
    }

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element String element) {
      return null;
    }

    @NewTracker
    public SomeRestrictionTracker newTracker(@Restriction SomeRestriction restriction) {
      return null;
    }
  }

  MockFn fn = mock(MockFn.class);
  when(fn.processElement(mockProcessContext, null)).thenReturn(resume());
  assertEquals(resume(), invokeProcessElement(fn));
}
 

@Test
public void testDefaultWatermarkEstimatorStateAndCoder() throws Exception {
  class MockFn extends DoFn<String, String> {
    @ProcessElement
    public void processElement(
        ProcessContext c, RestrictionTracker<RestrictionWithDefaultTracker, Void> tracker) {}

    @GetInitialRestriction
    public RestrictionWithDefaultTracker getInitialRestriction(@Element String element) {
      return null;
    }
  }

  MockFn fn = mock(MockFn.class);
  DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);

  CoderRegistry coderRegistry = CoderRegistry.createDefault();
  coderRegistry.registerCoderProvider(
      CoderProviders.fromStaticMethods(
          RestrictionWithDefaultTracker.class, CoderForDefaultTracker.class));
  assertEquals(VoidCoder.of(), invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry));
  assertNull(invoker.invokeGetInitialWatermarkEstimatorState(new FakeArgumentProvider<>()));
}
 

@Test
public void testMissingNewWatermarkEstimatorMethod() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context,
        RestrictionTracker<SomeRestriction, Void> tracker,
        ManualWatermarkEstimator<Instant> watermarkEstimator) {}

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction() {
      return null;
    }

    @GetInitialWatermarkEstimatorState
    public Instant getInitialWatermarkEstimatorState() {
      return null;
    }
  }

  thrown.expectMessage(
      "Splittable, either @NewWatermarkEstimator method must be defined or Instant must implement HasDefaultWatermarkEstimator.");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@Test
public void testSplittableMissingNewTrackerMethod() throws Exception {
  class OtherRestriction {}

  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<OtherRestriction, Void> tracker) {}

    @GetInitialRestriction
    public OtherRestriction getInitialRestriction() {
      return null;
    }
  }

  thrown.expectMessage(
      "Splittable, either @NewTracker method must be defined or OtherRestriction must implement HasDefaultTracker.");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@Test
public void testHasDefaultTracker() throws Exception {
  class Fn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext c, RestrictionTracker<RestrictionWithDefaultTracker, Void> tracker) {}

    @GetInitialRestriction
    public RestrictionWithDefaultTracker getInitialRestriction(@Element Integer element) {
      return null;
    }
  }

  DoFnSignature signature = DoFnSignatures.getSignature(Fn.class);
  assertEquals(RestrictionTracker.class, signature.processElement().trackerT().getRawType());
}
 

@Test
public void testNewTrackerReturnsWrongType() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {}

    @NewTracker
    public void newTracker(@Restriction SomeRestriction restriction) {}

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element Integer element) {
      return null;
    }
  }

  thrown.expectMessage(
      "Returns void, but must return a subtype of RestrictionTracker<SomeRestriction, ?>");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@Test
public void testNewWatermarkEstimatorReturnsWrongType() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {}

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element Integer element) {
      return null;
    }

    @NewWatermarkEstimator
    public void newWatermarkEstimator() {}
  }

  thrown.expectMessage("Returns void, but must return a subtype of WatermarkEstimator<Void>");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@Test
public void testGetInitialRestrictionMismatchesNewTracker() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {}

    @NewTracker
    public SomeRestrictionTracker newTracker(@Restriction SomeRestriction restriction) {
      return null;
    }

    @GetInitialRestriction
    public String getInitialRestriction(@Element Integer element) {
      return null;
    }
  }

  thrown.expectMessage("but must return a subtype of RestrictionTracker<String, ?>");
  thrown.expectMessage("newTracker(SomeRestriction): Returns SomeRestrictionTracker");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@ProcessElement
public ProcessContinuation processElement(
    ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
  int[] blockStarts = {-1, 0, 12, 123, 1234, 12345, 34567, MAX_INDEX};
  int trueStart = snapToNextBlock((int) tracker.currentRestriction().getFrom(), blockStarts);
  for (int i = trueStart, numIterations = 1;
      tracker.tryClaim((long) blockStarts[i]);
      ++i, ++numIterations) {
    for (int index = blockStarts[i]; index < blockStarts[i + 1]; ++index) {
      c.output(index);
    }
    if (numIterations == numClaimsPerCall) {
      return resume();
    }
  }
  return stop();
}
 

@ProcessElement
public ProcessContinuation processElement(
    ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
  int[] blockStarts = {-1, 0, 12, 123, 1234, 12345, 34567, MAX_INDEX};
  int trueStart = snapToNextBlock((int) tracker.currentRestriction().getFrom(), blockStarts);
  for (int i = trueStart, numIterations = 1;
      tracker.tryClaim((long) blockStarts[i]);
      ++i, ++numIterations) {
    for (int index = blockStarts[i]; index < blockStarts[i + 1]; ++index) {
      c.output(KV.of(c.sideInput(sideInput) + ":" + c.element(), index));
    }
    if (numIterations == numClaimsPerCall) {
      return resume();
    }
  }
  return stop();
}
 

@ProcessElement
public ProcessContinuation process(
    @Element String element,
    OutputReceiver<String> receiver,
    RestrictionTracker<OffsetRange, Long> tracker,
    BundleFinalizer bundleFinalizer)
    throws InterruptedException {
  if (wasFinalized.get()) {
    // Claim beyond the end now that we know we have been finalized.
    tracker.tryClaim(Long.MAX_VALUE);
    receiver.output(element);
    return stop();
  }
  if (tracker.tryClaim(tracker.currentRestriction().getFrom() + 1)) {
    bundleFinalizer.afterBundleCommit(
        Instant.now().plus(Duration.standardSeconds(MAX_ATTEMPTS)),
        () -> wasFinalized.set(true));
    // We sleep here instead of setting a resume time since the resume time doesn't need to
    // be honored.
    sleep(1000L); // 1 second
    return resume();
  }
  return stop();
}
 

@ProcessElement
public void processElement(
    @Element Read read,
    OutputReceiver<Result> out,
    RestrictionTracker<ByteKeyRange, ByteKey> tracker)
    throws Exception {
  Connection connection = ConnectionFactory.createConnection(read.getConfiguration());
  TableName tableName = TableName.valueOf(read.getTableId());
  Table table = connection.getTable(tableName);
  final ByteKeyRange range = tracker.currentRestriction();
  try (ResultScanner scanner =
      table.getScanner(HBaseUtils.newScanInRange(read.getScan(), range))) {
    for (Result result : scanner) {
      ByteKey key = ByteKey.copyFrom(result.getRow());
      if (!tracker.tryClaim(key)) {
        return;
      }
      out.output(result);
    }
    tracker.tryClaim(ByteKey.EMPTY);
  }
}
 

@ProcessElement
public ProcessContinuation process(
    ProcessContext context, RestrictionTracker<OffsetRange, Long> tracker) {
  Uninterruptibles.sleepUninterruptibly(
      sleepBeforeFirstClaim.getMillis(), TimeUnit.MILLISECONDS);
  for (long i = tracker.currentRestriction().getFrom(), numIterations = 1;
      tracker.tryClaim(i);
      ++i, ++numIterations) {
    Uninterruptibles.sleepUninterruptibly(
        sleepBeforeEachOutput.getMillis(), TimeUnit.MILLISECONDS);
    context.output("" + i);
    if (numIterations == numOutputsPerProcessCall) {
      return resume();
    }
  }
  return stop();
}
 

@Test
public void testInvokeProcessElementOutputDisallowedBeforeTryClaim() throws Exception {
  DoFn<Void, String> brokenFn =
      new DoFn<Void, String>() {
        @ProcessElement
        public void process(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
          c.output("foo");
        }

        @GetInitialRestriction
        public OffsetRange getInitialRestriction(@Element Void element) {
          throw new UnsupportedOperationException("Should not be called in this test");
        }
      };
  e.expectMessage("Output is not allowed before tryClaim()");
  runTest(brokenFn, new OffsetRange(0, 5));
}
 

@Test
public void testHasDefaultWatermarkEstimator() throws Exception {
  class Fn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext c,
        RestrictionTracker<SomeRestriction, Void> tracker,
        WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator>
            watermarkEstimator) {}

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element Integer element) {
      return null;
    }

    @GetInitialWatermarkEstimatorState
    public WatermarkEstimatorStateWithDefaultWatermarkEstimator
        getInitialWatermarkEstimatorState() {
      return null;
    }
  }

  DoFnSignature signature = DoFnSignatures.getSignature(Fn.class);
  assertEquals(
      WatermarkEstimator.class, signature.processElement().watermarkEstimatorT().getRawType());
}
 

@Test
public void testGetWatermarkEstimatorStateCoderReturnsWrongType() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {}

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element Integer element) {
      return null;
    }

    @GetWatermarkEstimatorStateCoder
    public KvCoder getWatermarkEstimatorStateCoder() {
      return null;
    }
  }

  thrown.expectMessage(
      "getWatermarkEstimatorStateCoder() returns KvCoder which is not a subtype of Coder<Void>");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@Test
public void testGetRestrictionCoderReturnsWrongType() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {}

    @NewTracker
    public SomeRestrictionTracker newTracker(@Restriction SomeRestriction restriction) {
      return null;
    }

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element Integer element) {
      return null;
    }

    @GetRestrictionCoder
    public KvCoder getRestrictionCoder() {
      return null;
    }
  }

  thrown.expectMessage(
      "getRestrictionCoder() returns KvCoder which is not a subtype of Coder<SomeRestriction>");
  DoFnSignatures.getSignature(BadFn.class);
}
 

/**
 * Returns a thread safe {@link RestrictionTracker} which reports all claim attempts to the
 * specified {@link ClaimObserver}.
 */
public static <RestrictionT, PositionT> RestrictionTracker<RestrictionT, PositionT> observe(
    RestrictionTracker<RestrictionT, PositionT> restrictionTracker,
    ClaimObserver<PositionT> claimObserver) {
  if (restrictionTracker instanceof RestrictionTracker.HasProgress) {
    return new RestrictionTrackerObserverWithProgress<>(restrictionTracker, claimObserver);
  } else {
    return new RestrictionTrackerObserver<>(restrictionTracker, claimObserver);
  }
}
 

@NewTracker
public RestrictionTracker<OffsetRange, Long> newTracker(@Restriction OffsetRange restriction) {
  // Use a modified OffsetRangeTracker with only support for checkpointing.
  return new OffsetRangeTracker(restriction) {
    @Override
    public SplitResult<OffsetRange> trySplit(double fractionOfRemainder) {
      return super.trySplit(0);
    }
  };
}
 

/**
 * List messages.
 *
 * @param hl7v2Store the HL7v2 store to list messages from
 * @throws IOException the io exception
 */
@ProcessElement
public void listMessages(
    @Element String hl7v2Store,
    RestrictionTracker<OffsetRange, Long> tracker,
    OutputReceiver<HL7v2Message> outputReceiver)
    throws IOException {
  OffsetRange currentRestriction = (OffsetRange) tracker.currentRestriction();
  Instant startRestriction = Instant.ofEpochMilli(currentRestriction.getFrom());
  Instant endRestriction = Instant.ofEpochMilli(currentRestriction.getTo());
  HttpHealthcareApiClient.HL7v2MessagePages pages =
      new HttpHealthcareApiClient.HL7v2MessagePages(
          client, hl7v2Store, startRestriction, endRestriction, filter.get(), "sendTime");
  Instant cursor;
  long lastClaimedMilliSecond = startRestriction.getMillis() - 1;
  for (HL7v2Message msg : FluentIterable.concat(pages)) {
    cursor = Instant.parse(msg.getSendTime());
    if (cursor.getMillis() > lastClaimedMilliSecond) {
      // Return early after the first claim failure preventing us from iterating
      // through the remaining messages.
      if (!tracker.tryClaim(cursor.getMillis())) {
        return;
      }
      lastClaimedMilliSecond = cursor.getMillis();
    }

    outputReceiver.output(msg);
  }

  // We've paginated through all messages for this restriction but the last message may be
  // before the end of the restriction
  tracker.tryClaim(currentRestriction.getTo());
}
 

@ProcessElement
public ProcessContinuation process(
    ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
  for (long i = tracker.currentRestriction().getFrom(), numIterations = 0;
      tracker.tryClaim(i);
      ++i, ++numIterations) {
    c.output(KV.of(c.element(), (int) i));
    if (numIterations % 3 == 0) {
      return resume();
    }
  }
  return stop();
}
 

@ProcessElement
public ProcessContinuation processElement(
    ProcessContext context,
    RestrictionTracker<OffsetRange, Long> tracker,
    ManualWatermarkEstimator<Instant> watermarkEstimator)
    throws Exception {
  long checkpointUpperBound = CHECKPOINT_UPPER_BOUND;
  long position = tracker.currentRestriction().getFrom();
  boolean claimStatus;
  while (true) {
    claimStatus = (tracker.tryClaim(position));
    if (!claimStatus) {
      break;
    } else if (position == SPLIT_ELEMENT) {
      enableAndWaitForTrySplitToHappen();
    }
    context.outputWithTimestamp(
        context.element() + ":" + position, GlobalWindow.TIMESTAMP_MIN_VALUE.plus(position));
    watermarkEstimator.setWatermark(GlobalWindow.TIMESTAMP_MIN_VALUE.plus(position));
    position += 1L;
    if (position == checkpointUpperBound) {
      break;
    }
  }
  if (!claimStatus) {
    return ProcessContinuation.stop();
  } else {
    return ProcessContinuation.resume().withResumeDelay(Duration.millis(54321L));
  }
}
 

@Test
public void
    testWatermarkEstimatorStateHasDefaultWatermarkEstimatorProcessUsesWrongWatermarkEstimator()
        throws Exception {
  class Fn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext c,
        RestrictionTracker<SomeRestriction, Void> tracker,
        SomeDefaultWatermarkEstimator watermarkEstimator) {}

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(@Element Integer element) {
      return null;
    }

    @GetInitialWatermarkEstimatorState
    public WatermarkEstimatorStateWithDefaultWatermarkEstimator
        getInitialWatermarkEstimatorState() {
      return null;
    }
  }

  thrown.expectMessage(
      "Has watermark estimator type SomeDefaultWatermarkEstimator, but the DoFn's watermark estimator type must be one of [WatermarkEstimator, ManualWatermarkEstimator] types.");
  DoFnSignature signature = DoFnSignatures.getSignature(Fn.class);
}
 

@ProcessElement
public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
  // create the channel
  String fileName = c.element().getKey();
  try (SeekableByteChannel channel = getReader(c.element().getValue())) {
    ByteBuffer readBuffer = ByteBuffer.allocate(BATCH_SIZE);
    ByteString buffer = ByteString.EMPTY;
    for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
      long startOffset = (i * BATCH_SIZE) - BATCH_SIZE;
      channel.position(startOffset);
      readBuffer = ByteBuffer.allocate(BATCH_SIZE);
      buffer = ByteString.EMPTY;
      channel.read(readBuffer);
      readBuffer.flip();
      buffer = ByteString.copyFrom(readBuffer);
      readBuffer.clear();
      LOG.debug(
          "Current Restriction {}, Content Size{}",
          tracker.currentRestriction(),
          buffer.size());
      c.output(KV.of(fileName, buffer.toStringUtf8().trim()));
    }
  } catch (Exception e) {

    c.output(textReaderFailedElements, e.getMessage());
  }
}
 

public ProcessContext(
    WindowedValue<InputT> element,
    RestrictionTracker<RestrictionT, PositionT> tracker,
    WatermarkEstimator<WatermarkEstimatorStateT> watermarkEstimator) {
  fn.super();
  this.element = element;
  this.tracker = RestrictionTrackers.observe(tracker, this);
  this.watermarkEstimator = WatermarkEstimators.threadSafe(watermarkEstimator);
}
 

@Test
public void testSplittableMissingGetInitialRestrictionMethod() throws Exception {
  class BadFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(
        ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {}
  }

  thrown.expectMessage(
      "Splittable, but does not define the required @GetInitialRestriction method.");
  DoFnSignatures.getSignature(BadFn.class);
}
 

@ProcessElement
public void process(ProcessContext c, RestrictionTracker<SomeRestriction, Void> tracker) {
  checkState(tracker.tryClaim(null));
  c.output(c.element().toString() + "a");
  c.output(c.element().toString() + "b");
  c.output(c.element().toString() + "c");
}
 

@ProcessElement
public void process(
    ProcessContext c,
    RestrictionTracker<OffsetRange, Long> tracker,
    ManualWatermarkEstimator<Instant> watermarkEstimator) {
  for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
    watermarkEstimator.setWatermark(c.element().plus(Duration.standardSeconds(i)));
    c.output(String.valueOf(i));
  }
}
 

private Progress getProgress() {
  synchronized (splitLock) {
    if (currentTracker instanceof RestrictionTracker.HasProgress) {
      return ((HasProgress) currentTracker).getProgress();
    }
  }
  return null;
}