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

static Factory createStateInternalFactory(
    String id,
    Coder<?> keyCoder,
    TaskContext context,
    SamzaPipelineOptions pipelineOptions,
    DoFnSignature signature) {
  final int batchGetSize = pipelineOptions.getStoreBatchGetSize();
  final Map<String, KeyValueStore<ByteArray, byte[]>> stores = new HashMap<>();
  stores.put(BEAM_STORE, getBeamStore(context));

  final Coder stateKeyCoder;
  if (keyCoder != null) {
    signature
        .stateDeclarations()
        .keySet()
        .forEach(
            stateId ->
                stores.put(
                    stateId, (KeyValueStore<ByteArray, byte[]>) context.getStore(stateId)));
    stateKeyCoder = keyCoder;
  } else {
    stateKeyCoder = VoidCoder.of();
  }
  return new Factory<>(Objects.toString(id), stores, stateKeyCoder, batchGetSize);
}
 

@Test
public void testTrivialProcessFnPropagatesOutputWindowAndTimestamp() throws Exception {
  // Tests that ProcessFn correctly propagates the window and timestamp of the element
  // inside the KeyedWorkItem.
  // The underlying DoFn is actually monolithic, so this doesn't test splitting.
  DoFn<Integer, String> fn = new ToStringFn();

  Instant base = Instant.now();

  IntervalWindow w =
      new IntervalWindow(
          base.minus(Duration.standardMinutes(1)), base.plus(Duration.standardMinutes(1)));

  ProcessFnTester<Integer, String, SomeRestriction, Void, Void> tester =
      new ProcessFnTester<>(
          base,
          fn,
          BigEndianIntegerCoder.of(),
          SerializableCoder.of(SomeRestriction.class),
          VoidCoder.of(),
          MAX_OUTPUTS_PER_BUNDLE,
          MAX_BUNDLE_DURATION);
  tester.startElement(
      WindowedValue.of(
          KV.of(42, new SomeRestriction()),
          base,
          Collections.singletonList(w),
          PaneInfo.ON_TIME_AND_ONLY_FIRING));

  assertEquals(
      Arrays.asList(
          TimestampedValue.of("42a", base),
          TimestampedValue.of("42b", base),
          TimestampedValue.of("42c", base)),
      tester.peekOutputElementsInWindow(w));
}
 

@Test
public void testResumeSetsTimer() throws Exception {
  DoFn<Integer, String> fn = new SelfInitiatedResumeFn();
  Instant base = Instant.now();
  dateTimeProvider.setDateTimeFixed(base.getMillis());
  ProcessFnTester<Integer, String, SomeRestriction, Void, Void> tester =
      new ProcessFnTester<>(
          base,
          fn,
          BigEndianIntegerCoder.of(),
          SerializableCoder.of(SomeRestriction.class),
          VoidCoder.of(),
          MAX_OUTPUTS_PER_BUNDLE,
          MAX_BUNDLE_DURATION);

  tester.startElement(42, new SomeRestriction());
  assertThat(tester.takeOutputElements(), contains("42"));

  // Should resume after 5 seconds: advancing by 3 seconds should have no effect.
  assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
  assertTrue(tester.takeOutputElements().isEmpty());

  // 6 seconds should be enough  should invoke the fn again.
  assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
  assertThat(tester.takeOutputElements(), contains("42"));

  // Should again resume after 5 seconds: advancing by 3 seconds should again have no effect.
  assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
  assertTrue(tester.takeOutputElements().isEmpty());

  // 6 seconds should again be enough.
  assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
  assertThat(tester.takeOutputElements(), contains("42"));
}
 

@Test
public void testResumeCarriesOverState() throws Exception {
  DoFn<Integer, String> fn = new CounterFn(1);
  Instant base = Instant.now();
  dateTimeProvider.setDateTimeFixed(base.getMillis());
  ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
      new ProcessFnTester<>(
          base,
          fn,
          BigEndianIntegerCoder.of(),
          SerializableCoder.of(OffsetRange.class),
          VoidCoder.of(),
          MAX_OUTPUTS_PER_BUNDLE,
          MAX_BUNDLE_DURATION);

  tester.startElement(42, new OffsetRange(0, 3));
  assertThat(tester.takeOutputElements(), contains("42"));
  assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
  assertThat(tester.takeOutputElements(), contains("43"));
  assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
  assertThat(tester.takeOutputElements(), contains("44"));
  // Should not resume the null residual.
  assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
  // After outputting all 3 items, should not output anything more.
  assertEquals(0, tester.takeOutputElements().size());
  // Should also not ask to resume.
  assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
}
 

@Test
public void testCheckpointsAfterDuration() throws Exception {
  // Don't bound number of outputs.
  int max = Integer.MAX_VALUE;
  // But bound bundle duration - the bundle should terminate.
  Duration maxBundleDuration = Duration.standardSeconds(1);
  // Create an fn that attempts to 2x output more than checkpointing allows.
  DoFn<Integer, String> fn = new CounterFn(Integer.MAX_VALUE);
  Instant base = Instant.now();
  int baseIndex = 42;

  ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
      new ProcessFnTester<>(
          base,
          fn,
          BigEndianIntegerCoder.of(),
          SerializableCoder.of(OffsetRange.class),
          VoidCoder.of(),
          max,
          maxBundleDuration);

  List<String> elements;

  tester.startElement(baseIndex, new OffsetRange(0, Long.MAX_VALUE));
  // Bundle should terminate, and should do at least some processing.
  elements = tester.takeOutputElements();
  assertFalse(elements.isEmpty());
  // Bundle should have run for at least the requested duration.
  assertThat(
      Instant.now().getMillis() - base.getMillis(),
      greaterThanOrEqualTo(maxBundleDuration.getMillis()));
}
 

@Test
public void testInvokesLifecycleMethods() throws Exception {
  DoFn<Integer, String> fn = new LifecycleVerifyingFn();
  try (ProcessFnTester<Integer, String, SomeRestriction, Void, Void> tester =
      new ProcessFnTester<>(
          Instant.now(),
          fn,
          BigEndianIntegerCoder.of(),
          SerializableCoder.of(SomeRestriction.class),
          VoidCoder.of(),
          MAX_OUTPUTS_PER_BUNDLE,
          MAX_BUNDLE_DURATION)) {
    tester.startElement(42, new SomeRestriction());
  }
}
 

@Test
public void testNullKey() {
  Void key = null;
  VoidCoder coder = VoidCoder.of();

  ByteBuffer byteBuffer = FlinkKeyUtils.encodeKey(key, coder);
  assertThat(FlinkKeyUtils.decodeKey(byteBuffer, coder), is(nullValue()));
}
 

@Test
public void sdkErrorsSurfaceOnClose() throws Exception {
  TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
  DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
  ExecutableStageDoFnOperator<Integer, Integer> operator =
      getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
      new OneInputStreamOperatorTestHarness<>(operator);

  testHarness.open();

  @SuppressWarnings("unchecked")
  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  @SuppressWarnings("unchecked")
  FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
  when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

  Exception expected = new RuntimeException(new Exception());
  doThrow(expected).when(bundle).close();
  thrown.expectCause(is(expected));

  operator.processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(0)));
  testHarness.close();
}
 

@Test
public void expectedInputsAreSent() throws Exception {
  TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
  DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
  ExecutableStageDoFnOperator<Integer, Integer> operator =
      getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

  @SuppressWarnings("unchecked")
  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  @SuppressWarnings("unchecked")
  FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
  when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

  WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
  WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
      new OneInputStreamOperatorTestHarness<>(operator);

  testHarness.open();

  testHarness.processElement(new StreamRecord<>(one));
  testHarness.processElement(new StreamRecord<>(two));
  testHarness.processElement(new StreamRecord<>(three));

  verify(receiver).accept(one);
  verify(receiver).accept(two);
  verify(receiver).accept(three);
  verifyNoMoreInteractions(receiver);

  testHarness.close();
}
 

@Test
public void testStageBundleClosed() throws Exception {
  TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
  DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
  ExecutableStageDoFnOperator<Integer, Integer> operator =
      getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
      new OneInputStreamOperatorTestHarness<>(operator);

  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(bundle.getInputReceivers())
      .thenReturn(
          ImmutableMap.<String, FnDataReceiver<WindowedValue>>builder()
              .put("input", Mockito.mock(FnDataReceiver.class))
              .build());
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  testHarness.open();
  testHarness.close();

  verify(stageBundleFactory).close();
  verify(stageContext).close();
  verifyNoMoreInteractions(stageBundleFactory);

  // close() will also call dispose(), but call again to verify no new bundle
  // is created afterwards
  operator.dispose();
  verifyNoMoreInteractions(bundle);
}
 

@SuppressWarnings("unchecked")
private <T> Coder<T> getDefaultCoder(Class<T> c) {
    if (Writable.class.isAssignableFrom(c)) {
        Class<? extends Writable> writableClass = (Class<? extends Writable>) c;
        return (Coder<T>) WritableCoder.of(writableClass);
    } else if (Void.class.equals(c)) {
        return (Coder<T>) VoidCoder.of();
    }
    // TODO: how to use registered coders here?
    throw new IllegalStateException("Cannot find coder for " + c);
}
 

/**
 * Attempt to infer the type for some very common Apache Beam parameterized types.
 *
 * <p>TODO: Instead, build a TypeDescriptor so that the {@link CoderRegistry} is invoked for the
 * type instead of hard coding the coders for common types.
 */
private static Coder<?> inferCoderFromObject(
    CoderRegistry coderRegistry, SchemaRegistry schemaRegistry, Object o)
    throws CannotProvideCoderException {

  if (o == null) {
    return VoidCoder.of();
  }

  try {
    return SchemaCoder.of(
        schemaRegistry.getSchema(o.getClass()),
        TypeDescriptor.of(o.getClass()),
        (SerializableFunction) schemaRegistry.getToRowFunction(o.getClass()),
        (SerializableFunction) schemaRegistry.getFromRowFunction(o.getClass()));
  } catch (NoSuchSchemaException e) {
    // No schema.
  }

  if (o instanceof TimestampedValue) {
    return TimestampedValueCoder.of(
        inferCoderFromObject(coderRegistry, schemaRegistry, ((TimestampedValue) o).getValue()));
  } else if (o instanceof List) {
    return ListCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Set) {
    return SetCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Collection) {
    return CollectionCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Iterable) {
    return IterableCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Map) {
    return MapCoder.of(
        inferCoderFromObjects(coderRegistry, schemaRegistry, ((Map) o).keySet()),
        inferCoderFromObjects(coderRegistry, schemaRegistry, ((Map) o).entrySet()));
  } else if (o instanceof KV) {
    return KvCoder.of(
        inferCoderFromObject(coderRegistry, schemaRegistry, ((KV) o).getKey()),
        inferCoderFromObject(coderRegistry, schemaRegistry, ((KV) o).getValue()));
  } else {
    return coderRegistry.getCoder(o.getClass());
  }
}
 

@Override
public Coder<Void> getAccumulatorCoder(CoderRegistry registry, Coder<Integer> inputCoder) {
  return (Coder) VoidCoder.of();
}
 

@Test
public void testCheckpointsAfterNumOutputs() throws Exception {
  int max = 100;
  DoFn<Integer, String> fn = new CounterFn(Integer.MAX_VALUE);
  Instant base = Instant.now();
  int baseIndex = 42;

  ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
      new ProcessFnTester<>(
          base,
          fn,
          BigEndianIntegerCoder.of(),
          SerializableCoder.of(OffsetRange.class),
          VoidCoder.of(),
          max,
          MAX_BUNDLE_DURATION);

  List<String> elements;

  // Create an fn that attempts to 2x output more than checkpointing allows.
  tester.startElement(baseIndex, new OffsetRange(0, 2 * max + max / 2));
  elements = tester.takeOutputElements();
  assertEquals(max, elements.size());
  // Should output the range [0, max)
  assertThat(elements, hasItem(String.valueOf(baseIndex)));
  assertThat(elements, hasItem(String.valueOf(baseIndex + max - 1)));

  assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
  elements = tester.takeOutputElements();
  assertEquals(max, elements.size());
  // Should output the range [max, 2*max)
  assertThat(elements, hasItem(String.valueOf(baseIndex + max)));
  assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max - 1)));

  assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
  elements = tester.takeOutputElements();
  assertEquals(max / 2, elements.size());
  // Should output the range [2*max, 2*max + max/2)
  assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max)));
  assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max + max / 2 - 1)));
  assertThat(elements, not(hasItem(String.valueOf(baseIndex + 2 * max + max / 2))));
}