Java Code Examples for org.apache.beam.sdk.util.WindowedValue#withValue()

private void processElementForSplitRestriction(
    WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>> elem) {
  currentElement = elem.withValue(elem.getValue().getKey());
  currentRestriction = elem.getValue().getValue().getKey();
  currentWatermarkEstimatorState = elem.getValue().getValue().getValue();
  try {
    doFnInvoker.invokeSplitRestriction(processContext);
  } finally {
    currentElement = null;
    currentRestriction = null;
    currentWatermarkEstimatorState = null;
  }

  // TODO(BEAM-10212): Support caching state data across bundle boundaries.
  this.stateAccessor.finalizeState();
}
 

private void processElementForWindowObservingSplitRestriction(
    WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>> elem) {
  currentElement = elem.withValue(elem.getValue().getKey());
  currentRestriction = elem.getValue().getValue().getKey();
  currentWatermarkEstimatorState = elem.getValue().getValue().getValue();
  try {
    Iterator<BoundedWindow> windowIterator =
        (Iterator<BoundedWindow>) elem.getWindows().iterator();
    while (windowIterator.hasNext()) {
      currentWindow = windowIterator.next();
      doFnInvoker.invokeSplitRestriction(processContext);
    }
  } finally {
    currentElement = null;
    currentRestriction = null;
    currentWatermarkEstimatorState = null;
    currentWindow = null;
  }

  // TODO(BEAM-10212): Support caching state data across bundle boundaries.
  this.stateAccessor.finalizeState();
}
 

@Override
public WindowedValue<KV<K, Iterable<V>>> next() {
  while (inner.hasNext()) {
    final ByteArray nextKey = inner.peek()._1;
    if (nextKey.equals(currentKey)) {
      // we still did not see all values for a given key
      inner.next();
      continue;
    }
    currentKey = nextKey;
    final WindowedValue<KV<K, V>> decodedItem = decodeItem(inner.peek());
    return decodedItem.withValue(
        KV.of(decodedItem.getValue().getKey(), new ValueIterator(inner, currentKey)));
  }
  hasNext = false;
  return null;
}
 

@Override
public void flatMap(
    WindowedValue<KV<K, InputT>> inWithMultipleWindows,
    Collector<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> out)
    throws Exception {

  // we need to wrap each one work item per window for now
  // since otherwise the PushbackSideInputRunner will not correctly
  // determine whether side inputs are ready
  //
  // this is tracked as https://issues.apache.org/jira/browse/BEAM-1850
  for (WindowedValue<KV<K, InputT>> in : inWithMultipleWindows.explodeWindows()) {
    SingletonKeyedWorkItem<K, InputT> workItem =
        new SingletonKeyedWorkItem<>(
            in.getValue().getKey(), in.withValue(in.getValue().getValue()));

    out.collect(WindowedValue.valueInGlobalWindow(workItem));
  }
}
 

@Override
public void flatMap(
    WindowedValue<KV<K, InputT>> inWithMultipleWindows,
    Collector<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> out)
    throws Exception {

  // we need to wrap each one work item per window for now
  // since otherwise the PushbackSideInputRunner will not correctly
  // determine whether side inputs are ready
  //
  // this is tracked as https://issues.apache.org/jira/browse/BEAM-1850
  for (WindowedValue<KV<K, InputT>> in : inWithMultipleWindows.explodeWindows()) {
    SingletonKeyedWorkItem<K, InputT> workItem =
        new SingletonKeyedWorkItem<>(
            in.getValue().getKey(), in.withValue(in.getValue().getValue()));

    out.collect(in.withValue(workItem));
  }
}
 

@Override
public UncommittedBundle<T> add(WindowedValue<T> element) {
  try {
    // Use the cloned value to ensure that if the coder behaves poorly (e.g. a NoOpCoder that
    // does not expect to be used) that is reflected in the values given to downstream
    // transforms
    WindowedValue<T> clone = element.withValue(CoderUtils.clone(coder, element.getValue()));
    underlying.add(clone);
  } catch (CoderException e) {
    throw UserCodeException.wrap(e);
  }
  return this;
}
 

/**
 * Resumes from a timer and returns the current element/restriction pair (with an up-to-date value
 * of the restriction).
 */
public WindowedValue<KV<InputT, RestrictionT>> resume(TimerData timer) {
  initState(timer.getNamespace());
  WindowedValue<KV<InputT, RestrictionT>> seed = seedState.read();
  inputTimestamp = seed.getTimestamp();
  return seed.withValue(KV.of(seed.getValue().getKey(), restrictionState.read()));
}
 

@Override
public WindowedValue<ValueWithRecordId<T>> getCurrent() throws NoSuchElementException {
  WindowedValue<T> result =
      WindowedValue.timestampedValueInGlobalWindow(
          reader.getCurrent(), reader.getCurrentTimestamp());
  return result.withValue(
      new ValueWithRecordId<>(result.getValue(), reader.getCurrentRecordId()));
}
 

@Override
public boolean advance() throws IOException {
  if (!iterator.hasNext()) {
    current = null;
    return false;
  }
  ShuffleEntry record = iterator.next();
  K key = CoderUtils.decodeFromByteArray(shuffleReader.keyCoder, record.getKey());
  WindowedValue<V> windowedValue =
      CoderUtils.decodeFromByteArray(shuffleReader.windowedValueCoder, record.getValue());
  shuffleReader.notifyElementRead(record.length());
  current = windowedValue.withValue(KV.of(key, windowedValue.getValue()));
  return true;
}
 

@Override
public void processElement(
    WindowedValue<KV<K, V>> inputElement, OpEmitter<KeyedWorkItem<K, V>> emitter) {
  final KV<K, V> kv = inputElement.getValue();
  for (WindowedValue<KV<K, V>> windowedValue : inputElement.explodeWindows()) {
    final KeyedWorkItem<K, V> workItem =
        new SingletonKeyedWorkItem<>(kv.getKey(), windowedValue.withValue(kv.getValue()));
    emitter.emitElement(windowedValue.withValue(workItem));
  }
}
 

@Override
public WindowedValue<KV<Void, T>> map(WindowedValue<T> value) {
  return value.withValue(KV.of(null, value.getValue()));
}
 

@Override
public Object getKeyFromInputPair(Object pair) {
  @SuppressWarnings("unchecked")
  WindowedValue<KV<?, ?>> windowedKv = (WindowedValue<KV<?, ?>>) pair;
  return windowedKv.withValue(windowedKv.getValue().getKey());
}
 

@Override
public Object makeOutputPair(Object key, Object values) {
  WindowedValue<?> windowedKey = (WindowedValue<?>) key;
  return windowedKey.withValue(KV.of(windowedKey.getValue(), values));
}
 

@Override
public WindowedValue<KV<K, V>> call(KV<K, WindowedValue<V>> kv) {
  WindowedValue<V> wv = kv.getValue();
  return wv.withValue(KV.of(kv.getKey(), wv.getValue()));
}
 

public ThrowingFunction<WindowedValue<String>, WindowedValue<String>>
    createMapFunctionForPTransform(String ptransformId, PTransform pTransform) {
  assertEquals(EXPECTED_ID, ptransformId);
  assertEquals(EXPECTED_PTRANSFORM, pTransform);
  return (WindowedValue<String> str) -> str.withValue(str.getValue().toUpperCase());
}
 

@Override
public Object makeOutputPair(Object key, Object values) {
  WindowedValue<?> windowedKey = (WindowedValue<?>) key;
  return windowedKey.withValue(KV.of(windowedKey.getValue(), values));
}
 

@Override
public Object getKeyFromInputPair(Object pair) {
  @SuppressWarnings("unchecked")
  WindowedValue<KV<?, ?>> windowedKv = (WindowedValue<KV<?, ?>>) pair;
  return windowedKv.withValue(windowedKv.getValue().getKey());
}