ReplacementOutput mapping = originalToReplacement.get(output.getValue()); if (mapping != null) { if (this.equals(producers.get(mapping.getReplacement().getValue()))) { producerInput.remove(mapping.getReplacement().getValue()); producers.remove(mapping.getReplacement().getValue()); producers.put(mapping.getOriginal().getValue(), this); mapping.getReplacement(), mapping.getOriginal()); newOutputsBuilder.put(output.getKey(), mapping.getOriginal().getValue()); } else { newOutputsBuilder.put(output);
resultBuilder.put( replacementValue.getValue(), ReplacementOutput.of( mapped, TaggedPValue.of(replacementValue.getKey(), replacementValue.getValue()))); missingTags.remove(replacementValue.getKey());
Collections.singletonMap( replacementOutput.get(longs), ReplacementOutput.of( TaggedPValue.ofExpandedValue(output), TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
Collections.singletonMap( replacementOutput.get(longs), ReplacementOutput.of( TaggedPValue.ofExpandedValue(output), TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
Collections.singletonMap( replacementOutput, ReplacementOutput.of(TaggedPValue.ofExpandedValue(originalOutput), taggedReplacement)); hierarchy.replaceOutputs(replacementOutputs);
intsReplacement, equalTo( ReplacementOutput.of( TaggedPValue.of(intsTag, ints), TaggedPValue.of(intsTag, replacementInts)))); assertThat( strsReplacement, equalTo( ReplacementOutput.of( TaggedPValue.of(strsTag, strs), TaggedPValue.of(strsTag, replacementStrs)))); assertThat( moreIntsReplacement, equalTo( ReplacementOutput.of( TaggedPValue.of(moreIntsTag, moreInts), TaggedPValue.of(moreIntsTag, moreReplacementInts))));
private void testStreamingWriteOverride(PipelineOptions options, int expectedNumShards) { TestPipeline p = TestPipeline.fromOptions(options); StreamingShardedWriteFactory<Object, Void, Object> factory = new StreamingShardedWriteFactory<>(p.getOptions()); WriteFiles<Object, Void, Object> original = WriteFiles.to(new TestSink(tmpFolder.toString())); PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>> originalApplication = AppliedPTransform.of("writefiles", objs.expand(), Collections.emptyMap(), original, p); WriteFiles<Object, Void, Object> replacement = (WriteFiles<Object, Void, Object>) factory.getReplacementTransform(originalApplication).getTransform(); assertThat(replacement, not(equalTo((Object) original))); assertThat(replacement.getNumShardsProvider().get(), equalTo(expectedNumShards)); WriteFilesResult<Void> originalResult = objs.apply(original); WriteFilesResult<Void> replacementResult = objs.apply(replacement); Map<PValue, ReplacementOutput> res = factory.mapOutputs(originalResult.expand(), replacementResult); assertEquals(1, res.size()); assertEquals( originalResult.getPerDestinationOutputFilenames(), res.get(replacementResult.getPerDestinationOutputFilenames()).getOriginal().getValue()); }
@Test public void testMapOutputs() { PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3)); PCollection<Integer> output = input.apply("Map", MapElements.via(fn)); PCollection<Integer> reappliedOutput = input.apply("ReMap", MapElements.via(fn)); Map<PValue, ReplacementOutput> replacementMap = factory.mapOutputs(output.expand(), reappliedOutput); assertThat( replacementMap, Matchers.hasEntry( reappliedOutput, ReplacementOutput.of( TaggedPValue.ofExpandedValue(output), TaggedPValue.ofExpandedValue(reappliedOutput)))); }
@Test public void outputMapping() { final PCollectionList<String> inputList = PCollectionList.of(first).and(second).and(first).and(first); PCollection<String> original = inputList.apply(Flatten.pCollections()); PCollection<String> replacement = inputList.apply(new FlattenWithoutDuplicateInputs<>()); assertThat( factory.mapOutputs(original.expand(), replacement), Matchers.hasEntry( replacement, ReplacementOutput.of( TaggedPValue.ofExpandedValue(original), TaggedPValue.ofExpandedValue(replacement)))); } }
@Test public void mapOutputsSucceeds() { PCollection<Long> original = pipeline.apply("Original", GenerateSequence.from(0)); PCollection<Long> replacement = pipeline.apply("Replacement", GenerateSequence.from(0)); Map<PValue, ReplacementOutput> mapping = factory.mapOutputs(original.expand(), replacement); assertThat( mapping, Matchers.hasEntry( replacement, ReplacementOutput.of( TaggedPValue.ofExpandedValue(original), TaggedPValue.ofExpandedValue(replacement)))); }
@Override public Map<PValue, ReplacementOutput> mapOutputs( Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) { Map.Entry<TupleTag<?>, PValue> original = Iterables.getOnlyElement(outputs.entrySet()); Map.Entry<TupleTag<?>, PValue> replacement = Iterables.getOnlyElement(newOutput.expand().entrySet()); return Collections.singletonMap( newOutput, ReplacementOutput.of( TaggedPValue.of(original.getKey(), original.getValue()), TaggedPValue.of(replacement.getKey(), replacement.getValue()))); } }
@Override public Map<PValue, ReplacementOutput> mapOutputs( Map<TupleTag<?>, PValue> outputs, PCollection<Long> newOutput) { Map.Entry<TupleTag<?>, PValue> original = Iterables.getOnlyElement(outputs.entrySet()); Map.Entry<TupleTag<?>, PValue> replacement = Iterables.getOnlyElement(newOutput.expand().entrySet()); return Collections.singletonMap( newOutput, ReplacementOutput.of( TaggedPValue.of(original.getKey(), original.getValue()), TaggedPValue.of(replacement.getKey(), replacement.getValue()))); } }
public static Map<PValue, ReplacementOutput> singleton( Map<TupleTag<?>, PValue> original, PValue replacement) { Entry<TupleTag<?>, PValue> originalElement = Iterables.getOnlyElement(original.entrySet()); TupleTag<?> replacementTag = Iterables.getOnlyElement(replacement.expand().entrySet()).getKey(); return Collections.singletonMap( replacement, ReplacementOutput.of( TaggedPValue.of(originalElement.getKey(), originalElement.getValue()), TaggedPValue.of(replacementTag, replacement))); }
@Test public void singletonSucceeds() { Map<PValue, ReplacementOutput> replacements = ReplacementOutputs.singleton(ints.expand(), replacementInts); assertThat(replacements, Matchers.hasKey(replacementInts)); ReplacementOutput replacement = replacements.get(replacementInts); Map.Entry<TupleTag<?>, PValue> taggedInts = Iterables.getOnlyElement(ints.expand().entrySet()); assertThat(replacement.getOriginal().getTag(), equalTo(taggedInts.getKey())); assertThat(replacement.getOriginal().getValue(), equalTo(taggedInts.getValue())); assertThat(replacement.getReplacement().getValue(), equalTo(replacementInts)); }