org.apache.beam.sdk.values.ValueInSingleWindow Java Examples

public PCollectionSingletonIterableAssert(
    PCollection<Iterable<T>> actual,
    AssertionWindows rewindowingStrategy,
    SimpleFunction<Iterable<ValueInSingleWindow<Iterable<T>>>, Iterable<Iterable<T>>>
        paneExtractor,
    PAssertionSite site) {
  this.actual = actual;

  @SuppressWarnings("unchecked")
  Coder<T> typedCoder = (Coder<T>) actual.getCoder().getCoderArguments().get(0);
  this.elementCoder = typedCoder;

  this.rewindowingStrategy = rewindowingStrategy;
  this.paneExtractor = paneExtractor;
  this.site = site;
}
 

@Test
public void onlyPaneMultiplePanesFails() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.onlyPane(PAssert.PAssertionSite.capture(""));
  Iterable<ValueInSingleWindow<Integer>> multipleFiring =
      ImmutableList.of(
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)),
          ValueInSingleWindow.of(
              2,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              1,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)));

  thrown.expectMessage("trigger that fires at most once");
  extractor.apply(multipleFiring);
}
 

@Test
public void onTimePane() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.onTimePane();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              2,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)));

  assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4));
}
 

@Test
public void finalPane() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.finalPane();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              8,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, true, Timing.LATE, 2L, 1L)),
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              1,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)));

  assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(8));
}
 

@Test
public void finalPaneNoExplicitFinalEmpty() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.finalPane();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              8,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)),
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              1,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)));

  assertThat(extractor.apply(onlyOnTime), emptyIterable());
}
 

@Test
@Category({NeedsRunner.class, UsesTestStream.class})
public void globalWindowNoKeys() {
  PCollection<ValueInSingleWindow<String>> result =
      pipeline
          .apply(
              TestStream.create(StringUtf8Coder.of())
                  .addElements(TimestampedValue.of("dei", new Instant(123L)))
                  .advanceWatermarkToInfinity())
          .apply(Reify.windows());
  PAssert.that(result)
      .containsInAnyOrder(
          ValueInSingleWindow.of(
              "dei", new Instant(123L), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING));
  pipeline.run();
}
 

@Test
public void nonLatePanesSingleEarly() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.nonLatePanes();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              8,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)),
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)));

  assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(4, 8));
}
 

@Override
public PCollection<Iterable<ValueInSingleWindow<T>>> expand(PCollection<T> input) {
  WindowFn<?, ?> originalWindowFn = input.getWindowingStrategy().getWindowFn();

  return input
      .apply(Reify.windows())
      .apply(
          WithKeys.<Integer, ValueInSingleWindow<T>>of(0)
              .withKeyType(new TypeDescriptor<Integer>() {}))
      .apply(
          Window.into(
                  new IdentityWindowFn<KV<Integer, ValueInSingleWindow<T>>>(
                      originalWindowFn.windowCoder()))
              .triggering(Never.ever())
              .withAllowedLateness(input.getWindowingStrategy().getAllowedLateness())
              .discardingFiredPanes())
      // all values have the same key so they all appear as a single output element
      .apply(GroupByKey.create())
      .apply(Values.create())
      .setWindowingStrategyInternal(input.getWindowingStrategy());
}
 

public long insertAll(
    TableReference ref, List<TableRow> rowList, @Nullable List<String> insertIdList)
    throws IOException, InterruptedException {
  List<ValueInSingleWindow<TableRow>> windowedRows = Lists.newArrayList();
  for (TableRow row : rowList) {
    windowedRows.add(
        ValueInSingleWindow.of(
            row,
            GlobalWindow.TIMESTAMP_MAX_VALUE,
            GlobalWindow.INSTANCE,
            PaneInfo.ON_TIME_AND_ONLY_FIRING));
  }
  return insertAll(
      ref,
      windowedRows,
      insertIdList,
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
}
 

@Test
public void allPanesMultiplePanes() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.allPanes();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              8,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)),
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              1,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)));

  assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(4, 8, 1));
}
 

/**
 * input - a tupel that contains the data element (TableRow), the window, the timestamp, and the pane
 */

@Override
public TableDestination apply(ValueInSingleWindow<TableRow> input) {
    
	String partition;
	
	if (this.isTimeField) {
     String sTime = (String) input.getValue().get(this.fieldName);
     Instant time = Instant.parse(sTime);
     partition = time.toString(partitionFormatter);
	} else {
		partition = ((Integer) input.getValue().get(this.fieldName)).toString();
	}
	
    TableReference reference = new TableReference();
    reference.setProjectId(this.projectId);
    reference.setDatasetId(this.datasetId);
    reference.setTableId(this.partitionPrefix + partition);
    return new TableDestination(reference, null);
}
 

/** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
@ProcessElement
public void processElement(
    @Element KV<ShardedKey<String>, TableRowInfo<ElementT>> element,
    @Timestamp Instant timestamp,
    BoundedWindow window,
    PaneInfo pane) {
  String tableSpec = element.getKey().getKey();
  List<ValueInSingleWindow<TableRow>> rows =
      BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
  List<String> uniqueIds =
      BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);

  TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
  rows.add(ValueInSingleWindow.of(tableRow, timestamp, window, pane));
  uniqueIds.add(element.getValue().uniqueId);
}
 

@Override
public <T> long insertAll(
    TableReference ref,
    List<ValueInSingleWindow<TableRow>> rowList,
    @Nullable List<String> insertIdList,
    InsertRetryPolicy retryPolicy,
    List<ValueInSingleWindow<T>> failedInserts,
    ErrorContainer<T> errorContainer,
    boolean skipInvalidRows,
    boolean ignoreUnknownValues,
    boolean ignoreInsertIds)
    throws IOException, InterruptedException {
  return insertAll(
      ref,
      rowList,
      insertIdList,
      BackOffAdapter.toGcpBackOff(INSERT_BACKOFF_FACTORY.backoff()),
      Sleeper.DEFAULT,
      retryPolicy,
      failedInserts,
      errorContainer,
      skipInvalidRows,
      ignoreUnknownValues,
      ignoreInsertIds);
}
 

@Test
public void onTimePaneOnlyEarlyAndLate() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.onTimePane();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              8,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)),
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              2,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              1,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)));

  assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4));
}
 

@Test
public void testE2EBigQueryClusteringTableFunction() throws Exception {
  String tableName = "weather_stations_clustered_table_function_" + System.currentTimeMillis();

  Pipeline p = Pipeline.create(options);

  p.apply(BigQueryIO.readTableRows().from(options.getBqcInput()))
      .apply(ParDo.of(new KeepStationNumberAndConvertDate()))
      .apply(
          BigQueryIO.writeTableRows()
              .to(
                  (ValueInSingleWindow<TableRow> vsw) ->
                      new TableDestination(
                          String.format("%s.%s", DATASET_NAME, tableName),
                          null,
                          TIME_PARTITIONING,
                          CLUSTERING))
              .withClustering()
              .withSchema(SCHEMA)
              .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
              .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE));

  p.run().waitUntilFinish();

  Table table = bqClient.tables().get(options.getProject(), DATASET_NAME, tableName).execute();

  Assert.assertEquals(table.getClustering(), CLUSTERING);
}
 

@Override
public TableDestination getDestination(ValueInSingleWindow<T> element) {
  TableDestination res = tableFunction.apply(element);
  checkArgument(
      res != null,
      "result of tableFunction can not be null, but %s returned null for element: %s",
      tableFunction,
      element);
  return res;
}
 

@Test
public void allPanesSinglePane() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.allPanes();
  Iterable<ValueInSingleWindow<Integer>> onlyOnTime =
      ImmutableList.of(
          ValueInSingleWindow.of(
              8, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING),
          ValueInSingleWindow.of(
              4, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING),
          ValueInSingleWindow.of(
              2, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING));

  assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4, 8));
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries quota exceeded attempts. */
@Test
public void testInsertQuotaExceededRetry() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 quota exceeded, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("quotaExceeded", 403)))
      .thenReturn(toStream(new TableDataInsertAllResponse()));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

@Test
public void testDecodeEncodeEqual() throws Exception {
  Instant now = Instant.now();
  ValueInSingleWindow<String> value =
      ValueInSingleWindow.of(
          "foo",
          now,
          new IntervalWindow(now, now.plus(Duration.standardSeconds(10))),
          PaneInfo.NO_FIRING);

  CoderProperties.coderDecodeEncodeEqual(
      ValueInSingleWindow.Coder.of(StringUtf8Coder.of(), IntervalWindow.getCoder()), value);
}
 

@Test
public void allPanesEmpty() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.allPanes();
  Iterable<ValueInSingleWindow<Integer>> noPanes = ImmutableList.of();

  assertThat(extractor.apply(noPanes), emptyIterable());
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries rate limited attempts. */
@Test
public void testInsertRateLimitRetry() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 rate limited, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
      .thenReturn(toStream(new TableDataInsertAllResponse()));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

@Test
@Category(NeedsRunner.class)
public void multiplePanesMultipleReifiedPane() {
  PCollection<Long> someElems = p.apply("someLongs", GenerateSequence.from(0).to(20000));
  PCollection<Long> otherElems = p.apply("otherLongs", GenerateSequence.from(0).to(20000));
  PCollection<Iterable<ValueInSingleWindow<Iterable<Long>>>> accumulatedPanes =
      PCollectionList.of(someElems)
          .and(otherElems)
          .apply(Flatten.pCollections())
          .apply(WithTimestamps.of(input -> new Instant(input * 10)))
          .apply(
              Window.<Long>into(FixedWindows.of(Duration.standardMinutes(1)))
                  .triggering(
                      AfterWatermark.pastEndOfWindow()
                          .withEarlyFirings(AfterPane.elementCountAtLeast(1)))
                  .withAllowedLateness(Duration.ZERO)
                  .discardingFiredPanes())
          .apply(WithKeys.<Void, Long>of((Void) null).withKeyType(new TypeDescriptor<Void>() {}))
          .apply(GroupByKey.create())
          .apply(Values.create())
          .apply(GatherAllPanes.globally());

  PAssert.that(accumulatedPanes)
      .satisfies(
          input -> {
            for (Iterable<ValueInSingleWindow<Iterable<Long>>> windowedInput : input) {
              if (Iterables.size(windowedInput) > 1) {
                return null;
              }
            }
            fail("Expected at least one window to have multiple panes");
            return null;
          });

  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void singlePaneSingleReifiedPane() {
  PCollection<Iterable<ValueInSingleWindow<Iterable<Long>>>> accumulatedPanes =
      p.apply(GenerateSequence.from(0).to(20000))
          .apply(WithTimestamps.of(input -> new Instant(input * 10)))
          .apply(
              Window.<Long>into(FixedWindows.of(Duration.standardMinutes(1)))
                  .triggering(AfterWatermark.pastEndOfWindow())
                  .withAllowedLateness(Duration.ZERO)
                  .discardingFiredPanes())
          .apply(WithKeys.<Void, Long>of((Void) null).withKeyType(new TypeDescriptor<Void>() {}))
          .apply(GroupByKey.create())
          .apply(Values.create())
          .apply(GatherAllPanes.globally());

  PAssert.that(accumulatedPanes)
      .satisfies(
          input -> {
            for (Iterable<ValueInSingleWindow<Iterable<Long>>> windowedInput : input) {
              if (Iterables.size(windowedInput) > 1) {
                fail("Expected all windows to have exactly one pane, got " + windowedInput);
                return null;
              }
            }
            return null;
          });

  p.run();
}
 

private CreateActual(
    PCollection<T> actual,
    AssertionWindows rewindowActuals,
    SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> extractPane,
    PTransform<PCollection<T>, PCollectionView<ActualT>> actualView) {
  this.actual = actual;
  this.rewindowActuals = rewindowActuals;
  this.extractPane = extractPane;
  this.actualView = actualView;
}
 

PCollectionSingletonAssert(
    PCollection<T> actual,
    AssertionWindows rewindowingStrategy,
    SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> paneExtractor,
    PAssertionSite site) {
  this.actual = actual;
  this.coder = actual.getCoder();
  this.rewindowingStrategy = rewindowingStrategy;
  this.paneExtractor = paneExtractor;
  this.site = site;
}
 

@Test
@Category(NeedsRunner.class)
public void windowsInValueSucceeds() {
  PCollection<KV<String, Integer>> timestamped =
      pipeline
          .apply(Create.of(KV.of("foo", 0), KV.of("foo", 1), KV.of("bar", 2), KV.of("baz", 3)))
          .apply(TIMESTAMP_FROM_V);

  PCollection<KV<String, ValueInSingleWindow<Integer>>> reified =
      timestamped.apply(Reify.windowsInValue());

  PAssert.that(reified)
      .containsInAnyOrder(
          KV.of(
              "foo",
              ValueInSingleWindow.of(
                  0, new Instant(0), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)),
          KV.of(
              "foo",
              ValueInSingleWindow.of(
                  1, new Instant(1), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)),
          KV.of(
              "bar",
              ValueInSingleWindow.of(
                  2, new Instant(2), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)),
          KV.of(
              "baz",
              ValueInSingleWindow.of(
                  3, new Instant(3), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)));

  pipeline.run();
}
 

private PCollectionSingletonAssert<T> withPanes(
    BoundedWindow window,
    SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> paneExtractor) {
  @SuppressWarnings({"unchecked", "rawtypes"})
  Coder<BoundedWindow> windowCoder =
      (Coder) actual.getWindowingStrategy().getWindowFn().windowCoder();
  return new PCollectionSingletonAssert<>(
      actual, IntoStaticWindows.of(windowCoder, window), paneExtractor, site);
}
 

private PCollectionViewAssert(
    PCollection<ElemT> actual,
    PTransform<PCollection<ElemT>, PCollectionView<ViewT>> view,
    AssertionWindows rewindowActuals,
    SimpleFunction<Iterable<ValueInSingleWindow<ElemT>>, Iterable<ElemT>> paneExtractor,
    Coder<ViewT> coder,
    PAssertionSite site) {
  this.actual = actual;
  this.view = view;
  this.rewindowActuals = rewindowActuals;
  this.paneExtractor = paneExtractor;
  this.coder = coder;
  this.site = site;
}
 

@Override
public KV<String, TableRow> getDestination(ValueInSingleWindow<KV<String, TableRow>> element) {
  String key = element.getValue().getKey();
  String tableName = String.format("%s:%s.%s", projectId, datasetName, key);
  LOG.debug("Table Name {}", tableName);
  return KV.of(tableName, element.getValue().getValue());
}
 

private GroupThenAssert(
    SerializableFunction<Iterable<T>, Void> checkerFn,
    AssertionWindows rewindowingStrategy,
    SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> paneExtractor,
    PAssertionSite site) {
  this.checkerFn = checkerFn;
  this.rewindowingStrategy = rewindowingStrategy;
  this.paneExtractor = paneExtractor;
  this.site = site;
}