Java Code Examples for org.apache.beam.sdk.coders.StringUtf8Coder#of()

@Test
public void testGet() throws Exception {
  FakeBeamFnStateClient fakeBeamFnStateClient =
      new FakeBeamFnStateClient(
          ImmutableMap.of(
              key("A"), encode("A1", "A2", "A3"),
              key("B"), encode("B1", "B2")));

  MultimapSideInput<String, String> multimapSideInput =
      new MultimapSideInput<>(
          fakeBeamFnStateClient,
          "instructionId",
          "ptransformId",
          "sideInputId",
          ByteString.copyFromUtf8("encodedWindow"),
          StringUtf8Coder.of(),
          StringUtf8Coder.of());
  assertArrayEquals(
      new String[] {"A1", "A2", "A3"},
      Iterables.toArray(multimapSideInput.get("A"), String.class));
  assertArrayEquals(
      new String[] {"B1", "B2"}, Iterables.toArray(multimapSideInput.get("B"), String.class));
  assertArrayEquals(
      new String[] {}, Iterables.toArray(multimapSideInput.get("unknown"), String.class));
}
 

@Test
public void testReiteration() throws Exception {
  FakeBeamFnStateClient fakeBeamFnStateClient =
      new FakeBeamFnStateClient(
          ImmutableMap.of(
              key("nonEmptySuffix"), encode("C", "D", "E", "F", "G", "H", "I", "J", "K"),
              key("emptySuffix"), encode()));

  StateBackedIterable<String> iterable =
      new StateBackedIterable<>(
          fakeBeamFnStateClient,
          "instruction",
          encode(suffixKey),
          StringUtf8Coder.of(),
          prefix);

  assertEquals(expected, Lists.newArrayList(iterable));
  assertEquals(expected, Lists.newArrayList(iterable));
  assertEquals(expected, Lists.newArrayList(iterable));
}
 

@Override
protected byte[][] getSortedTestData() {
  StringUtf8Coder coder = StringUtf8Coder.of();

  try {
    return new byte[][] {
      CoderUtils.encodeToByteArray(coder, ""),
      CoderUtils.encodeToByteArray(coder, "Lorem Ipsum Dolor Omit Longer"),
      CoderUtils.encodeToByteArray(coder, "aaaa"),
      CoderUtils.encodeToByteArray(coder, "abcd"),
      CoderUtils.encodeToByteArray(coder, "abce"),
      CoderUtils.encodeToByteArray(coder, "abdd"),
      CoderUtils.encodeToByteArray(coder, "accd"),
      CoderUtils.encodeToByteArray(coder, "bbcd")
    };
  } catch (CoderException e) {
    throw new RuntimeException("Could not encode values.", e);
  }
}
 

@Test
public void testCreateWithCombinerAndStreaming() throws Exception {
  StreamingOptions options = PipelineOptionsFactory.as(StreamingOptions.class);
  options.setStreaming(true);

  Coder keyCoder = StringUtf8Coder.of();
  Coder valueCoder = BigEndianIntegerCoder.of();
  KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder);

  TestOutputReceiver receiver =
      new TestOutputReceiver(
          new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)),
          counterSet,
          NameContextsForTests.nameContextForTest());

  ParDoFn pgbk =
      PartialGroupByKeyParDoFns.create(
          options,
          kvCoder,
          AppliedCombineFn.withInputCoder(
              Sum.ofIntegers(), CoderRegistry.createDefault(), kvCoder),
          NullSideInputReader.empty(),
          receiver,
          null);
  assertTrue(pgbk instanceof SimplePartialGroupByKeyParDoFn);
}
 

@Test
public void testGet() throws Exception {
  FakeBeamFnStateClient fakeClient =
      new FakeBeamFnStateClient(ImmutableMap.of(key("A"), encode("A1", "A2", "A3")));
  BagUserState<String> userState =
      new BagUserState<>(
          fakeClient,
          "instructionId",
          "ptransformId",
          "stateId",
          ByteString.copyFromUtf8("encodedWindow"),
          encode("A"),
          StringUtf8Coder.of());
  assertArrayEquals(
      new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class));

  userState.asyncClose();
  thrown.expect(IllegalStateException.class);
  userState.get();
}
 

@Test
public void testRequestDynamicSplit() throws Exception {
  // Note that exhaustive tests for AvroSource's split behavior exist in {@link AvroSourceTest}.
  List<List<String>> elements = generateInputBlocks(10, 100 * 100, 100);
  Coder<String> coder = StringUtf8Coder.of();
  AvroFileInfo<String> fileInfo = initInputFile(elements, coder);
  AvroByteReader<String> reader =
      new AvroByteReader<String>(
          fileInfo.filename, 0L, Long.MAX_VALUE, coder, PipelineOptionsFactory.create());
  // Read most of the records before the proposed split point.
  testRequestDynamicSplitInternal(reader, 0.5F, 490, SplitVerificationBehavior.VERIFY_SUCCESS);
  // Read a single record.
  testRequestDynamicSplitInternal(reader, 0.5F, 1, SplitVerificationBehavior.VERIFY_SUCCESS);
  // Read zero records.
  testRequestDynamicSplitInternal(reader, 0.5F, 0, SplitVerificationBehavior.VERIFY_FAILURE);
  // Read almost the entire input.
  testRequestDynamicSplitInternal(reader, 0.5F, 900, SplitVerificationBehavior.VERIFY_FAILURE);
  // Read the entire input.
  testRequestDynamicSplitInternal(reader, 0.5F, 2000, SplitVerificationBehavior.VERIFY_FAILURE);
}
 

@Test
public void testClear() throws Exception {
  FakeBeamFnStateClient fakeClient =
      new FakeBeamFnStateClient(ImmutableMap.of(key("A"), encode("A1", "A2", "A3")));
  BagUserState<String> userState =
      new BagUserState<>(
          fakeClient,
          "instructionId",
          "ptransformId",
          "stateId",
          ByteString.copyFromUtf8("encodedWindow"),
          encode("A"),
          StringUtf8Coder.of());
  assertArrayEquals(
      new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class));
  userState.clear();
  assertFalse(userState.get().iterator().hasNext());
  userState.append("A4");
  assertArrayEquals(new String[] {"A4"}, Iterables.toArray(userState.get(), String.class));
  userState.clear();
  assertFalse(userState.get().iterator().hasNext());
  userState.asyncClose();

  assertNull(fakeClient.getData().get(key("A")));
  thrown.expect(IllegalStateException.class);
  userState.clear();
}
 

@Test
public void testCheckpointCoderNulls() throws Exception {
  CheckpointCoder<String> coder = new CheckpointCoder<>(StringUtf8Coder.of());
  Checkpoint<String> emptyCheckpoint = new Checkpoint<>(null, null);
  Checkpoint<String> decodedEmptyCheckpoint =
      CoderUtils.decodeFromByteArray(coder, CoderUtils.encodeToByteArray(coder, emptyCheckpoint));
  assertNull(decodedEmptyCheckpoint.getResidualElements());
  assertNull(decodedEmptyCheckpoint.getResidualSource());
}
 

@Override
public void encode(InvoiceGroupingKey value, OutputStream outStream) throws IOException {
  Coder<String> stringCoder = StringUtf8Coder.of();
  stringCoder.encode(value.startDate(), outStream);
  stringCoder.encode(value.endDate(), outStream);
  stringCoder.encode(value.productAccountKey(), outStream);
  stringCoder.encode(value.usageGroupingKey(), outStream);
  stringCoder.encode(value.description(), outStream);
  stringCoder.encode(String.valueOf(value.unitPrice()), outStream);
  stringCoder.encode(value.unitPriceCurrency(), outStream);
  stringCoder.encode(value.poNumber(), outStream);
}
 

@Test
public void testCreateWithCombinerAndBatchSideInputs() throws Exception {
  PipelineOptions options = PipelineOptionsFactory.create();

  Coder keyCoder = StringUtf8Coder.of();
  Coder valueCoder = BigEndianIntegerCoder.of();
  KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder);

  TestOutputReceiver receiver =
      new TestOutputReceiver(
          new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)),
          counterSet,
          NameContextsForTests.nameContextForTest());

  StepContext stepContext =
      BatchModeExecutionContext.forTesting(options, "testStage")
          .getStepContext(TestOperationContext.create(counterSet));

  when(mockSideInputReader.isEmpty()).thenReturn(false);

  ParDoFn pgbk =
      PartialGroupByKeyParDoFns.create(
          options,
          kvCoder,
          AppliedCombineFn.withInputCoder(
              Sum.ofIntegers(),
              CoderRegistry.createDefault(),
              kvCoder,
              ImmutableList.<PCollectionView<?>>of(),
              WindowingStrategy.globalDefault()),
          mockSideInputReader,
          receiver,
          stepContext);
  assertTrue(pgbk instanceof BatchSideInputPGBKParDoFn);
}
 

private <ReceiverT> StreamingSideInputFetcher<String, IntervalWindow> createFetcher(
    List<PCollectionView<String>> views) throws Exception {
  @SuppressWarnings({"unchecked", "rawtypes"})
  Iterable<PCollectionView<?>> typedViews = (Iterable) views;

  return new StreamingSideInputFetcher<String, IntervalWindow>(
      typedViews, StringUtf8Coder.of(), WindowingStrategy.of(WINDOW_FN), stepContext);
}
 

@Test
public void testCreateWithCombinerAndStreamingSideInputs() throws Exception {
  StreamingOptions options = PipelineOptionsFactory.as(StreamingOptions.class);
  options.setStreaming(true);

  Coder keyCoder = StringUtf8Coder.of();
  Coder valueCoder = BigEndianIntegerCoder.of();
  KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder);

  TestOutputReceiver receiver =
      new TestOutputReceiver(
          new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)),
          counterSet,
          NameContextsForTests.nameContextForTest());

  when(mockSideInputReader.isEmpty()).thenReturn(false);
  when(mockStreamingStepContext.stateInternals()).thenReturn((StateInternals) mockStateInternals);
  when(mockStateInternals.state(Matchers.<StateNamespace>any(), Matchers.<StateTag>any()))
      .thenReturn(mockState);
  when(mockState.read()).thenReturn(Maps.newHashMap());

  ParDoFn pgbk =
      PartialGroupByKeyParDoFns.create(
          options,
          kvCoder,
          AppliedCombineFn.withInputCoder(
              Sum.ofIntegers(),
              CoderRegistry.createDefault(),
              kvCoder,
              ImmutableList.<PCollectionView<?>>of(),
              WindowingStrategy.globalDefault()),
          mockSideInputReader,
          receiver,
          mockStreamingStepContext);
  assertTrue(pgbk instanceof StreamingSideInputPGBKParDoFn);
}
 

@Test(expected = IOException.class)
public void sourceThrowingException() throws Exception {
  final TestUnboundedSource<String> source = new TestUnboundedSource<>(StringUtf8Coder.of());
  source.advanceWatermarkToInfinity = true;
  source.throwOnClose = true;
  processElement(source);
}
 

public TestOutputCounter(CounterSet counters, NameContext nameContext) {
  this(
      "output_name",
      new ElementByteSizeObservableCoder<>(StringUtf8Coder.of()),
      counters,
      nameContext);
}
 

@Override
public Coder<String> getOutputCoder() {
  return StringUtf8Coder.of();
}
 

@Test
public void testStateRestore() throws Exception {
  DoFn<KV<String, Long>, KV<String, Long>> filterElementsEqualToCountFn =
      new DoFn<KV<String, Long>, KV<String, Long>>() {

        @StateId("counter")
        private final StateSpec<ValueState<Long>> counterSpec =
            StateSpecs.value(VarLongCoder.of());

        @ProcessElement
        public void processElement(
            ProcessContext context, @StateId("counter") ValueState<Long> count) {
          long currentCount = Optional.ofNullable(count.read()).orElse(0L);
          currentCount = currentCount + 1;
          count.write(currentCount);

          KV<String, Long> currentElement = context.element();
          if (currentCount == currentElement.getValue()) {
            context.output(currentElement);
          }
        }
      };

  WindowingStrategy<Object, GlobalWindow> windowingStrategy = WindowingStrategy.globalDefault();

  TupleTag<KV<String, Long>> outputTag = new TupleTag<>("main-output");

  StringUtf8Coder keyCoder = StringUtf8Coder.of();
  KvToByteBufferKeySelector keySelector = new KvToByteBufferKeySelector<>(keyCoder);
  KvCoder<String, Long> coder = KvCoder.of(keyCoder, VarLongCoder.of());

  FullWindowedValueCoder<KV<String, Long>> kvCoder =
      WindowedValue.getFullCoder(coder, windowingStrategy.getWindowFn().windowCoder());

  CoderTypeInformation<ByteBuffer> keyCoderInfo =
      new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of());

  OneInputStreamOperatorTestHarness<
          WindowedValue<KV<String, Long>>, WindowedValue<KV<String, Long>>>
      testHarness =
          createTestHarness(
              windowingStrategy,
              filterElementsEqualToCountFn,
              kvCoder,
              kvCoder,
              keyCoder,
              outputTag,
              keyCoderInfo,
              keySelector);
  testHarness.open();

  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));

  OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
  testHarness.close();

  testHarness =
      createTestHarness(
          windowingStrategy,
          filterElementsEqualToCountFn,
          kvCoder,
          kvCoder,
          keyCoder,
          outputTag,
          keyCoderInfo,
          keySelector);
  testHarness.initializeState(snapshot);
  testHarness.open();

  // after restore: counter = 2
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 4L))));
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 5L))));
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      contains(
          WindowedValue.valueInGlobalWindow(KV.of("a", 4L)),
          WindowedValue.valueInGlobalWindow(KV.of("a", 5L))));

  testHarness.close();
}
 

@Test
public void shouldWriteAndReadSnapshotForConcreteClassCoder() throws Exception {
  Coder<String> concreteClassCoder = StringUtf8Coder.of();
  testWriteAndReadConfigSnapshot(concreteClassCoder);
}
 

@Override
public Coder<String> getOutputCoder() {
  return StringUtf8Coder.of();
}
 

@Override
public Coder<String> getOutputCoder() {
  return StringUtf8Coder.of();
}
 

@Test
public void testSessions() throws Exception {
  TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
  ListOutputManager outputManager = new ListOutputManager();
  DoFnRunner<KeyedWorkItem<String, String>, KV<String, Iterable<String>>> runner =
      makeRunner(
          outputTag,
          outputManager,
          WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10)))
              .withTimestampCombiner(TimestampCombiner.EARLIEST));

  when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(0));

  runner.startBundle();

  WorkItem.Builder workItem1 = WorkItem.newBuilder();
  workItem1.setKey(ByteString.copyFromUtf8(KEY));
  workItem1.setWorkToken(WORK_TOKEN);
  InputMessageBundle.Builder messageBundle = workItem1.addMessageBundlesBuilder();
  messageBundle.setSourceComputationId(SOURCE_COMPUTATION_ID);

  Coder<String> valueCoder = StringUtf8Coder.of();
  addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(0), valueCoder, "v1");
  addElement(messageBundle, Arrays.asList(window(5, 15)), new Instant(5), valueCoder, "v2");
  addElement(messageBundle, Arrays.asList(window(15, 25)), new Instant(15), valueCoder, "v3");

  addElement(messageBundle, Arrays.asList(window(3, 13)), new Instant(3), valueCoder, "v0");

  runner.processElement(createValue(workItem1, valueCoder));

  runner.finishBundle();
  runner.startBundle();

  WorkItem.Builder workItem2 = WorkItem.newBuilder();
  workItem2.setKey(ByteString.copyFromUtf8(KEY));
  workItem2.setWorkToken(WORK_TOKEN);
  // Note that the WATERMARK timer for Instant(9) will have been deleted by
  // ReduceFnRunner when window(0, 10) was merged away.
  addTimer(workItem2, window(0, 15), new Instant(14), Timer.Type.WATERMARK);
  addTimer(workItem2, window(15, 25), new Instant(24), Timer.Type.WATERMARK);
  when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(25));

  runner.processElement(createValue(workItem2, valueCoder));

  runner.finishBundle();

  List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);

  assertThat(result.size(), equalTo(2));

  assertThat(
      result,
      containsInAnyOrder(
          WindowMatchers.isSingleWindowedValue(
              isKv(equalTo(KEY), containsInAnyOrder("v0", "v1", "v2")),
              equalTo(new Instant(0)),
              equalTo(window(0, 15))),
          WindowMatchers.isSingleWindowedValue(
              isKv(equalTo(KEY), containsInAnyOrder("v3")),
              equalTo(new Instant(15)),
              equalTo(window(15, 25)))));
}