http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkParDoBoundWrapper.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkParDoBoundWrapper.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkParDoBoundWrapper.java deleted file mode 100644 index 6be94b2..0000000 --- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkParDoBoundWrapper.java +++ /dev/null @@ -1,104 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.beam.runners.flink.translation.wrappers.streaming; - -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.options.PipelineOptions; -import org.apache.beam.sdk.transforms.OldDoFn; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.transforms.windowing.PaneInfo; -import org.apache.beam.sdk.util.TimerInternals; -import org.apache.beam.sdk.util.WindowedValue; -import org.apache.beam.sdk.util.WindowingInternals; -import org.apache.beam.sdk.util.WindowingStrategy; -import org.apache.beam.sdk.util.state.StateInternals; -import org.apache.beam.sdk.values.PCollectionView; -import org.apache.beam.sdk.values.TupleTag; - -import org.apache.flink.util.Collector; -import org.joda.time.Instant; - -import java.io.IOException; -import java.util.Collection; - -/** - * A wrapper for the {@link org.apache.beam.sdk.transforms.ParDo.Bound} Beam transformation. - * */ -public class FlinkParDoBoundWrapper<IN, OUT> extends FlinkAbstractParDoWrapper<IN, OUT, OUT> { - - public FlinkParDoBoundWrapper(PipelineOptions options, WindowingStrategy<?, ?> windowingStrategy, OldDoFn<IN, OUT> doFn) { - super(options, windowingStrategy, doFn); - } - - @Override - public void outputWithTimestampHelper(WindowedValue<IN> inElement, OUT output, Instant timestamp, Collector<WindowedValue<OUT>> collector) { - checkTimestamp(inElement, timestamp); - collector.collect(makeWindowedValue( - output, - timestamp, - inElement.getWindows(), - inElement.getPane())); - } - - @Override - public <T> void sideOutputWithTimestampHelper(WindowedValue<IN> inElement, T output, Instant timestamp, Collector<WindowedValue<OUT>> outCollector, TupleTag<T> tag) { - // ignore the side output, this can happen when a user does not register - // side outputs but then outputs using a freshly created TupleTag. - throw new RuntimeException("sideOutput() not not available in ParDo.Bound()."); - } - - @Override - public WindowingInternals<IN, OUT> windowingInternalsHelper(final WindowedValue<IN> inElement, final Collector<WindowedValue<OUT>> collector) { - return new WindowingInternals<IN, OUT>() { - @Override - public StateInternals stateInternals() { - throw new NullPointerException("StateInternals are not available for ParDo.Bound()."); - } - - @Override - public void outputWindowedValue(OUT output, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo pane) { - collector.collect(makeWindowedValue(output, timestamp, windows, pane)); - } - - @Override - public TimerInternals timerInternals() { - throw new NullPointerException("TimeInternals are not available for ParDo.Bound()."); - } - - @Override - public Collection<? extends BoundedWindow> windows() { - return inElement.getWindows(); - } - - @Override - public PaneInfo pane() { - return inElement.getPane(); - } - - @Override - public <T> void writePCollectionViewData(TupleTag<?> tag, Iterable<WindowedValue<T>> data, Coder<T> elemCoder) throws IOException { - throw new RuntimeException("writePCollectionViewData() not supported in Streaming mode."); - } - - @Override - public <T> T sideInput(PCollectionView<T> view, BoundedWindow mainInputWindow) { - throw new RuntimeException("sideInput() not implemented."); - } - }; - } -}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkStateInternals.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkStateInternals.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkStateInternals.java new file mode 100644 index 0000000..a3cf2e2 --- /dev/null +++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkStateInternals.java @@ -0,0 +1,1038 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import org.apache.beam.runners.flink.translation.types.CoderTypeInformation; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.CombineWithContext; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.OutputTimeFn; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.CombineContextFactory; +import org.apache.beam.sdk.util.state.AccumulatorCombiningState; +import org.apache.beam.sdk.util.state.BagState; +import org.apache.beam.sdk.util.state.ReadableState; +import org.apache.beam.sdk.util.state.State; +import org.apache.beam.sdk.util.state.StateContext; +import org.apache.beam.sdk.util.state.StateContexts; +import org.apache.beam.sdk.util.state.StateInternals; +import org.apache.beam.sdk.util.state.StateNamespace; +import org.apache.beam.sdk.util.state.StateTag; +import org.apache.beam.sdk.util.state.ValueState; +import org.apache.beam.sdk.util.state.WatermarkHoldState; + +import com.google.common.collect.Iterables; +import com.google.common.collect.Lists; + +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.runtime.state.AbstractStateBackend; +import org.joda.time.Instant; + +import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; + +/** + * {@link StateInternals} that uses a Flink {@link AbstractStateBackend} to + * manage state. + * + * <p>Note: In the Flink streaming runner the key is always encoded + * using an {@link Coder} and stored in a {@link ByteBuffer}. + */ +public class FlinkStateInternals<K> implements StateInternals<K> { + + private final Coder<K> keyCoder; + + private final AbstractStateBackend flinkStateBackend; + + // on recovery, these will no be properly set because we don't + // know which watermark hold states there are in the Flink State Backend + private final Map<String, Instant> watermarkHolds = new HashMap<>(); + + public FlinkStateInternals(AbstractStateBackend flinkStateBackend, Coder<K> keyCoder) { + this.flinkStateBackend = flinkStateBackend; + this.keyCoder = keyCoder; + } + + /** + * Returns the minimum over all watermark holds. + */ + public Instant watermarkHold() { + long min = Long.MAX_VALUE; + for (Instant hold: watermarkHolds.values()) { + min = Math.min(min, hold.getMillis()); + } + return new Instant(min); + } + + @Override + public K getKey() { + ByteBuffer keyBytes = (ByteBuffer) flinkStateBackend.getCurrentKey(); + try { + return CoderUtils.decodeFromByteArray(keyCoder, keyBytes.array()); + } catch (CoderException e) { + throw new RuntimeException("Error decoding key.", e); + } + } + + @Override + public <T extends State> T state( + final StateNamespace namespace, + StateTag<? super K, T> address) { + + return state(namespace, address, StateContexts.nullContext()); + } + + @Override + public <T extends State> T state( + final StateNamespace namespace, + StateTag<? super K, T> address, + final StateContext<?> context) { + + return address.bind(new StateTag.StateBinder<K>() { + + @Override + public <T> ValueState<T> bindValue( + StateTag<? super K, ValueState<T>> address, + Coder<T> coder) { + + return new FlinkValueState<>(flinkStateBackend, address, namespace, coder); + } + + @Override + public <T> BagState<T> bindBag( + StateTag<? super K, BagState<T>> address, + Coder<T> elemCoder) { + + return new FlinkBagState<>(flinkStateBackend, address, namespace, elemCoder); + } + + @Override + public <InputT, AccumT, OutputT> + AccumulatorCombiningState<InputT, AccumT, OutputT> + bindCombiningValue( + StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address, + Coder<AccumT> accumCoder, + Combine.CombineFn<InputT, AccumT, OutputT> combineFn) { + + return new FlinkAccumulatorCombiningState<>( + flinkStateBackend, address, combineFn, namespace, accumCoder); + } + + @Override + public <InputT, AccumT, OutputT> + AccumulatorCombiningState<InputT, AccumT, OutputT> bindKeyedCombiningValue( + StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address, + Coder<AccumT> accumCoder, + final Combine.KeyedCombineFn<? super K, InputT, AccumT, OutputT> combineFn) { + return new FlinkKeyedAccumulatorCombiningState<>( + flinkStateBackend, + address, + combineFn, + namespace, + accumCoder, + FlinkStateInternals.this); + } + + @Override + public <InputT, AccumT, OutputT> + AccumulatorCombiningState<InputT, AccumT, OutputT> bindKeyedCombiningValueWithContext( + StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address, + Coder<AccumT> accumCoder, + CombineWithContext.KeyedCombineFnWithContext< + ? super K, InputT, AccumT, OutputT> combineFn) { + return new FlinkAccumulatorCombiningStateWithContext<>( + flinkStateBackend, + address, + combineFn, + namespace, + accumCoder, + FlinkStateInternals.this, + CombineContextFactory.createFromStateContext(context)); + } + + @Override + public <W extends BoundedWindow> WatermarkHoldState<W> bindWatermark( + StateTag<? super K, WatermarkHoldState<W>> address, + OutputTimeFn<? super W> outputTimeFn) { + + return new FlinkWatermarkHoldState<>( + flinkStateBackend, FlinkStateInternals.this, address, namespace, outputTimeFn); + } + }); + } + + private static class FlinkValueState<K, T> implements ValueState<T> { + + private final StateNamespace namespace; + private final StateTag<? super K, ValueState<T>> address; + private final ValueStateDescriptor<T> flinkStateDescriptor; + private final AbstractStateBackend flinkStateBackend; + + FlinkValueState( + AbstractStateBackend flinkStateBackend, + StateTag<? super K, ValueState<T>> address, + StateNamespace namespace, + Coder<T> coder) { + + this.namespace = namespace; + this.address = address; + this.flinkStateBackend = flinkStateBackend; + + CoderTypeInformation<T> typeInfo = new CoderTypeInformation<>(coder); + + flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null); + } + + @Override + public void write(T input) { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).update(input); + } catch (Exception e) { + throw new RuntimeException("Error updating state.", e); + } + } + + @Override + public ValueState<T> readLater() { + return this; + } + + @Override + public T read() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public void clear() { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).clear(); + } catch (Exception e) { + throw new RuntimeException("Error clearing state.", e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + FlinkValueState<?, ?> that = (FlinkValueState<?, ?>) o; + + return namespace.equals(that.namespace) && address.equals(that.address); + + } + + @Override + public int hashCode() { + int result = namespace.hashCode(); + result = 31 * result + address.hashCode(); + return result; + } + } + + private static class FlinkBagState<K, T> implements BagState<T> { + + private final StateNamespace namespace; + private final StateTag<? super K, BagState<T>> address; + private final ListStateDescriptor<T> flinkStateDescriptor; + private final AbstractStateBackend flinkStateBackend; + + FlinkBagState( + AbstractStateBackend flinkStateBackend, + StateTag<? super K, BagState<T>> address, + StateNamespace namespace, + Coder<T> coder) { + + this.namespace = namespace; + this.address = address; + this.flinkStateBackend = flinkStateBackend; + + CoderTypeInformation<T> typeInfo = new CoderTypeInformation<>(coder); + + flinkStateDescriptor = new ListStateDescriptor<>(address.getId(), typeInfo); + } + + @Override + public void add(T input) { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).add(input); + } catch (Exception e) { + throw new RuntimeException("Error adding to bag state.", e); + } + } + + @Override + public BagState<T> readLater() { + return this; + } + + @Override + public Iterable<T> read() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).get(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public ReadableState<Boolean> isEmpty() { + return new ReadableState<Boolean>() { + @Override + public Boolean read() { + try { + Iterable<T> result = flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).get(); + return Iterables.isEmpty(result); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + + } + + @Override + public ReadableState<Boolean> readLater() { + return this; + } + }; + } + + @Override + public void clear() { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).clear(); + } catch (Exception e) { + throw new RuntimeException("Error clearing state.", e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + FlinkBagState<?, ?> that = (FlinkBagState<?, ?>) o; + + return namespace.equals(that.namespace) && address.equals(that.address); + + } + + @Override + public int hashCode() { + int result = namespace.hashCode(); + result = 31 * result + address.hashCode(); + return result; + } + } + + private static class FlinkAccumulatorCombiningState<K, InputT, AccumT, OutputT> + implements AccumulatorCombiningState<InputT, AccumT, OutputT> { + + private final StateNamespace namespace; + private final StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address; + private final Combine.CombineFn<InputT, AccumT, OutputT> combineFn; + private final ValueStateDescriptor<AccumT> flinkStateDescriptor; + private final AbstractStateBackend flinkStateBackend; + + FlinkAccumulatorCombiningState( + AbstractStateBackend flinkStateBackend, + StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address, + Combine.CombineFn<InputT, AccumT, OutputT> combineFn, + StateNamespace namespace, + Coder<AccumT> accumCoder) { + + this.namespace = namespace; + this.address = address; + this.combineFn = combineFn; + this.flinkStateBackend = flinkStateBackend; + + CoderTypeInformation<AccumT> typeInfo = new CoderTypeInformation<>(accumCoder); + + flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null); + } + + @Override + public AccumulatorCombiningState<InputT, AccumT, OutputT> readLater() { + return this; + } + + @Override + public void add(InputT value) { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT current = state.value(); + if (current == null) { + current = combineFn.createAccumulator(); + } + current = combineFn.addInput(current, value); + state.update(current); + } catch (Exception e) { + throw new RuntimeException("Error adding to state." , e); + } + } + + @Override + public void addAccum(AccumT accum) { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT current = state.value(); + if (current == null) { + state.update(accum); + } else { + current = combineFn.mergeAccumulators(Lists.newArrayList(current, accum)); + state.update(current); + } + } catch (Exception e) { + throw new RuntimeException("Error adding to state.", e); + } + } + + @Override + public AccumT getAccum() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public AccumT mergeAccumulators(Iterable<AccumT> accumulators) { + return combineFn.mergeAccumulators(accumulators); + } + + @Override + public OutputT read() { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT accum = state.value(); + if (accum != null) { + return combineFn.extractOutput(accum); + } else { + return combineFn.extractOutput(combineFn.createAccumulator()); + } + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public ReadableState<Boolean> isEmpty() { + return new ReadableState<Boolean>() { + @Override + public Boolean read() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value() == null; + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + + } + + @Override + public ReadableState<Boolean> readLater() { + return this; + } + }; + } + + @Override + public void clear() { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).clear(); + } catch (Exception e) { + throw new RuntimeException("Error clearing state.", e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + FlinkAccumulatorCombiningState<?, ?, ?, ?> that = + (FlinkAccumulatorCombiningState<?, ?, ?, ?>) o; + + return namespace.equals(that.namespace) && address.equals(that.address); + + } + + @Override + public int hashCode() { + int result = namespace.hashCode(); + result = 31 * result + address.hashCode(); + return result; + } + } + + private static class FlinkKeyedAccumulatorCombiningState<K, InputT, AccumT, OutputT> + implements AccumulatorCombiningState<InputT, AccumT, OutputT> { + + private final StateNamespace namespace; + private final StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address; + private final Combine.KeyedCombineFn<? super K, InputT, AccumT, OutputT> combineFn; + private final ValueStateDescriptor<AccumT> flinkStateDescriptor; + private final AbstractStateBackend flinkStateBackend; + private final FlinkStateInternals<K> flinkStateInternals; + + FlinkKeyedAccumulatorCombiningState( + AbstractStateBackend flinkStateBackend, + StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address, + Combine.KeyedCombineFn<? super K, InputT, AccumT, OutputT> combineFn, + StateNamespace namespace, + Coder<AccumT> accumCoder, + FlinkStateInternals<K> flinkStateInternals) { + + this.namespace = namespace; + this.address = address; + this.combineFn = combineFn; + this.flinkStateBackend = flinkStateBackend; + this.flinkStateInternals = flinkStateInternals; + + CoderTypeInformation<AccumT> typeInfo = new CoderTypeInformation<>(accumCoder); + + flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null); + } + + @Override + public AccumulatorCombiningState<InputT, AccumT, OutputT> readLater() { + return this; + } + + @Override + public void add(InputT value) { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT current = state.value(); + if (current == null) { + current = combineFn.createAccumulator(flinkStateInternals.getKey()); + } + current = combineFn.addInput(flinkStateInternals.getKey(), current, value); + state.update(current); + } catch (Exception e) { + throw new RuntimeException("Error adding to state." , e); + } + } + + @Override + public void addAccum(AccumT accum) { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT current = state.value(); + if (current == null) { + state.update(accum); + } else { + current = combineFn.mergeAccumulators( + flinkStateInternals.getKey(), + Lists.newArrayList(current, accum)); + state.update(current); + } + } catch (Exception e) { + throw new RuntimeException("Error adding to state.", e); + } + } + + @Override + public AccumT getAccum() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public AccumT mergeAccumulators(Iterable<AccumT> accumulators) { + return combineFn.mergeAccumulators(flinkStateInternals.getKey(), accumulators); + } + + @Override + public OutputT read() { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT accum = state.value(); + if (accum != null) { + return combineFn.extractOutput(flinkStateInternals.getKey(), accum); + } else { + return combineFn.extractOutput( + flinkStateInternals.getKey(), + combineFn.createAccumulator(flinkStateInternals.getKey())); + } + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public ReadableState<Boolean> isEmpty() { + return new ReadableState<Boolean>() { + @Override + public Boolean read() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value() == null; + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + + } + + @Override + public ReadableState<Boolean> readLater() { + return this; + } + }; + } + + @Override + public void clear() { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).clear(); + } catch (Exception e) { + throw new RuntimeException("Error clearing state.", e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + FlinkKeyedAccumulatorCombiningState<?, ?, ?, ?> that = + (FlinkKeyedAccumulatorCombiningState<?, ?, ?, ?>) o; + + return namespace.equals(that.namespace) && address.equals(that.address); + + } + + @Override + public int hashCode() { + int result = namespace.hashCode(); + result = 31 * result + address.hashCode(); + return result; + } + } + + private static class FlinkAccumulatorCombiningStateWithContext<K, InputT, AccumT, OutputT> + implements AccumulatorCombiningState<InputT, AccumT, OutputT> { + + private final StateNamespace namespace; + private final StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address; + private final CombineWithContext.KeyedCombineFnWithContext< + ? super K, InputT, AccumT, OutputT> combineFn; + private final ValueStateDescriptor<AccumT> flinkStateDescriptor; + private final AbstractStateBackend flinkStateBackend; + private final FlinkStateInternals<K> flinkStateInternals; + private final CombineWithContext.Context context; + + FlinkAccumulatorCombiningStateWithContext( + AbstractStateBackend flinkStateBackend, + StateTag<? super K, AccumulatorCombiningState<InputT, AccumT, OutputT>> address, + CombineWithContext.KeyedCombineFnWithContext< + ? super K, InputT, AccumT, OutputT> combineFn, + StateNamespace namespace, + Coder<AccumT> accumCoder, + FlinkStateInternals<K> flinkStateInternals, + CombineWithContext.Context context) { + + this.namespace = namespace; + this.address = address; + this.combineFn = combineFn; + this.flinkStateBackend = flinkStateBackend; + this.flinkStateInternals = flinkStateInternals; + this.context = context; + + CoderTypeInformation<AccumT> typeInfo = new CoderTypeInformation<>(accumCoder); + + flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null); + } + + @Override + public AccumulatorCombiningState<InputT, AccumT, OutputT> readLater() { + return this; + } + + @Override + public void add(InputT value) { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT current = state.value(); + if (current == null) { + current = combineFn.createAccumulator(flinkStateInternals.getKey(), context); + } + current = combineFn.addInput(flinkStateInternals.getKey(), current, value, context); + state.update(current); + } catch (Exception e) { + throw new RuntimeException("Error adding to state." , e); + } + } + + @Override + public void addAccum(AccumT accum) { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT current = state.value(); + if (current == null) { + state.update(accum); + } else { + current = combineFn.mergeAccumulators( + flinkStateInternals.getKey(), + Lists.newArrayList(current, accum), + context); + state.update(current); + } + } catch (Exception e) { + throw new RuntimeException("Error adding to state.", e); + } + } + + @Override + public AccumT getAccum() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public AccumT mergeAccumulators(Iterable<AccumT> accumulators) { + return combineFn.mergeAccumulators(flinkStateInternals.getKey(), accumulators, context); + } + + @Override + public OutputT read() { + try { + org.apache.flink.api.common.state.ValueState<AccumT> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + AccumT accum = state.value(); + return combineFn.extractOutput(flinkStateInternals.getKey(), accum, context); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public ReadableState<Boolean> isEmpty() { + return new ReadableState<Boolean>() { + @Override + public Boolean read() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value() == null; + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + + } + + @Override + public ReadableState<Boolean> readLater() { + return this; + } + }; + } + + @Override + public void clear() { + try { + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).clear(); + } catch (Exception e) { + throw new RuntimeException("Error clearing state.", e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + FlinkAccumulatorCombiningStateWithContext<?, ?, ?, ?> that = + (FlinkAccumulatorCombiningStateWithContext<?, ?, ?, ?>) o; + + return namespace.equals(that.namespace) && address.equals(that.address); + + } + + @Override + public int hashCode() { + int result = namespace.hashCode(); + result = 31 * result + address.hashCode(); + return result; + } + } + + private static class FlinkWatermarkHoldState<K, W extends BoundedWindow> + implements WatermarkHoldState<W> { + private final StateTag<? super K, WatermarkHoldState<W>> address; + private final OutputTimeFn<? super W> outputTimeFn; + private final StateNamespace namespace; + private final AbstractStateBackend flinkStateBackend; + private final FlinkStateInternals<K> flinkStateInternals; + private final ValueStateDescriptor<Instant> flinkStateDescriptor; + + public FlinkWatermarkHoldState( + AbstractStateBackend flinkStateBackend, + FlinkStateInternals<K> flinkStateInternals, + StateTag<? super K, WatermarkHoldState<W>> address, + StateNamespace namespace, + OutputTimeFn<? super W> outputTimeFn) { + this.address = address; + this.outputTimeFn = outputTimeFn; + this.namespace = namespace; + this.flinkStateBackend = flinkStateBackend; + this.flinkStateInternals = flinkStateInternals; + + CoderTypeInformation<Instant> typeInfo = new CoderTypeInformation<>(InstantCoder.of()); + flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null); + } + + @Override + public OutputTimeFn<? super W> getOutputTimeFn() { + return outputTimeFn; + } + + @Override + public WatermarkHoldState<W> readLater() { + return this; + } + + @Override + public ReadableState<Boolean> isEmpty() { + return new ReadableState<Boolean>() { + @Override + public Boolean read() { + try { + return flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor).value() == null; + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public ReadableState<Boolean> readLater() { + return this; + } + }; + + } + + @Override + public void add(Instant value) { + try { + org.apache.flink.api.common.state.ValueState<Instant> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + + Instant current = state.value(); + if (current == null) { + state.update(value); + flinkStateInternals.watermarkHolds.put(namespace.stringKey(), value); + } else { + Instant combined = outputTimeFn.combine(current, value); + state.update(combined); + flinkStateInternals.watermarkHolds.put(namespace.stringKey(), combined); + } + } catch (Exception e) { + throw new RuntimeException("Error updating state.", e); + } + } + + @Override + public Instant read() { + try { + org.apache.flink.api.common.state.ValueState<Instant> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + return state.value(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public void clear() { + flinkStateInternals.watermarkHolds.remove(namespace.stringKey()); + try { + org.apache.flink.api.common.state.ValueState<Instant> state = + flinkStateBackend.getPartitionedState( + namespace.stringKey(), + StringSerializer.INSTANCE, + flinkStateDescriptor); + state.clear(); + } catch (Exception e) { + throw new RuntimeException("Error reading state.", e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + FlinkWatermarkHoldState<?, ?> that = (FlinkWatermarkHoldState<?, ?>) o; + + if (!address.equals(that.address)) { + return false; + } + if (!outputTimeFn.equals(that.outputTimeFn)) { + return false; + } + return namespace.equals(that.namespace); + + } + + @Override + public int hashCode() { + int result = address.hashCode(); + result = 31 * result + outputTimeFn.hashCode(); + result = 31 * result + namespace.hashCode(); + return result; + } + } +} http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItem.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItem.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItem.java new file mode 100644 index 0000000..94bf3af --- /dev/null +++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItem.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import org.apache.beam.sdk.util.KeyedWorkItem; +import org.apache.beam.sdk.util.TimerInternals; +import org.apache.beam.sdk.util.WindowedValue; + +import java.util.Collections; + +public class SingletonKeyedWorkItem<K, ElemT> implements KeyedWorkItem<K, ElemT> { + + final K key; + final WindowedValue<ElemT> value; + + public SingletonKeyedWorkItem(K key, WindowedValue<ElemT> value) { + this.key = key; + this.value = value; + } + + @Override + public K key() { + return key; + } + + public WindowedValue<ElemT> value() { + return value; + } + + @Override + public Iterable<TimerInternals.TimerData> timersIterable() { + return Collections.EMPTY_LIST; + } + + @Override + public Iterable<WindowedValue<ElemT>> elementsIterable() { + return Collections.singletonList(value); + } +} http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItemCoder.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItemCoder.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItemCoder.java new file mode 100644 index 0000000..323f572 --- /dev/null +++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SingletonKeyedWorkItemCoder.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import static com.google.common.base.Preconditions.checkArgument; + +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.StandardCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.KeyedWorkItem; +import org.apache.beam.sdk.util.KeyedWorkItemCoder; +import org.apache.beam.sdk.util.PropertyNames; +import org.apache.beam.sdk.util.WindowedValue; + +import com.google.common.collect.ImmutableList; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.List; + +public class SingletonKeyedWorkItemCoder<K, ElemT> extends StandardCoder<SingletonKeyedWorkItem<K, ElemT>> { + /** + * Create a new {@link KeyedWorkItemCoder} with the provided key coder, element coder, and window + * coder. + */ + public static <K, ElemT> SingletonKeyedWorkItemCoder<K, ElemT> of( + Coder<K> keyCoder, Coder<ElemT> elemCoder, Coder<? extends BoundedWindow> windowCoder) { + return new SingletonKeyedWorkItemCoder<>(keyCoder, elemCoder, windowCoder); + } + + @JsonCreator + public static <K, ElemT> SingletonKeyedWorkItemCoder<K, ElemT> of( + @JsonProperty(PropertyNames.COMPONENT_ENCODINGS) List<Coder<?>> components) { + checkArgument(components.size() == 3, "Expecting 3 components, got %s", components.size()); + @SuppressWarnings("unchecked") + Coder<K> keyCoder = (Coder<K>) components.get(0); + @SuppressWarnings("unchecked") + Coder<ElemT> elemCoder = (Coder<ElemT>) components.get(1); + @SuppressWarnings("unchecked") + Coder<? extends BoundedWindow> windowCoder = (Coder<? extends BoundedWindow>) components.get(2); + return new SingletonKeyedWorkItemCoder<>(keyCoder, elemCoder, windowCoder); + } + + private final Coder<K> keyCoder; + private final Coder<ElemT> elemCoder; + private final Coder<? extends BoundedWindow> windowCoder; + private final WindowedValue.FullWindowedValueCoder<ElemT> valueCoder; + + private SingletonKeyedWorkItemCoder( + Coder<K> keyCoder, Coder<ElemT> elemCoder, Coder<? extends BoundedWindow> windowCoder) { + this.keyCoder = keyCoder; + this.elemCoder = elemCoder; + this.windowCoder = windowCoder; + valueCoder= WindowedValue.FullWindowedValueCoder.of(elemCoder, windowCoder); + } + + public Coder<K> getKeyCoder() { + return keyCoder; + } + + public Coder<ElemT> getElementCoder() { + return elemCoder; + } + + @Override + public void encode(SingletonKeyedWorkItem<K, ElemT> value, OutputStream outStream, Context context) + throws CoderException, IOException { + Context nestedContext = context.nested(); + keyCoder.encode(value.key(), outStream, nestedContext); + valueCoder.encode(value.value, outStream, nestedContext); + } + + @Override + public SingletonKeyedWorkItem<K, ElemT> decode(InputStream inStream, Context context) + throws CoderException, IOException { + Context nestedContext = context.nested(); + K key = keyCoder.decode(inStream, nestedContext); + WindowedValue<ElemT> value = valueCoder.decode(inStream, nestedContext); + return new SingletonKeyedWorkItem<>(key, value); + } + + @Override + public List<? extends Coder<?>> getCoderArguments() { + return ImmutableList.of(keyCoder, elemCoder, windowCoder); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + keyCoder.verifyDeterministic(); + elemCoder.verifyDeterministic(); + windowCoder.verifyDeterministic(); + } + + /** + * {@inheritDoc}. + * + * {@link KeyedWorkItemCoder} is not consistent with equals as it can return a + * {@link KeyedWorkItem} of a type different from the originally encoded type. + */ + @Override + public boolean consistentWithEquals() { + return false; + } + +} http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java new file mode 100644 index 0000000..0a279cc --- /dev/null +++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java @@ -0,0 +1,326 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import org.apache.beam.runners.core.GroupAlsoByWindowViaWindowSetDoFn; +import org.apache.beam.runners.flink.translation.wrappers.DataInputViewWrapper; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.OldDoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.ExecutionContext; +import org.apache.beam.sdk.util.KeyedWorkItem; +import org.apache.beam.sdk.util.KeyedWorkItems; +import org.apache.beam.sdk.util.SystemReduceFn; +import org.apache.beam.sdk.util.TimeDomain; +import org.apache.beam.sdk.util.TimerInternals; +import org.apache.beam.sdk.util.WindowedValue; +import org.apache.beam.sdk.util.WindowingStrategy; +import org.apache.beam.sdk.util.state.StateInternals; +import org.apache.beam.sdk.util.state.StateInternalsFactory; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.TupleTag; + +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.runtime.state.AbstractStateBackend; +import org.apache.flink.runtime.state.StateHandle; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.tasks.StreamTaskState; +import org.joda.time.Instant; + +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Queue; +import java.util.Set; +import javax.annotation.Nullable; + +/** + * Flink operator for executing window {@link DoFn DoFns}. + * + * @param <InputT> + * @param <OutputT> + */ +public class WindowDoFnOperator<K, InputT, OutputT> + extends DoFnOperator<KeyedWorkItem<K, InputT>, KV<K, OutputT>, WindowedValue<KV<K, OutputT>>> { + + /** + * To keep track of the current watermark so that we can immediately fire if a trigger + * registers an event time callback for a timestamp that lies in the past. + */ + private transient long currentWatermark = Long.MIN_VALUE; + + private final Coder<K> keyCoder; + private final TimerInternals.TimerDataCoder timerCoder; + + private transient Set<Tuple2<ByteBuffer, TimerInternals.TimerData>> watermarkTimers; + private transient Queue<Tuple2<ByteBuffer, TimerInternals.TimerData>> watermarkTimersQueue; + + private FlinkStateInternals<K> stateInternals; + + private final SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> systemReduceFn; + + public WindowDoFnOperator( + SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> systemReduceFn, + TupleTag<KV<K, OutputT>> mainOutputTag, + List<TupleTag<?>> sideOutputTags, + OutputManagerFactory<WindowedValue<KV<K, OutputT>>> outputManagerFactory, + WindowingStrategy<?, ?> windowingStrategy, + Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputs, + PipelineOptions options, + Coder<K> keyCoder) { + super( + null, + mainOutputTag, + sideOutputTags, + outputManagerFactory, + windowingStrategy, + sideInputs, + options); + + this.systemReduceFn = systemReduceFn; + + this.keyCoder = keyCoder; + this.timerCoder = + TimerInternals.TimerDataCoder.of(windowingStrategy.getWindowFn().windowCoder()); + } + + @Override + protected OldDoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> getDoFn() { + StateInternalsFactory<K> stateInternalsFactory = new StateInternalsFactory<K>() { + @Override + public StateInternals<K> stateInternalsForKey(K key) { + //this will implicitly be keyed by the key of the incoming + // element or by the key of a firing timer + return stateInternals; + } + }; + + // we have to do the unchecked cast because GroupAlsoByWindowViaWindowSetDoFn.create + // has the window type as generic parameter while WindowingStrategy is almost always + // untyped. + @SuppressWarnings("unchecked") + OldDoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> doFn = + GroupAlsoByWindowViaWindowSetDoFn.create( + windowingStrategy, stateInternalsFactory, (SystemReduceFn) systemReduceFn); + return doFn; + } + + + @Override + public void open() throws Exception { + + // might already be initialized from restoreTimers() + if (watermarkTimers == null) { + watermarkTimers = new HashSet<>(); + + watermarkTimersQueue = new PriorityQueue<>( + 10, + new Comparator<Tuple2<ByteBuffer, TimerInternals.TimerData>>() { + @Override + public int compare( + Tuple2<ByteBuffer, TimerInternals.TimerData> o1, + Tuple2<ByteBuffer, TimerInternals.TimerData> o2) { + return o1.f1.compareTo(o2.f1); + } + }); + } + + stateInternals = new FlinkStateInternals<>(getStateBackend(), keyCoder); + + // call super at the end because this will call getDoFn() which requires stateInternals + // to be set + super.open(); + } + + @Override + protected ExecutionContext.StepContext createStepContext() { + return new WindowDoFnOperator.StepContext(); + } + + private void registerEventTimeTimer(TimerInternals.TimerData timer) { + Tuple2<ByteBuffer, TimerInternals.TimerData> keyedTimer = + new Tuple2<>((ByteBuffer) getStateBackend().getCurrentKey(), timer); + if (watermarkTimers.add(keyedTimer)) { + watermarkTimersQueue.add(keyedTimer); + } + } + + private void deleteEventTimeTimer(TimerInternals.TimerData timer) { + Tuple2<ByteBuffer, TimerInternals.TimerData> keyedTimer = + new Tuple2<>((ByteBuffer) getStateBackend().getCurrentKey(), timer); + if (watermarkTimers.remove(keyedTimer)) { + watermarkTimersQueue.remove(keyedTimer); + } + + } + + @Override + public void processWatermark(Watermark mark) throws Exception { + this.currentWatermark = mark.getTimestamp(); + + boolean fire; + + do { + Tuple2<ByteBuffer, TimerInternals.TimerData> timer = watermarkTimersQueue.peek(); + if (timer != null && timer.f1.getTimestamp().getMillis() <= mark.getTimestamp()) { + fire = true; + + watermarkTimersQueue.remove(); + watermarkTimers.remove(timer); + + setKeyContext(timer.f0); + + doFnRunner.processElement(WindowedValue.valueInGlobalWindow( + KeyedWorkItems.<K, InputT>timersWorkItem( + stateInternals.getKey(), + Collections.singletonList(timer.f1)))); + + } else { + fire = false; + } + } while (fire); + + Instant watermarkHold = stateInternals.watermarkHold(); + + long outputWatermark = Math.min(currentWatermark, watermarkHold.getMillis()); + + output.emitWatermark(new Watermark(outputWatermark)); + } + + @Override + public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception { + StreamTaskState result = super.snapshotOperatorState(checkpointId, timestamp); + + AbstractStateBackend.CheckpointStateOutputView outputView = + getStateBackend().createCheckpointStateOutputView(checkpointId, timestamp); + + snapshotTimers(outputView); + + StateHandle<DataInputView> handle = outputView.closeAndGetHandle(); + + // this might overwrite stuff that super checkpointed + result.setOperatorState(handle); + + return result; + } + + @Override + public void restoreState(StreamTaskState state, long recoveryTimestamp) throws Exception { + super.restoreState(state, recoveryTimestamp); + + @SuppressWarnings("unchecked") + StateHandle<DataInputView> operatorState = + (StateHandle<DataInputView>) state.getOperatorState(); + + DataInputView in = operatorState.getState(getUserCodeClassloader()); + + restoreTimers(new DataInputViewWrapper(in)); + } + + private void restoreTimers(InputStream in) throws IOException { + DataInputStream dataIn = new DataInputStream(in); + int numWatermarkTimers = dataIn.readInt(); + + watermarkTimers = new HashSet<>(numWatermarkTimers); + watermarkTimersQueue = new PriorityQueue<>(Math.max(numWatermarkTimers, 1)); + + for (int i = 0; i < numWatermarkTimers; i++) { + int length = dataIn.readInt(); + byte[] keyBytes = new byte[length]; + dataIn.readFully(keyBytes); + TimerInternals.TimerData timerData = timerCoder.decode(dataIn, Coder.Context.NESTED); + Tuple2<ByteBuffer, TimerInternals.TimerData> keyedTimer = + new Tuple2<>(ByteBuffer.wrap(keyBytes), timerData); + if (watermarkTimers.add(keyedTimer)) { + watermarkTimersQueue.add(keyedTimer); + } + } + } + + private void snapshotTimers(OutputStream out) throws IOException { + DataOutputStream dataOut = new DataOutputStream(out); + dataOut.writeInt(watermarkTimersQueue.size()); + for (Tuple2<ByteBuffer, TimerInternals.TimerData> timer : watermarkTimersQueue) { + dataOut.writeInt(timer.f0.limit()); + dataOut.write(timer.f0.array(), 0, timer.f0.limit()); + timerCoder.encode(timer.f1, dataOut, Coder.Context.NESTED); + } + } + + /** + * {@link StepContext} for running {@link DoFn DoFns} on Flink. This does now allow + * accessing state or timer internals. + */ + protected class StepContext extends DoFnOperator.StepContext { + + @Override + public TimerInternals timerInternals() { + return new TimerInternals() { + @Override + public void setTimer(TimerData timerKey) { + if (timerKey.getDomain().equals(TimeDomain.EVENT_TIME)) { + registerEventTimeTimer(timerKey); + } else { + throw new UnsupportedOperationException("Processing-time timers not supported."); + } + } + + @Override + public void deleteTimer(TimerData timerKey) { + deleteEventTimeTimer(timerKey); + } + + @Override + public Instant currentProcessingTime() { + return Instant.now(); + } + + @Nullable + @Override + public Instant currentSynchronizedProcessingTime() { + return Instant.now(); + } + + @Override + public Instant currentInputWatermarkTime() { + return new Instant(currentWatermark); + } + + @Nullable + @Override + public Instant currentOutputWatermarkTime() { + return new Instant(currentWatermark); + } + }; + } + } + +} http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WorkItemKeySelector.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WorkItemKeySelector.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WorkItemKeySelector.java new file mode 100644 index 0000000..2bbed58 --- /dev/null +++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WorkItemKeySelector.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.KeyedWorkItem; +import org.apache.beam.sdk.util.WindowedValue; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.typeutils.GenericTypeInfo; +import org.apache.flink.api.java.typeutils.ResultTypeQueryable; + +import java.nio.ByteBuffer; + +/** + * {@link KeySelector} that retrieves a key from a {@link KeyedWorkItem}. This will return + * the key as encoded by the provided {@link Coder} in a {@link ByteBuffer}. This ensures + * that all key comparisons/hashing happen on the encoded form. + */ +public class WorkItemKeySelector<K, V> + implements KeySelector<WindowedValue<SingletonKeyedWorkItem<K, V>>, ByteBuffer>, ResultTypeQueryable<ByteBuffer> { + + private final Coder<K> keyCoder; + + public WorkItemKeySelector(Coder<K> keyCoder) { + this.keyCoder = keyCoder; + } + + @Override + public ByteBuffer getKey(WindowedValue<SingletonKeyedWorkItem<K, V>> value) throws Exception { + K key = value.getValue().key(); + byte[] keyBytes = CoderUtils.encodeToByteArray(keyCoder, key); + return ByteBuffer.wrap(keyBytes); + } + + @Override + public TypeInformation<ByteBuffer> getProducedType() { + return new GenericTypeInfo<>(ByteBuffer.class); + } +} http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/1de76b7a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/AbstractFlinkTimerInternals.java ---------------------------------------------------------------------- diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/AbstractFlinkTimerInternals.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/AbstractFlinkTimerInternals.java deleted file mode 100644 index a0b33f8..0000000 --- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/AbstractFlinkTimerInternals.java +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.beam.runners.flink.translation.wrappers.streaming.state; - -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.transforms.OldDoFn; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.util.TimerInternals; - -import org.joda.time.Instant; - -import java.io.IOException; -import java.io.Serializable; - -import javax.annotation.Nullable; - -/** - * An implementation of Beam's {@link TimerInternals}, that also provides serialization functionality. - * The latter is used when snapshots of the current state are taken, for fault-tolerance. - * */ -public abstract class AbstractFlinkTimerInternals<K, VIN> implements TimerInternals, Serializable { - private Instant currentInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - private Instant currentOutputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - - public void setCurrentInputWatermark(Instant watermark) { - checkIfValidInputWatermark(watermark); - this.currentInputWatermark = watermark; - } - - public void setCurrentOutputWatermark(Instant watermark) { - checkIfValidOutputWatermark(watermark); - this.currentOutputWatermark = watermark; - } - - private void setCurrentInputWatermarkAfterRecovery(Instant watermark) { - if (!currentInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MIN_VALUE)) { - throw new RuntimeException("Explicitly setting the input watermark is only allowed on " + - "initialization after recovery from a node failure. Apparently this is not " + - "the case here as the watermark is already set."); - } - this.currentInputWatermark = watermark; - } - - private void setCurrentOutputWatermarkAfterRecovery(Instant watermark) { - if (!currentOutputWatermark.isEqual(BoundedWindow.TIMESTAMP_MIN_VALUE)) { - throw new RuntimeException("Explicitly setting the output watermark is only allowed on " + - "initialization after recovery from a node failure. Apparently this is not " + - "the case here as the watermark is already set."); - } - this.currentOutputWatermark = watermark; - } - - @Override - public Instant currentProcessingTime() { - return Instant.now(); - } - - @Override - public Instant currentInputWatermarkTime() { - return currentInputWatermark; - } - - @Nullable - @Override - public Instant currentSynchronizedProcessingTime() { - // TODO - return null; - } - - @Override - public Instant currentOutputWatermarkTime() { - return currentOutputWatermark; - } - - private void checkIfValidInputWatermark(Instant newWatermark) { - if (currentInputWatermark.isAfter(newWatermark)) { - throw new IllegalArgumentException(String.format( - "Cannot set current input watermark to %s. Newer watermarks " + - "must be no earlier than the current one (%s).", - newWatermark, currentInputWatermark)); - } - } - - private void checkIfValidOutputWatermark(Instant newWatermark) { - if (currentOutputWatermark.isAfter(newWatermark)) { - throw new IllegalArgumentException(String.format( - "Cannot set current output watermark to %s. Newer watermarks " + - "must be no earlier than the current one (%s).", - newWatermark, currentOutputWatermark)); - } - } - - public void encodeTimerInternals(OldDoFn.ProcessContext context, - StateCheckpointWriter writer, - KvCoder<K, VIN> kvCoder, - Coder<? extends BoundedWindow> windowCoder) throws IOException { - if (context == null) { - throw new RuntimeException("The Context has not been initialized."); - } - - writer.setTimestamp(currentInputWatermark); - writer.setTimestamp(currentOutputWatermark); - } - - public void restoreTimerInternals(StateCheckpointReader reader, - KvCoder<K, VIN> kvCoder, - Coder<? extends BoundedWindow> windowCoder) throws IOException { - setCurrentInputWatermarkAfterRecovery(reader.getTimestamp()); - setCurrentOutputWatermarkAfterRecovery(reader.getTimestamp()); - } -}
