/** Produces a {@link RunnerApi.Components} which contains only primitive transforms. */ @VisibleForTesting static Collection<String> getPrimitiveTransformIds(RunnerApi.Components components) { Collection<String> ids = new LinkedHashSet<>(); for (Map.Entry<String, PTransform> transformEntry : components.getTransformsMap().entrySet()) { PTransform transform = transformEntry.getValue(); boolean isPrimitive = isPrimitiveTransform(transform); if (isPrimitive) { // Sometimes "primitive" transforms have sub-transforms (and even deeper-nested descendents), due to runners // either rewriting them in terms of runner-specific transforms, or SDKs constructing them in terms of other // underlying transforms (see https://issues.apache.org/jira/browse/BEAM-5441). // We consider any "leaf" descendents of these "primitive" transforms to be the true "primitives" that we // preserve here; in the common case, this is just the "primitive" itself, which has no descendents). Deque<String> transforms = new ArrayDeque<>(); transforms.push(transformEntry.getKey()); while (!transforms.isEmpty()) { String id = transforms.pop(); PTransform next = components.getTransformsMap().get(id); List<String> subtransforms = next.getSubtransformsList(); if (subtransforms.isEmpty()) { ids.add(id); } else { transforms.addAll(subtransforms); } } } } return ids; }
transformBuilder.addAllSubtransforms( transform .getSubtransformsList() .stream() .filter(id -> !viewTransforms.contains(id))
private static void validateTransform(String id, PTransform transform, Components components) { for (String subtransformId : transform.getSubtransformsList()) { checkArgument( components.containsTransforms(subtransformId),
assertThat(updatedSecond.getSubtransformsList(), contains("second_sub")); assertThat(updatedSecond.getSpec().getPayload(), equalTo(newPayload)); assertThat(updatedThird.getSubtransformsList(), contains("third_sub")); assertThat(updatedThird.getSpec().getPayload(), equalTo(newPayload));
equalTo("beam:composite")); assertThat( pipeline.getComponents().getTransformsOrThrow("root").getSubtransformsList(), contains("foo", "bar"));
private void removeDescendants(RunnerApi.Pipeline.Builder pipeline, String parentId) { RunnerApi.PTransform parentProto = pipeline.getComponents().getTransformsOrDefault(parentId, null); if (parentProto != null) { for (String childId : parentProto.getSubtransformsList()) { removeDescendants(pipeline, childId); pipeline.getComponentsBuilder().removeTransforms(childId); } pipeline .getComponentsBuilder() .putTransforms(parentId, parentProto.toBuilder().clearSubtransforms().build()); } }
/** * Remove all subtransforms of the provided transform recursively.A {@link PTransform} can be the * subtransform of only one enclosing transform. */ private static void removeSubtransforms(PTransform pt, Components.Builder target) { for (String subtransformId : pt.getSubtransformsList()) { PTransform subtransform = target.getTransformsOrThrow(subtransformId); removeSubtransforms(subtransform, target); target.removeTransforms(subtransformId); // TODO: remove PCollections not produced by 'pt' here. } }