Java Code Examples for org.apache.beam.sdk.values.PCollectionList#empty()

@Override
public PCollectionList<T> expand(PCollection<T> in) {
  final TupleTagList outputTags = partitionDoFn.getOutputTags();

  PCollectionTuple outputs =
      in.apply(
          ParDo.of(partitionDoFn)
              .withOutputTags(new TupleTag<Void>() {}, outputTags)
              .withSideInputs(partitionDoFn.getSideInputs()));

  PCollectionList<T> pcs = PCollectionList.empty(in.getPipeline());
  Coder<T> coder = in.getCoder();

  for (TupleTag<?> outputTag : outputTags.getAll()) {
    // All the tuple tags are actually TupleTag<T>
    // And all the collections are actually PCollection<T>
    @SuppressWarnings("unchecked")
    TupleTag<T> typedOutputTag = (TupleTag<T>) outputTag;
    pcs = pcs.and(outputs.get(typedOutputTag).setCoder(coder));
  }
  return pcs;
}
 

@Override
public PCollection<T> expand(PCollectionList<T> input) {
  Map<PCollection<T>, Integer> instances = new HashMap<>();
  for (PCollection<T> pCollection : input.getAll()) {
    int existing = instances.get(pCollection) == null ? 0 : instances.get(pCollection);
    instances.put(pCollection, existing + 1);
  }
  PCollectionList<T> output = PCollectionList.empty(input.getPipeline());
  for (Map.Entry<PCollection<T>, Integer> instanceEntry : instances.entrySet()) {
    if (instanceEntry.getValue().equals(1)) {
      output = output.and(instanceEntry.getKey());
    } else {
      String duplicationName = String.format("Multiply %s", instanceEntry.getKey().getName());
      PCollection<T> duplicated =
          instanceEntry
              .getKey()
              .apply(duplicationName, ParDo.of(new DuplicateFn<>(instanceEntry.getValue())));
      output = output.and(duplicated);
    }
  }
  return output.apply(Flatten.pCollections());
}
 

@Test
@Category(NeedsRunner.class)
public void testOverride() {
  PCollectionList<Long> empty = PCollectionList.empty(pipeline);
  PCollection<Long> emptyFlattened =
      empty.apply(
          factory
              .getReplacementTransform(
                  AppliedPTransform.of(
                      "nonEmptyInput",
                      Collections.emptyMap(),
                      Collections.emptyMap(),
                      Flatten.pCollections(),
                      pipeline))
              .getTransform());
  PAssert.that(emptyFlattened).empty();
  pipeline.run();
}
 

@Override
public PCollection<KV<URI, String>> expand(PBegin input) {
  Pipeline pipeline = input.getPipeline();

  // Create one TextIO.Read transform for each document
  // and add its output to a PCollectionList
  PCollectionList<KV<URI, String>> urisToLines = PCollectionList.empty(pipeline);

  // TextIO.Read supports:
  //  - file: URIs and paths locally
  //  - gs: URIs on the service
  for (final URI uri : uris) {
    String uriString;
    if ("file".equals(uri.getScheme())) {
      uriString = new File(uri).getPath();
    } else {
      uriString = uri.toString();
    }

    PCollection<KV<URI, String>> oneUriToLines =
        pipeline
            .apply("TextIO.Read(" + uriString + ")", TextIO.read().from(uriString))
            .apply("WithKeys(" + uriString + ")", WithKeys.of(uri))
            .setCoder(KvCoder.of(StringDelegateCoder.of(URI.class), StringUtf8Coder.of()));

    urisToLines = urisToLines.and(oneUriToLines);
  }

  return urisToLines.apply(Flatten.pCollections());
}
 

/** Transforms the inputs into a PInput. */
private static PCollectionList<Row> buildPCollectionList(
    List<RelNode> inputRels, Pipeline pipeline, Map<Integer, PCollection<Row>> cache) {
  if (inputRels.isEmpty()) {
    return PCollectionList.empty(pipeline);
  } else {
    return PCollectionList.of(
        inputRels.stream()
            .map(input -> BeamSqlRelUtils.toPCollection(pipeline, (BeamRelNode) input, cache))
            .collect(Collectors.toList()));
  }
}
 

@Override
public PCollectionList<Long> createInput(Pipeline p, Map<String, PCollection<?>> inputs) {
  PCollectionList<Long> inputList = PCollectionList.empty(p);
  for (PCollection<?> collection : inputs.values()) {
    inputList = inputList.and((PCollection<Long>) collection);
  }
  return inputList;
}
 

@Test
@Category({
  ValidatesRunner.class,
  UsesStatefulParDo.class, // This test fails if State is unsupported despite no direct usage.
  DataflowPortabilityApiUnsupported.class
})
public void testUnboundedSourceSplits() throws Exception {
  long numElements = 1000;
  int numSplits = 10;

  UnboundedSource<Long, ?> initial = CountingSource.unbounded();
  List<? extends UnboundedSource<Long, ?>> splits = initial.split(numSplits, p.getOptions());
  assertEquals("Expected exact splitting", numSplits, splits.size());

  long elementsPerSplit = numElements / numSplits;
  assertEquals("Expected even splits", numElements, elementsPerSplit * numSplits);
  PCollectionList<Long> pcollections = PCollectionList.empty(p);
  for (int i = 0; i < splits.size(); ++i) {
    pcollections =
        pcollections.and(
            p.apply("split" + i, Read.from(splits.get(i)).withMaxNumRecords(elementsPerSplit)));
  }
  PCollection<Long> input = pcollections.apply(Flatten.pCollections());

  addCountingAsserts(input, numElements);
  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testUnboundedSourceRateSplits() throws Exception {
  int elementsPerPeriod = 10;
  Duration period = Duration.millis(5);

  long numElements = 1000;
  int numSplits = 10;

  UnboundedCountingSource initial =
      CountingSource.createUnboundedFrom(0).withRate(elementsPerPeriod, period);
  List<? extends UnboundedSource<Long, ?>> splits = initial.split(numSplits, p.getOptions());
  assertEquals("Expected exact splitting", numSplits, splits.size());

  long elementsPerSplit = numElements / numSplits;
  assertEquals("Expected even splits", numElements, elementsPerSplit * numSplits);
  PCollectionList<Long> pcollections = PCollectionList.empty(p);
  for (int i = 0; i < splits.size(); ++i) {
    pcollections =
        pcollections.and(
            p.apply("split" + i, Read.from(splits.get(i)).withMaxNumRecords(elementsPerSplit)));
  }
  PCollection<Long> input = pcollections.apply(Flatten.pCollections());

  addCountingAsserts(input, numElements);
  Instant startTime = Instant.now();
  p.run();
  Instant endTime = Instant.now();
  // 500 ms if the readers are all initialized in parallel; 5000 ms if they are evaluated serially
  long expectedMinimumMillis = (numElements * period.getMillis()) / elementsPerPeriod;
  assertThat(expectedMinimumMillis, lessThan(endTime.getMillis() - startTime.getMillis()));
}
 

@Test
public void testUnboundedSourceSplits() throws Exception {

  int numElements = 1000;
  int numSplits = 10;

  // Coders must be specified explicitly here due to the way the transform
  // is used in the test.
  UnboundedSource<KafkaRecord<Integer, Long>, ?> initial =
      mkKafkaReadTransform(numElements, null)
          .withKeyDeserializerAndCoder(IntegerDeserializer.class, BigEndianIntegerCoder.of())
          .withValueDeserializerAndCoder(LongDeserializer.class, BigEndianLongCoder.of())
          .makeSource();

  List<? extends UnboundedSource<KafkaRecord<Integer, Long>, ?>> splits =
      initial.split(numSplits, p.getOptions());
  assertEquals("Expected exact splitting", numSplits, splits.size());

  long elementsPerSplit = numElements / numSplits;
  assertEquals("Expected even splits", numElements, elementsPerSplit * numSplits);
  PCollectionList<Long> pcollections = PCollectionList.empty(p);
  for (int i = 0; i < splits.size(); ++i) {
    pcollections =
        pcollections.and(
            p.apply("split" + i, Read.from(splits.get(i)).withMaxNumRecords(elementsPerSplit))
                .apply("Remove Metadata " + i, ParDo.of(new RemoveKafkaMetadata<>()))
                .apply("collection " + i, Values.create()));
  }
  PCollection<Long> input = pcollections.apply(Flatten.pCollections());

  addCountingAsserts(input, numElements);
  p.run();
}
 

/**
 * {@inheritDoc}.
 *
 * <p>The input {@link PCollectionList} that is constructed will have the same values in the same
 */
private PCollectionList<T> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
  PCollectionList<T> pCollections = PCollectionList.empty(p);
  for (PValue input : inputs.values()) {
    PCollection<T> pcollection = (PCollection<T>) input;
    pCollections = pCollections.and(pcollection);
  }
  return pCollections;
}
 

@Test
public void testFlattenInMemoryEvaluatorWithEmptyPCollectionList() throws Exception {
  PCollectionList<Integer> list = PCollectionList.empty(p);

  PCollection<Integer> flattened = list.apply(Flatten.pCollections());
  flattened.setCoder(VarIntCoder.of());

  EvaluationContext evaluationContext = mock(EvaluationContext.class);
  when(evaluationContext.createBundle(flattened))
      .thenReturn(bundleFactory.createBundle(flattened));

  FlattenEvaluatorFactory factory = new FlattenEvaluatorFactory(evaluationContext);
  AppliedPTransform<?, ?, ?> flattendProducer = DirectGraphs.getProducer(flattened);
  TransformEvaluator<Integer> emptyEvaluator =
      factory.forApplication(
          flattendProducer,
          bundleFactory.createRootBundle().commit(BoundedWindow.TIMESTAMP_MAX_VALUE));

  TransformResult<Integer> leftSideResult = emptyEvaluator.finishBundle();

  CommittedBundle<?> outputBundle =
      Iterables.getOnlyElement(leftSideResult.getOutputBundles()).commit(Instant.now());
  assertThat(outputBundle.getElements(), emptyIterable());
  assertThat(
      leftSideResult.getTransform(),
      Matchers.<AppliedPTransform<?, ?, ?>>equalTo(flattendProducer));
}
 

@Override
public PCollection<T> expand(PBegin input) {
  PCollectionList<T> empty = PCollectionList.empty(input.getPipeline());
  return empty.apply(Flatten.pCollections());
}