Java Code Examples for org.apache.beam.sdk.transforms.Create#Values

/** Tests that trying to register a transform which has unregistered children throws. */
@Test
public void registerTransformWithUnregisteredChildren() throws IOException {
  Create.Values<Long> create = Create.of(1L, 2L, 3L);
  GenerateSequence createChild = GenerateSequence.from(0);

  PCollection<Long> pt = pipeline.apply(create);
  String userName = "my_transform";
  String childUserName = "my_transform/my_nesting";
  AppliedPTransform<?, ?, ?> transform =
      AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline);
  AppliedPTransform<?, ?, ?> childTransform =
      AppliedPTransform.of(
          childUserName, pipeline.begin().expand(), pt.expand(), createChild, pipeline);

  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage(childTransform.toString());
  components.registerPTransform(transform, Collections.singletonList(childTransform));
}
 

@Test
public void parDo() {
  Create.Values<String> lines =
      Create.of(
          "apple orange grape banana apple banana",
          "banana orange banana papaya");

  PCollection<String> linesPColl = testPipeline.apply(lines);

  PCollection<String> results = Task.applyTransform(linesPColl);

  PAssert.that(results)
      .containsInAnyOrder(
          "apple:2",
          "banana:4",
          "grape:1",
          "orange:2",
          "papaya:1"
      );

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void producingOwnOutputWithCompositeFails() {
  final PCollection<Long> comp =
      PCollection.createPrimitiveOutputInternal(
          pipeline, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarLongCoder.of());
  PTransform<PBegin, PCollection<Long>> root =
      new PTransform<PBegin, PCollection<Long>>() {
        @Override
        public PCollection<Long> expand(PBegin input) {
          return comp;
        }
      };
  hierarchy.pushNode("Composite", PBegin.in(pipeline), root);

  Create.Values<Integer> create = Create.of(1);
  hierarchy.pushNode("Create", PBegin.in(pipeline), create);
  hierarchy.setOutput(pipeline.apply(create));
  hierarchy.popNode();

  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("contains a primitive POutput produced by it");
  thrown.expectMessage("primitive transforms are permitted to produce");
  thrown.expectMessage("Composite");
  hierarchy.setOutput(comp);
}
 

/**
 * Traverse the specified {@link PTransform}, collecting {@link DisplayData} registered on the
 * inner primitive {@link PTransform PTransforms}.
 *
 * @param root The root {@link PTransform} to traverse
 * @param inputCoder The coder to set for the {@link PTransform} input, or null to infer the
 *     default coder.
 * @return the set of {@link DisplayData} for primitive {@link PTransform PTransforms}.
 */
public <InputT> Set<DisplayData> displayDataForPrimitiveTransforms(
    final PTransform<? super PCollection<InputT>, ? extends POutput> root,
    Coder<InputT> inputCoder) {

  Create.Values<InputT> input;
  if (inputCoder != null) {
    input = Create.empty(inputCoder);
  } else {
    // These types don't actually work, but the pipeline will never be run
    input = (Create.Values<InputT>) Create.empty(VoidCoder.of());
  }

  Pipeline pipeline = Pipeline.create(options);
  pipeline.apply("Input", input).apply("Transform", root);

  return displayDataForPipeline(pipeline, root);
}
 

@SuppressWarnings("unchecked")
@Test
public void combine_combineFn() {
  Create.Values<KV<String, Integer>> values = Create.of(
      KV.of(Task.PLAYER_1, 15), KV.of(Task.PLAYER_2, 10), KV.of(Task.PLAYER_1, 100),
      KV.of(Task.PLAYER_3, 25), KV.of(Task.PLAYER_2, 75)
  );
  PCollection<KV<String, Integer>> numbers = testPipeline.apply(values);

  PCollection<KV<String, Integer>> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(
          KV.of(Task.PLAYER_1, 115), KV.of(Task.PLAYER_2, 85), KV.of(Task.PLAYER_3, 25)
      );

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void combine_simpleFn() {
  Create.Values<Integer> values = Create.of(10, 30, 50, 70, 90);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Integer> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(250);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void combine_binaryCombineFn_lambda() {
  Create.Values<BigInteger> values = Create.of(
      BigInteger.valueOf(10), BigInteger.valueOf(20), BigInteger.valueOf(30),
      BigInteger.valueOf(40), BigInteger.valueOf(50)
  );
  PCollection<BigInteger> numbers = testPipeline.apply(values);

  PCollection<BigInteger> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(BigInteger.valueOf(150));

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void combine_combineFn() {
  Create.Values<Integer> values = Create.of(10, 20, 50, 70, 90);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Double> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(48.0);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void flatMapElements() {
  Create.Values<String> values = Create.of("Apache Beam", "Unified Batch and Streaming");
  PCollection<String> numbers = testPipeline.apply(values);

  PCollection<String> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder("Apache", "Beam", "Unified", "Batch", "and", "Streaming");

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void parDo() {
  Create.Values<Integer> values = Create.of(1, 2, 3, 4, 5);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Integer> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(10, 20, 30, 40, 50);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void compositeTransform() {
  Create.Values<String> values = Create.of("1,2,3,4,5", "6,7,8,9,10");

  PCollection<Integer> results =
      testPipeline
          .apply(values)
          .apply(new ExtractAndMultiplyNumbers());

  PAssert.that(results)
      .containsInAnyOrder(10, 20, 30, 40, 50, 60, 70, 80, 90, 100);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void filter_parDo() {
  Create.Values<Integer> values = Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Integer> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(1, 3, 5, 7, 9);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void sum() {
  Create.Values<Integer> values = Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Integer> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(55);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void max() {
  Create.Values<Integer> values = Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Integer> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(10);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void count() {
  Create.Values<Integer> values = Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Long> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(10L);

  testPipeline.run().waitUntilFinish();
}
 

@Test
public void registerTransformNullComponents() throws IOException {
  Create.Values<Integer> create = Create.of(1, 2, 3);
  PCollection<Integer> pt = pipeline.apply(create);
  String userName = "my_transform/my_nesting";
  AppliedPTransform<?, ?, ?> transform =
      AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline);
  thrown.expect(NullPointerException.class);
  thrown.expectMessage("child nodes may not be null");
  components.registerPTransform(transform, null);
}
 

@Test
public void mean() {
  Create.Values<Integer> values = Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
  PCollection<Integer> numbers = testPipeline.apply(values);

  PCollection<Double> results = Task.applyTransform(numbers);

  PAssert.that(results)
      .containsInAnyOrder(5.5);

  testPipeline.run().waitUntilFinish();
}
 

/**
 * Runs a write test against a BigQuery table to check that SchemaUpdateOption sets are taking
 * effect.
 *
 * <p>Attempt write a row via BigQueryIO.writeTables with the given params, then run the given
 * query, and finaly check the results of the query.
 *
 * @param schemaUpdateOptions The SchemaUpdateOption set to use
 * @param tableName The table to write to
 * @param schema The schema to use for the table
 * @param rowToInsert The row to insert
 * @param testQuery A testing SQL query to run after writing the row
 * @param expectedResult The expected result of the query as a nested list of column values (one
 *     list per result row)
 */
private void runWriteTest(
    Set<SchemaUpdateOption> schemaUpdateOptions,
    String tableName,
    TableSchema schema,
    TableRow rowToInsert,
    String testQuery,
    List<List<String>> expectedResult)
    throws Exception {
  Options options = TestPipeline.testingPipelineOptions().as(Options.class);
  options.setTempLocation(options.getTempRoot() + "/bq_it_temp");

  Pipeline p = Pipeline.create(options);
  Create.Values<TableRow> input = Create.<TableRow>of(rowToInsert);

  Write<TableRow> writer =
      BigQueryIO.writeTableRows()
          .to(String.format("%s:%s.%s", options.getProject(), BIG_QUERY_DATASET_ID, tableName))
          .withSchema(schema)
          .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
          .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND)
          .withSchemaUpdateOptions(schemaUpdateOptions);

  p.apply(input).apply(writer);
  p.run().waitUntilFinish();

  QueryResponse response = BQ_CLIENT.queryWithRetries(testQuery, project);

  List<List<String>> result =
      response.getRows().stream()
          .map(
              row ->
                  row.getF().stream()
                      .map(cell -> cell.getV().toString())
                      .collect(Collectors.toList()))
          .collect(Collectors.toList());

  assertEquals(expectedResult, result);
}
 

@Test
public void replaceSucceeds() {
  PTransform<?, ?> enclosingPT =
      new PTransform<PInput, POutput>() {
        @Override
        public POutput expand(PInput input) {
          return PDone.in(input.getPipeline());
        }
      };

  TransformHierarchy.Node enclosing =
      hierarchy.pushNode("Enclosing", PBegin.in(pipeline), enclosingPT);

  Create.Values<Long> originalTransform = Create.of(1L);
  TransformHierarchy.Node original =
      hierarchy.pushNode("Create", PBegin.in(pipeline), originalTransform);
  assertThat(hierarchy.getCurrent(), equalTo(original));
  PCollection<Long> originalOutput = pipeline.apply(originalTransform);
  hierarchy.setOutput(originalOutput);
  hierarchy.popNode();
  assertThat(original.finishedSpecifying, is(true));
  hierarchy.setOutput(PDone.in(pipeline));
  hierarchy.popNode();

  assertThat(hierarchy.getCurrent(), not(equalTo(enclosing)));
  Read.Bounded<Long> replacementTransform = Read.from(CountingSource.upTo(1L));
  PCollection<Long> replacementOutput = pipeline.apply(replacementTransform);
  Node replacement = hierarchy.replaceNode(original, PBegin.in(pipeline), replacementTransform);
  assertThat(hierarchy.getCurrent(), equalTo(replacement));
  hierarchy.setOutput(replacementOutput);

  TaggedPValue taggedReplacement = TaggedPValue.ofExpandedValue(replacementOutput);
  Map<PValue, ReplacementOutput> replacementOutputs =
      Collections.singletonMap(
          replacementOutput,
          ReplacementOutput.of(TaggedPValue.ofExpandedValue(originalOutput), taggedReplacement));
  hierarchy.replaceOutputs(replacementOutputs);

  assertThat(replacement.getInputs(), equalTo(original.getInputs()));
  assertThat(replacement.getEnclosingNode(), equalTo(original.getEnclosingNode()));
  assertThat(replacement.getEnclosingNode(), equalTo(enclosing));
  assertThat(replacement.getTransform(), equalTo(replacementTransform));
  // THe tags of the replacement transform are matched to the appropriate PValues of the original
  assertThat(replacement.getOutputs().keySet(), Matchers.contains(taggedReplacement.getTag()));
  assertThat(replacement.getOutputs().values(), Matchers.contains(originalOutput));
  hierarchy.popNode();
}
 

private StreamingFnApiCreate(Create.Values<T> transform, PCollection<T> originalOutput) {
  this.transform = transform;
  this.originalOutput = originalOutput;
}