/** Creates a mapping from PCollection id to output tag integer. */ public static BiMap<String, Integer> createOutputMap(Iterable<String> localOutputs) { ImmutableBiMap.Builder<String, Integer> builder = ImmutableBiMap.builder(); int outputIndex = 0; // sort localOutputs for stable indexing for (String tag : Sets.newTreeSet(localOutputs)) { builder.put(tag, outputIndex); outputIndex++; } return builder.build(); }
FlinkStreamingPortablePipelineTranslator() { ImmutableMap.Builder<String, PTransformTranslator<StreamingTranslationContext>> translatorMap = ImmutableMap.builder(); translatorMap.put(PTransformTranslation.FLATTEN_TRANSFORM_URN, this::translateFlatten); translatorMap.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, this::translateGroupByKey); translatorMap.put(PTransformTranslation.IMPULSE_TRANSFORM_URN, this::translateImpulse); translatorMap.put(STREAMING_IMPULSE_TRANSFORM_URN, this::translateStreamingImpulse); translatorMap.put( PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN, this::translateAssignWindows); translatorMap.put(ExecutableStage.URN, this::translateExecutableStage); translatorMap.put(PTransformTranslation.RESHUFFLE_URN, this::translateReshuffle); this.urnToTransformTranslator = translatorMap.build(); }
@Override public Iterable<Class<? extends PipelineRunner<?>>> getPipelineRunners() { return ImmutableList.of(FlinkRunner.class, TestFlinkRunner.class); } }
private Iterable<WindowedValue<String>> stripStreamRecord(Iterable<?> input) { return FluentIterable.from(input) .filter(o -> o instanceof StreamRecord) .transform( new Function<Object, WindowedValue<String>>() { @Nullable @Override @SuppressWarnings({"unchecked", "rawtypes"}) public WindowedValue<String> apply(@Nullable Object o) { if (o instanceof StreamRecord) { return (WindowedValue<String>) ((StreamRecord) o).getValue(); } throw new RuntimeException("unreachable"); } }); }
private <T, W extends BoundedWindow> SideInputHandler<T, W> forIterableSideInput( List<WindowedValue<T>> broadcastVariable, Coder<T> elementCoder, Coder<W> windowCoder) { ImmutableMultimap.Builder<Object, T> windowToValuesBuilder = ImmutableMultimap.builder(); for (WindowedValue<T> windowedValue : broadcastVariable) { for (BoundedWindow boundedWindow : windowedValue.getWindows()) { @SuppressWarnings("unchecked") W window = (W) boundedWindow; windowToValuesBuilder.put(windowCoder.structuralValue(window), windowedValue.getValue()); } } ImmutableMultimap<Object, T> windowToValues = windowToValuesBuilder.build(); return new SideInputHandler<T, W>() { @Override public Iterable<T> get(byte[] key, W window) { return windowToValues.get(windowCoder.structuralValue(window)); } @Override public Coder<T> resultCoder() { return elementCoder; } }; }
@Override public BoundedWindow window() { return Iterables.getOnlyElement(value.getWindows()); } }
@Override public Iterable<T> get(byte[] key, W window) { return windowToValues.get(windowCoder.structuralValue(window)); }
@SuppressWarnings("unchecked") public MultiOutputOutputManagerFactory( TupleTag<OutputT> mainTag, Coder<WindowedValue<OutputT>> mainCoder) { this( mainTag, new HashMap<>(), ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder() .put(mainTag, (Coder) mainCoder) .build(), ImmutableMap.<TupleTag<?>, Integer>builder().put(mainTag, 0).build()); }
@Override public Iterable<Class<? extends PipelineOptions>> getPipelineOptions() { return ImmutableList.of(FlinkPipelineOptions.class); } }
@SuppressWarnings("Guava") static <T> Iterable<WindowedValue<T>> stripStreamRecordFromWindowedValue(Iterable<Object> input) { return FluentIterable.from(input) .filter( o -> o instanceof StreamRecord && ((StreamRecord) o).getValue() instanceof WindowedValue) .transform( new Function<Object, WindowedValue<T>>() { @Nullable @Override @SuppressWarnings({"unchecked", "rawtypes"}) public WindowedValue<T> apply(@Nullable Object o) { if (o instanceof StreamRecord && ((StreamRecord) o).getValue() instanceof WindowedValue) { return (WindowedValue) ((StreamRecord) o).getValue(); } throw new RuntimeException("unreachable"); } }); } }
private <K, V, W extends BoundedWindow> SideInputHandler<V, W> forMultimapSideInput( List<WindowedValue<KV<K, V>>> broadcastVariable, Coder<K> keyCoder, Coder<V> valueCoder, Coder<W> windowCoder) { ImmutableMultimap.Builder<SideInputKey, V> multimap = ImmutableMultimap.builder(); for (WindowedValue<KV<K, V>> windowedValue : broadcastVariable) { K key = windowedValue.getValue().getKey(); V value = windowedValue.getValue().getValue(); for (BoundedWindow boundedWindow : windowedValue.getWindows()) { @SuppressWarnings("unchecked") W window = (W) boundedWindow; multimap.put( SideInputKey.of(keyCoder.structuralValue(key), windowCoder.structuralValue(window)), value); } } return new MultimapSideInputHandler<>(multimap.build(), keyCoder, valueCoder, windowCoder); }
@SuppressWarnings("unchecked") <T extends PValue> T getInput(PTransform<T, ?> transform) { return (T) Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(currentTransform)); }
/** Creates a batch translator. */ public static FlinkBatchPortablePipelineTranslator createTranslator() { ImmutableMap.Builder<String, PTransformTranslator> translatorMap = ImmutableMap.builder(); translatorMap.put( PTransformTranslation.FLATTEN_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateFlatten); translatorMap.put( PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateGroupByKey); translatorMap.put( PTransformTranslation.IMPULSE_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateImpulse); translatorMap.put( PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateAssignWindows); translatorMap.put( ExecutableStage.URN, FlinkBatchPortablePipelineTranslator::translateExecutableStage); translatorMap.put( PTransformTranslation.RESHUFFLE_URN, FlinkBatchPortablePipelineTranslator::translateReshuffle); return new FlinkBatchPortablePipelineTranslator(translatorMap.build()); }
@Override public List<? extends Coder<?>> getCoderArguments() { return ImmutableList.of(keyCoder, elemCoder, windowCoder); }
@SuppressWarnings("unchecked") public <T extends PValue> T getInput(PTransform<T, ?> transform) { return (T) Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(currentTransform)); }
@Override public Map< ? extends Class<? extends PTransform>, ? extends PTransformTranslation.TransformPayloadTranslator> getTransformPayloadTranslators() { return ImmutableMap .<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder() .put( CreateStreamingFlinkView.CreateFlinkPCollectionView.class, new CreateStreamingFlinkViewPayloadTranslator()) .put( SplittableParDoViaKeyedWorkItems.ProcessElements.class, new SplittableParDoProcessElementsTranslator()) .put( SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class, new SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator()) .build(); } }
@Override public AccumT addInput( K key, AccumT accumulator, AccumT value, PipelineOptions options, SideInputReader sideInputReader, Collection<? extends BoundedWindow> windows) { return combineFnRunner.mergeAccumulators( ImmutableList.of(accumulator, value), options, sideInputReader, windows); }
@SuppressWarnings("unchecked") public <T extends PValue> T getOutput(PTransform<?, T> transform) { return (T) Iterables.getOnlyElement(currentTransform.getOutputs().values()); }
@Test public void testOutputMapCreation() { List<String> outputs = Arrays.asList("output1", "output2", "output3"); BiMap<String, Integer> outputMap = FlinkPipelineTranslatorUtils.createOutputMap(outputs); Map<Object, Object> expected = ImmutableMap.builder().put("output1", 0).put("output2", 1).put("output3", 2).build(); assertThat(outputMap, is(expected)); } }
@SuppressWarnings("unchecked") <T extends PValue> T getOutput(PTransform<?, T> transform) { return (T) Iterables.getOnlyElement(currentTransform.getOutputs().values()); } }