Java Code Examples for org.apache.beam.sdk.values.TimestampedValue#of()

/**
 * Verifies that later-ending merged windows from any of the timestamps hold up output of
 * earlier-ending windows, using the provided {@link WindowFn} and {@link TimestampCombiner}.
 *
 * <p>Given a list of lists of timestamps, where each list is expected to merge into a single
 * window with end times in ascending order, assigns and merges windows for each list (as though
 * each were a separate key/user session). Then combines each timestamp in the list according to
 * the provided {@link TimestampCombiner}.
 *
 * <p>Verifies that a overlapping windows do not hold each other up via the watermark.
 */
public static <T, W extends IntervalWindow> void validateGetOutputTimestamps(
    WindowFn<T, W> windowFn,
    TimestampCombiner timestampCombiner,
    List<List<Long>> timestampsPerWindow)
    throws Exception {

  List<List<TimestampedValue<T>>> timestampValuesPerWindow = new ArrayList<>();
  for (List<Long> timestamps : timestampsPerWindow) {
    List<TimestampedValue<T>> timestampedValues = new ArrayList<>();
    for (Long timestamp : timestamps) {
      TimestampedValue<T> tv = TimestampedValue.of(null, new Instant(timestamp));
      timestampedValues.add(tv);
    }
    timestampValuesPerWindow.add(timestampedValues);
  }
  validateGetOutputTimestampsWithValue(windowFn, timestampCombiner, timestampValuesPerWindow);
}
 

/** Update the per-seller running counts/sums. */
private void captureWinningBid(Auction auction, Bid bid, Instant timestamp) {
  NexmarkUtils.info("winning auction, bid: %s, %s", auction, bid);
  Queue<Bid> queue = winningBidsPerSeller.get(auction.seller);
  if (queue == null) {
    queue = new PriorityQueue<>(10, (Bid b1, Bid b2) -> b1.dateTime.compareTo(b2.dateTime));
  }
  Long total = totalWinningBidPricesPerSeller.get(auction.seller);
  if (total == null) {
    total = 0L;
  }
  int count = queue.size();
  if (count == 10) {
    total -= queue.remove().price;
  } else {
    count += 1;
  }
  queue.add(bid);
  total += bid.price;
  winningBidsPerSeller.put(auction.seller, queue);
  totalWinningBidPricesPerSeller.put(auction.seller, total);
  TimestampedValue<SellerPrice> intermediateResult =
      TimestampedValue.of(
          new SellerPrice(auction.seller, Math.round((double) total / count)), timestamp);
  addIntermediateResult(intermediateResult);
}
 

private void flushSession(long bidder) {
  List<TimestampedValue<Event>> session = activeSessions.get(bidder);
  checkState(session != null);

  Instant sessionStart =
      Ordering.<Instant>natural()
          .min(
              session.stream()
                  .<Instant>map(tsv -> tsv.getTimestamp())
                  .collect(Collectors.toList()));

  Instant sessionEnd =
      Ordering.<Instant>natural()
          .max(
              session.stream()
                  .<Instant>map(tsv -> tsv.getTimestamp())
                  .collect(Collectors.toList()))
          .plus(configuration.sessionGap);

  for (TimestampedValue<Event> timestampedEvent : session) {
    // Join to the side input is always a string representation of the id being looked up
    Bid bid = timestampedEvent.getValue().bid;
    Bid resultBid =
        new Bid(
            bid.auction,
            bid.bidder,
            bid.price,
            bid.dateTime,
            String.format(
                "%d:%s:%s",
                bid.bidder % configuration.sideInputRowCount, sessionStart, sessionEnd));
    TimestampedValue<Bid> result =
        TimestampedValue.of(resultBid, timestampedEvent.getTimestamp());
    addResult(result);
  }
  activeSessions.remove(bidder);
}
 

@Test
public void processElementWithOutputTimestamp() throws Exception {
  try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) {
    tester.processElement(1L);
    tester.processElement(2L);

    List<TimestampedValue<String>> peek = tester.peekOutputElementsWithTimestamp();
    TimestampedValue<String> one = TimestampedValue.of("1", new Instant(1000L));
    TimestampedValue<String> two = TimestampedValue.of("2", new Instant(2000L));
    assertThat(peek, hasItems(one, two));

    tester.processElement(3L);
    tester.processElement(4L);

    TimestampedValue<String> three = TimestampedValue.of("3", new Instant(3000L));
    TimestampedValue<String> four = TimestampedValue.of("4", new Instant(4000L));
    peek = tester.peekOutputElementsWithTimestamp();
    assertThat(peek, hasItems(one, two, three, four));
    List<TimestampedValue<String>> take = tester.takeOutputElementsWithTimestamp();
    assertThat(take, hasItems(one, two, three, four));

    // Following takeOutputElementsWithTimestamp(), neither takeOutputElementsWithTimestamp()
    // nor peekOutputElementsWithTimestamp() return anything.
    assertTrue(tester.takeOutputElementsWithTimestamp().isEmpty());
    assertTrue(tester.peekOutputElementsWithTimestamp().isEmpty());

    // peekOutputElements() and takeOutputElements() also return nothing.
    assertTrue(tester.peekOutputElements().isEmpty());
    assertTrue(tester.takeOutputElements().isEmpty());
  }
}
 

@Test
public void processTimestampedElement() throws Exception {
  try (DoFnTester<Long, TimestampedValue<Long>> tester = DoFnTester.of(new ReifyTimestamps())) {
    TimestampedValue<Long> input = TimestampedValue.of(1L, new Instant(100));
    tester.processTimestampedElement(input);
    assertThat(tester.takeOutputElements(), contains(input));
  }
}
 

@Test
public void testAddInputSameTimestamp() {
  TimestampedValue<Long> accum = TimestampedValue.of(100L, INSTANT);
  TimestampedValue<Long> input = TimestampedValue.of(200L, INSTANT);

  assertThat(
      "Latest for values with the same timestamp is chosen arbitrarily",
      fn.addInput(accum, input),
      isOneOf(accum, input));
}
 

/**
 * Adds the specified elements to the source with timestamp equal to the current watermark.
 *
 * @return A {@link TestStream.Builder} like this one that will add the provided elements after
 *     all earlier events have completed.
 */
@SafeVarargs
public final Builder<T> addElements(T element, T... elements) {
  TimestampedValue<T> firstElement = TimestampedValue.of(element, currentWatermark);
  @SuppressWarnings({"unchecked", "rawtypes"})
  TimestampedValue<T>[] remainingElements = new TimestampedValue[elements.length];
  for (int i = 0; i < elements.length; i++) {
    remainingElements[i] = TimestampedValue.of(elements[i], currentWatermark);
  }
  return addElements(firstElement, remainingElements);
}
 

private static TimestampedValue<KV<String, VisitOrImpression>> visitOrImpression(Visit visit, Impression impression) {
    checkArgument((visit == null) != (impression == null), "Either visit or impression must be null");
    VisitOrImpression voi = new VisitOrImpression(visit, impression);
    // Note that for this example we always use the same user, since the test
    // only looks at one user's data.
    return TimestampedValue.of(KV.of("UserX", voi), visit != null ? visit.timestamp() : impression.timestamp());
}
 

private TimestampedValue<GameActionInfo> event(
    TestUser user, int score, Duration baseTimeOffset) {
  return TimestampedValue.of(
      new GameActionInfo(
          user.getUser(), user.getTeam(), score, baseTime.plus(baseTimeOffset).getMillis()),
      baseTime.plus(baseTimeOffset));
}
 

private TimestampedValue<KV<String, GameActionInfo>> event(
    TestUser user, int score, Duration baseTimeOffset) {
  return TimestampedValue.of(
      KV.of(
          user.getTeam(),
          new GameActionInfo(
              user.getUser(), user.getTeam(), score, baseTime.plus(baseTimeOffset).getMillis())),
      baseTime.plus(baseTimeOffset));
}
 

private TimestampedValue<GameActionInfo> event(
    TestUser user, int score, Duration baseTimeOffset) {
  return TimestampedValue.of(
      new GameActionInfo(
          user.getUser(), user.getTeam(), score, baseTime.plus(baseTimeOffset).getMillis()),
      baseTime.plus(baseTimeOffset));
}
 

@Test
public void testExtractOutputNullValue() {
  TimestampedValue<Long> accum = TimestampedValue.of(null, baseTimestamp);
  assertEquals(null, fn.extractOutput(accum));
}
 

@Override
public boolean advance() {
  long now = System.currentTimeMillis();

  while (pendingEvent == null) {
    if (!generator.hasNext() && heldBackEvents.isEmpty()) {
      // No more events, EVER.
      if (isRateLimited) {
        updateBacklog(System.currentTimeMillis(), 0);
      }
      if (watermark < BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
        watermark = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis();
        LOG.trace("stopped unbounded generator {}", generator);
      }
      return false;
    }

    Generator.NextEvent next = heldBackEvents.peek();
    if (next != null && next.wallclockTimestamp <= now) {
      // Time to use the held-back event.
      heldBackEvents.poll();
      LOG.debug(
          "replaying held-back event {}ms behind watermark", watermark - next.eventTimestamp);
    } else if (generator.hasNext()) {
      next = generator.nextEvent();
      if (isRateLimited
          && config.getProbDelayedEvent() > 0.0
          && config.getOccasionalDelaySec() > 0
          && ThreadLocalRandom.current().nextDouble() < config.getProbDelayedEvent()) {
        // We'll hold back this event and go around again.
        long delayMs =
            ThreadLocalRandom.current().nextLong(config.getOccasionalDelaySec() * 1000) + 1L;
        LOG.debug("delaying event by {}ms", delayMs);
        heldBackEvents.add(next.withDelay(delayMs));
        continue;
      }
    } else {
      // Waiting for held-back event to fire.
      if (isRateLimited) {
        updateBacklog(now, 0);
      }
      return false;
    }

    pendingEventWallclockTime = next.wallclockTimestamp;
    pendingEvent = TimestampedValue.of(next.event, new Instant(next.eventTimestamp));
    long newWatermark =
        next.watermark - Duration.standardSeconds(watermarkHoldbackSec).getMillis();
    if (newWatermark > watermark) {
      watermark = newWatermark;
    }
  }

  if (isRateLimited) {
    if (pendingEventWallclockTime > now) {
      // We want this event to fire in the future. Try again later.
      updateBacklog(now, 0);
      return false;
    }
    updateBacklog(now, now - pendingEventWallclockTime);
  }

  // This event is ready to fire.
  currentEvent = pendingEvent;
  pendingEvent = null;
  return true;
}
 

/**
 * Test successful conversion of {@link FailsafeElement} records into {@link PubsubMessage} with
 * attributes.
 */
@Test
@Category(NeedsRunner.class)
public void testFailedStringToPubsubMessageFn() {

  String testMessage = "original-test-message";
  FailsafeElement<String, String> element = FailsafeElement.of(testMessage, testMessage);

  Instant expectedTimestamp = Instant.now();

  String errorMessage = "my-error-message";
  element.setErrorMessage(errorMessage);

  TimestampedValue<FailsafeElement<String, String>> input =
      TimestampedValue.of(element, expectedTimestamp);

  PCollection<PubsubMessage> pCollection =
      pipeline
          .apply(
              Create.timestamped(input)
                  .withCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
          .apply(ParDo.of(new FailedStringToPubsubMessageFn()));

  String expectedTimestampString =
      FailedStringToPubsubMessageFn.TIMESTAMP_FORMATTER.print(
          expectedTimestamp.toDateTime(DateTimeZone.UTC));

  PAssert.that(pCollection)
      .satisfies(
          collection -> {
            PubsubMessage actual = collection.iterator().next();
            assertThat(
                new String(actual.getPayload(), StandardCharsets.UTF_8),
                is(equalTo(testMessage)));
            assertThat(
                actual.getAttribute(FailedStringToPubsubMessageFn.ERROR_MESSAGE),
                is(equalTo(errorMessage)));
            assertThat(
                actual.getAttribute(FailedStringToPubsubMessageFn.TIMESTAMP),
                is(equalTo(expectedTimestampString)));
            return null;
          });

  pipeline.run();
}
 

@Override
public boolean tryClaim(TimestampedValue<T>[] position) {
  if (claimedAll) {
    return false;
  }
  try {
    if (currentReader == null) {
      currentReader = initialRestriction.createReader(pipelineOptions);
      if (!currentReader.start()) {
        claimedAll = true;
        try {
          currentReader.close();
        } finally {
          currentReader = null;
        }
        return false;
      }
      position[0] =
          TimestampedValue.of(
              currentReader.getCurrent(), currentReader.getCurrentTimestamp());
      return true;
    }
    if (!currentReader.advance()) {
      claimedAll = true;
      try {
        currentReader.close();
      } finally {
        currentReader = null;
      }
      return false;
    }
    position[0] =
        TimestampedValue.of(currentReader.getCurrent(), currentReader.getCurrentTimestamp());
    return true;
  } catch (IOException e) {
    if (currentReader != null) {
      try {
        currentReader.close();
      } catch (IOException closeException) {
        e.addSuppressed(closeException);
      } finally {
        currentReader = null;
      }
    }
    throw new RuntimeException(e);
  }
}
 

/** Helper method to easily create a timestamped value. */
private static TimestampedValue<Long> timestamped(Instant timestamp) {
  return TimestampedValue.of(uniqueLong.incrementAndGet(), timestamp);
}
 

@Test
@Category({ValidatesRunner.class, UsesSideInputs.class})
public void testWindowedCombineGloballyAsSingletonView() {
  FixedWindows windowFn = FixedWindows.of(Duration.standardMinutes(1));
  final PCollectionView<Integer> view =
      pipeline
          .apply(
              "CreateSideInput",
              Create.timestamped(
                  TimestampedValue.of(1, new Instant(100)),
                  TimestampedValue.of(3, new Instant(100))))
          .apply("WindowSideInput", Window.into(windowFn))
          .apply("CombineSideInput", Sum.integersGlobally().asSingletonView());

  TimestampedValue<Void> nonEmptyElement = TimestampedValue.of(null, new Instant(100));
  TimestampedValue<Void> emptyElement = TimestampedValue.atMinimumTimestamp(null);
  PCollection<Integer> output =
      pipeline
          .apply(
              "CreateMainInput",
              Create.timestamped(nonEmptyElement, emptyElement).withCoder(VoidCoder.of()))
          .apply("WindowMainInput", Window.into(windowFn))
          .apply(
              "OutputSideInput",
              ParDo.of(
                      new DoFn<Void, Integer>() {
                        @ProcessElement
                        public void processElement(ProcessContext c) {
                          c.output(c.sideInput(view));
                        }
                      })
                  .withSideInputs(view));

  PAssert.that(output).containsInAnyOrder(4, 0);
  PAssert.that(output)
      .inWindow(windowFn.assignWindow(nonEmptyElement.getTimestamp()))
      .containsInAnyOrder(4);
  PAssert.that(output)
      .inWindow(windowFn.assignWindow(emptyElement.getTimestamp()))
      .containsInAnyOrder(0);
  pipeline.run();
}
 

@Test
public void testAddInputNullValue() {
  TimestampedValue<Long> input = TimestampedValue.of(null, INSTANT.plus(10));
  assertEquals("Null values are allowed", input, fn.addInput(TV, input));
}
 

private TimestampedValue<PubsubMessage> message(String content, Duration baseTimeOffset) {
  return TimestampedValue.of(
      new PubsubMessage(content.getBytes(StandardCharsets.UTF_8), new HashMap<>()),
      baseTime.plus(baseTimeOffset));
}
 

private <T> TimestampedValue<T> event(T eventValue, Long timestamp) {

    return TimestampedValue.of(eventValue, START_TIME.plus(new Duration(timestamp)));
  }