@Override public PCollection<?> load(String id) throws Exception { checkState( pipeline != null, "%s Cannot rehydrate %s without a %s:" + " provide one via .withPipeline(...)", RehydratedComponents.class.getSimpleName(), PCollection.class.getSimpleName(), Pipeline.class.getSimpleName()); return PCollectionTranslation.fromProto( components.getPcollectionsOrThrow(id), pipeline, RehydratedComponents.this) .setName(id); } });
/** Creates a coder for a given PCollection id from the Proto definition. */ public static <T> Coder<WindowedValue<T>> instantiateCoder( String collectionId, RunnerApi.Components components) { PipelineNode.PCollectionNode collectionNode = PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId)); try { return WireCoders.instantiateRunnerWireCoder(collectionNode, components); } catch (IOException e) { throw new RuntimeException("Could not instantiate Coder", e); } }
/** Creates a coder for a given PCollection id from the Proto definition. */ public static <T> Coder<WindowedValue<T>> instantiateCoder( String collectionId, RunnerApi.Components components) { PipelineNode.PCollectionNode collectionNode = PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId)); try { return WireCoders.instantiateRunnerWireCoder(collectionNode, components); } catch (IOException e) { throw new RuntimeException("Could not instantiate Coder", e); } }
pipeline .getComponents() .getPcollectionsOrThrow(inputPCollectionId) .getWindowingStrategyId());
PCollectionNode longsOutput = PipelineNode.pCollection( longsOutputName, components.getPcollectionsOrThrow(longsOutputName)); String flattenOutputName = getOnlyElement( PCollectionNode flattenOutput = PipelineNode.pCollection( flattenOutputName, components.getPcollectionsOrThrow(flattenOutputName));
PCollectionNode collectionNode = PipelineNode.pCollection( outputCollectionId, components.getPcollectionsOrThrow(outputCollectionId)); Coder<WindowedValue<T>> outputCoder; try {
@Test public void forTransformsWithSubgraph() { Components components = Components.newBuilder() .putTransforms( "root", PTransform.newBuilder().putOutputs("output", "output.out").build()) .putPcollections( "output.out", RunnerApi.PCollection.newBuilder().setUniqueName("output.out").build()) .putTransforms( "consumer", PTransform.newBuilder().putInputs("input", "output.out").build()) .putTransforms( "ignored", PTransform.newBuilder().putInputs("input", "output.out").build()) .build(); QueryablePipeline pipeline = QueryablePipeline.forTransforms(ImmutableSet.of("root", "consumer"), components); assertThat( pipeline.getRootTransforms(), contains(PipelineNode.pTransform("root", components.getTransformsOrThrow("root")))); Set<PTransformNode> consumers = pipeline.getPerElementConsumers( PipelineNode.pCollection( "output.out", components.getPcollectionsOrThrow("output.out"))); assertThat( consumers, contains(PipelineNode.pTransform("consumer", components.getTransformsOrThrow("consumer")))); }
public Collection<UserStateReference> getUserStates(PTransformNode transform) { return getLocalUserStateNames(transform.getTransform()) .stream() .map( localName -> { String transformId = transform.getId(); PTransform transformProto = components.getTransformsOrThrow(transformId); // Get the main input PCollection id. String collectionId = transform .getTransform() .getInputsOrThrow( Iterables.getOnlyElement( Sets.difference( transform.getTransform().getInputsMap().keySet(), ImmutableSet.builder() .addAll(getLocalSideInputNames(transformProto)) .addAll(getLocalTimerNames(transformProto)) .build()))); PCollection collection = components.getPcollectionsOrThrow(collectionId); return UserStateReference.of( PipelineNode.pTransform(transformId, transformProto), localName, PipelineNode.pCollection(collectionId, collection)); }) .collect(Collectors.toSet()); }
/** * Tests that {@link QueryablePipeline#getPerElementConsumers(PCollectionNode)} returns a * transform that consumes the node more than once. */ @Test public void perElementConsumersWithConsumingMultipleTimes() { Pipeline p = Pipeline.create(); PCollection<Long> longs = p.apply("BoundedRead", Read.from(CountingSource.upTo(100L))); PCollectionList.of(longs).and(longs).and(longs).apply("flatten", Flatten.pCollections()); Components components = PipelineTranslation.toProto(p).getComponents(); // This breaks if the way that IDs are assigned to PTransforms changes in PipelineTranslation String readOutput = getOnlyElement(components.getTransformsOrThrow("BoundedRead").getOutputsMap().values()); QueryablePipeline qp = QueryablePipeline.forPrimitivesIn(components); Set<PTransformNode> consumers = qp.getPerElementConsumers( PipelineNode.pCollection(readOutput, components.getPcollectionsOrThrow(readOutput))); assertThat(consumers.size(), equalTo(1)); assertThat( getOnlyElement(consumers).getTransform().getSpec().getUrn(), equalTo(PTransformTranslation.FLATTEN_TRANSFORM_URN)); }
/** * Returns the {@link SideInputReference SideInputReferences} that the provided transform consumes * as side inputs. */ public Collection<SideInputReference> getSideInputs(PTransformNode transform) { return getLocalSideInputNames(transform.getTransform()) .stream() .map( localName -> { String transformId = transform.getId(); PTransform transformProto = components.getTransformsOrThrow(transformId); String collectionId = transform.getTransform().getInputsOrThrow(localName); PCollection collection = components.getPcollectionsOrThrow(collectionId); return SideInputReference.of( PipelineNode.pTransform(transformId, transformProto), localName, PipelineNode.pCollection(collectionId, collection)); }) .collect(Collectors.toSet()); }
@Before public void setup() throws IOException { Pipeline p = Pipeline.create(); PCollection<String> foo = p.apply("createFoo", Create.of("1", "2", "3")) .apply("windowFoo", Window.into(FixedWindows.of(Duration.standardMinutes(5L)))); PCollection<Integer> bar = p.apply("bar", Create.of(1, 2, 3)); SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); String fooId = sdkComponents.registerPCollection(foo); String barId = sdkComponents.registerPCollection(bar); baseComponents = sdkComponents.toComponents(); fooPC = PipelineNode.pCollection(fooId, baseComponents.getPcollectionsOrThrow(fooId)); barPC = PipelineNode.pCollection(barId, baseComponents.getPcollectionsOrThrow(barId)); outputBundles = new ArrayList<>(); factory = BundleFactoryOutputReceiverFactory.create( bundleFactory, baseComponents, outputBundles::add); }
/** Create a user state reference from a UserStateId proto and components. */ public static UserStateReference fromUserStateId( UserStateId userStateId, RunnerApi.Components components) { PTransform transform = components.getTransformsOrThrow(userStateId.getTransformId()); String mainInputCollectionId; try { mainInputCollectionId = transform.getInputsOrThrow(ParDoTranslation.getMainInputName(transform)); } catch (IOException e) { throw new RuntimeException(e); } return UserStateReference.of( PipelineNode.pTransform(userStateId.getTransformId(), transform), userStateId.getLocalName(), PipelineNode.pCollection( mainInputCollectionId, components.getPcollectionsOrThrow(mainInputCollectionId))); }
public static WindowingStrategy getWindowingStrategy( String pCollectionId, RunnerApi.Components components) { RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow( components.getPcollectionsOrThrow(pCollectionId).getWindowingStrategyId()); final WindowingStrategy<?, ?> windowingStrategy; try { return WindowingStrategyTranslation.fromProto( windowingStrategyProto, RehydratedComponents.forComponents(components)); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException( String.format( "Unable to hydrate windowing strategy %s for %s.", windowingStrategyProto, pCollectionId), e); } } }
@Test public void toAndFromProto() throws IOException { SdkComponents components = SdkComponents.create(spec.getTransform().getPipeline().getOptions()); RunnerApi.PTransform converted = convert(spec, components); Components protoComponents = components.toComponents(); // Sanity checks assertThat(converted.getInputsCount(), equalTo(spec.getTransform().getInputs().size())); assertThat(converted.getOutputsCount(), equalTo(spec.getTransform().getOutputs().size())); assertThat(converted.getSubtransformsCount(), equalTo(spec.getChildren().size())); assertThat(converted.getUniqueName(), equalTo(spec.getTransform().getFullName())); for (PValue inputValue : spec.getTransform().getInputs().values()) { PCollection<?> inputPc = (PCollection<?>) inputValue; protoComponents.getPcollectionsOrThrow(components.registerPCollection(inputPc)); } for (PValue outputValue : spec.getTransform().getOutputs().values()) { PCollection<?> outputPc = (PCollection<?>) outputValue; protoComponents.getPcollectionsOrThrow(components.registerPCollection(outputPc)); } }
public static WindowingStrategy getWindowingStrategy( String pCollectionId, RunnerApi.Components components) { RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow( components.getPcollectionsOrThrow(pCollectionId).getWindowingStrategyId()); final WindowingStrategy<?, ?> windowingStrategy; try { return WindowingStrategyTranslation.fromProto( windowingStrategyProto, RehydratedComponents.forComponents(components)); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException( String.format( "Unable to hydrate windowing strategy %s for %s.", windowingStrategyProto, pCollectionId), e); } } }
@Test public void registerPCollectionExistingNameCollision() throws IOException { PCollection<Long> pCollection = pipeline.apply("FirstCount", GenerateSequence.from(0)).setName("foo"); String firstId = components.registerPCollection(pCollection); PCollection<Long> duplicate = pipeline.apply("SecondCount", GenerateSequence.from(0)).setName("foo"); String secondId = components.registerPCollection(duplicate); assertThat(firstId, equalTo("foo")); assertThat(secondId, containsString("foo")); assertThat(secondId, not(equalTo("foo"))); components.toComponents().getPcollectionsOrThrow(firstId); components.toComponents().getPcollectionsOrThrow(secondId); }
@Override public <OutputT> FnDataReceiver<OutputT> create(String pCollectionId) { PCollectionNode pcollection = PipelineNode.pCollection(pCollectionId, components.getPcollectionsOrThrow(pCollectionId)); return create(pcollection); }
/** Create a side input reference from a SideInputId proto and components. */ public static SideInputReference fromSideInputId( SideInputId sideInputId, RunnerApi.Components components) { String transformId = sideInputId.getTransformId(); String localName = sideInputId.getLocalName(); String collectionId = components.getTransformsOrThrow(transformId).getInputsOrThrow(localName); PTransform transform = components.getTransformsOrThrow(transformId); PCollection collection = components.getPcollectionsOrThrow(collectionId); return SideInputReference.of( PipelineNode.pTransform(transformId, transform), localName, PipelineNode.pCollection(collectionId, collection)); }
@Test public void registerPCollection() throws IOException { PCollection<Long> pCollection = pipeline.apply(GenerateSequence.from(0)).setName("foo"); String id = components.registerPCollection(pCollection); assertThat(id, equalTo("foo")); components.toComponents().getPcollectionsOrThrow(id); }
public static RunnerApi.PCollection getMainInput( RunnerApi.PTransform ptransform, Components components) throws IOException { checkArgument( ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN), "Unexpected payload type %s", ptransform.getSpec().getUrn()); return components.getPcollectionsOrThrow( ptransform.getInputsOrThrow(getMainInputName(ptransform))); }