lukecwik commented on a change in pull request #12016:
URL: https://github.com/apache/beam/pull/12016#discussion_r447133216
##########
File path:
runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
##########
@@ -113,6 +113,8 @@
// SplittableParDoComponents
public static final String SPLITTABLE_PAIR_WITH_RESTRICTION_URN =
"beam:transform:sdf_pair_with_restriction:v1";
+ public static final String SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN =
+ "beam:transform:sdf_truncate_sized_restrictions:v1";
Review comment:
please add a checkState in static block like the others below
##########
File path:
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
##########
@@ -306,6 +308,32 @@ public static void
invokeSplitRestriction(DoFnInvoker.ArgumentProvider argumentP
}
}
+ /**
+ * Default implementation of {@link DoFn.TruncateSizedRestriction}, for
delegation by bytebuddy.
+ */
+ public static class DefaultTruncateSizedRestriction {
+ private final boolean isBoundedPerElement;
+
+ DefaultTruncateSizedRestriction(DoFnSignature doFnSignature) {
+ this.isBoundedPerElement =
doFnSignature.isBoundedPerElement().equals(IsBounded.BOUNDED);
+ }
+ /** Return the current restriction if it's bounded.Otherwise, return null.
*/
+ @SuppressWarnings("unused")
+ public void invokeTruncateSizedRestriction(DoFnInvoker.ArgumentProvider
argumentProvider) {
+ boolean isBounded;
+ try {
+ isBounded = argumentProvider.restrictionTracker().isBounded();
+ } catch (NotImplementedException expectedException) {
+ isBounded = this.isBoundedPerElement;
+ }
+ if (isBounded) {
+
argumentProvider.outputReceiver(null).output(argumentProvider.restriction());
+ } else {
+ argumentProvider.outputReceiver(null).output(null);
+ }
+ }
+ }
Review comment:
```suggestion
/**
* Default implementation of {@link DoFn.TruncateRestriction}, for
delegation by bytebuddy.
*/
public static class DefaultTruncateRestriction {
private final boolean isBoundedPerElement;
DefaultTruncateRestriction(DoFnSignature doFnSignature) {
this.isBoundedPerElement =
doFnSignature.isBoundedPerElement().equals(IsBounded.BOUNDED);
}
/** Return the current restriction if it's bounded.Otherwise, return
null. */
@SuppressWarnings("unused")
public void invokeTruncateRestriction(DoFnInvoker.ArgumentProvider
argumentProvider) {
boolean isBounded = this.isBoundedPerElement;
try {
isBounded = argumentProvider.restrictionTracker().isBounded();
} catch (NotImplementedException expectedException) {
}
if (isBounded) {
argumentProvider.outputReceiver(null).output(argumentProvider.restriction());
}
}
}
```
##########
File path:
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/GrowableOffsetRangeTracker.java
##########
@@ -140,4 +140,9 @@ public Progress getProgress() {
totalWork.subtract(workRemaining,
MathContext.DECIMAL128).doubleValue(),
workRemaining.doubleValue());
}
+
+ @Override
+ public boolean isBounded() {
+ return false;
Review comment:
Shouldn't this be true if `range.getTo() != Long.MAX_VALUE`?
##########
File path:
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
##########
@@ -574,6 +578,99 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
};
}
break;
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
+ if ((doFnSignature.truncateSizedRestriction() != null
+ && doFnSignature.truncateSizedRestriction().observesWindow())
+ || (doFnSignature.newTracker() != null &&
doFnSignature.newTracker().observesWindow())
+ || (doFnSignature.getSize() != null &&
doFnSignature.getSize().observesWindow())
+ || !sideInputMapping.isEmpty()) {
+ mainInputConsumer =
this::processElementForWindowObservingTruncateSizedRestriction;
+ // OutputT == RestrictionT
+ this.processContext =
+ new WindowObservingProcessBundleContext() {
+ @Override
+ public void outputWithTimestamp(OutputT output, Instant
timestamp) {
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this,
+
PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN
+ + "/GetSize") {
Review comment:
```suggestion
+ "/TruncateRestriction") {
```
##########
File path: sdks/python/apache_beam/runners/worker/bundle_processor.py
##########
@@ -1428,6 +1428,36 @@ def process(self, element_restriction, *args, **kwargs):
return _create_sdf_operation(SplitAndSizeRestrictions, *args)
+@BeamTransformFactory.register_urn(
+ common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn,
+ beam_runner_api_pb2.ParDoPayload)
+def create_truncate_sized_restriction(*args):
+ class TruncateAndSizeRestriction(beam.DoFn):
+ def __init__(self, fn, restriction_provider, watermark_estimator_provider):
+ self.restriction_provider = restriction_provider
+ self.dofn_signature = common.DoFnSignature(fn)
+
+ def process(self, element_restrictin, *args, **kwargs):
+ ((element, (restriction, estimator_state)), _) = element_restrictin
+ try:
+ truncated_restriciton = self.restriction_provider.truncate(
Review comment:
We should be returning the same as split_and_size allowing for 0
restrictions to be returned.
##########
File path: sdks/python/apache_beam/runners/worker/bundle_processor.py
##########
@@ -1428,6 +1428,36 @@ def process(self, element_restriction, *args, **kwargs):
return _create_sdf_operation(SplitAndSizeRestrictions, *args)
+@BeamTransformFactory.register_urn(
+ common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn,
+ beam_runner_api_pb2.ParDoPayload)
+def create_truncate_sized_restriction(*args):
+ class TruncateAndSizeRestriction(beam.DoFn):
+ def __init__(self, fn, restriction_provider, watermark_estimator_provider):
+ self.restriction_provider = restriction_provider
+ self.dofn_signature = common.DoFnSignature(fn)
+
+ def process(self, element_restrictin, *args, **kwargs):
+ ((element, (restriction, estimator_state)), _) = element_restrictin
+ try:
+ truncated_restriciton = self.restriction_provider.truncate(
Review comment:
truncated_restriciton -> truncated_restriction
##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/io/Read.java
##########
@@ -500,6 +505,19 @@ public void splitRestriction(
return new UnboundedSourceAsSDFRestrictionTracker(restriction,
pipelineOptions);
}
+ @TruncateSizedRestriction
+ public void truncateSizedRestriction(
+ @Restriction UnboundedSourceRestriction<OutputT, CheckpointT>
restriction,
+ OutputReceiver<UnboundedSourceRestriction<OutputT, CheckpointT>>
receiver) {
+ try {
+ restriction.getCheckpoint().finalizeCheckpoint();
+ } catch (Exception e) {
+ LOG.warn("Failed to finalize CheckpointMark {}",
restriction.getWatermark());
+ } finally {
+ receiver.output(null);
+ }
Review comment:
I don't think we produce any outputs and finalization is handled already
by the bundleFinalizer.afterBundleCommit callback in `@ProcessElement`
```suggestion
```
##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
##########
@@ -1034,6 +1034,47 @@ public Duration getAllowedTimestampSkew() {
* <li>If one of its arguments is tagged with the {@link Timestamp}
annotation, then it will be
* passed the timestamp of the current element being processed; the
argument must be of type
* {@link Instant}.
+ * <li>If one of its arguments is a subtype of {@link BoundedWindow}, then
it will be passed the
+ * window of the current element. When applied by {@link ParDo} the
subtype of {@link
+ * BoundedWindow} must match the type of windows on the input {@link
PCollection}. If the
+ * window is not accessed a runner may perform additional
optimizations.
+ * <li>If one of its arguments is of type {@link PaneInfo}, then it will
be passed information
+ * about the current triggering pane.
+ * <li>If one of the parameters is of type {@link PipelineOptions}, then
it will be passed the
+ * options for the current pipeline.
+ * </ul>
+ */
+ @Documented
+ @Retention(RetentionPolicy.RUNTIME)
+ @Target(ElementType.METHOD)
+ @Experimental(Kind.SPLITTABLE_DO_FN)
+ public @interface SplitRestriction {}
+
+ /**
+ * Annotation for the method that truncates restriction of a <a
Review comment:
We should add a comment to the `@ProcessElement` javadoc comment stating
something like:
`If may define a TruncateRestriction method to override the default
implementation where the default is ...`
##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
##########
@@ -1034,6 +1034,47 @@ public Duration getAllowedTimestampSkew() {
* <li>If one of its arguments is tagged with the {@link Timestamp}
annotation, then it will be
* passed the timestamp of the current element being processed; the
argument must be of type
* {@link Instant}.
+ * <li>If one of its arguments is a subtype of {@link BoundedWindow}, then
it will be passed the
+ * window of the current element. When applied by {@link ParDo} the
subtype of {@link
+ * BoundedWindow} must match the type of windows on the input {@link
PCollection}. If the
+ * window is not accessed a runner may perform additional
optimizations.
+ * <li>If one of its arguments is of type {@link PaneInfo}, then it will
be passed information
+ * about the current triggering pane.
+ * <li>If one of the parameters is of type {@link PipelineOptions}, then
it will be passed the
+ * options for the current pipeline.
+ * </ul>
+ */
+ @Documented
+ @Retention(RetentionPolicy.RUNTIME)
+ @Target(ElementType.METHOD)
+ @Experimental(Kind.SPLITTABLE_DO_FN)
+ public @interface SplitRestriction {}
+
+ /**
+ * Annotation for the method that truncates restriction of a <a
+ * href="https://s.apache.org/splittable-do-fn";>splittable</a> {@link DoFn}
into bounded one to be
+ * processed when pipeline starts to drain.
+ *
+ * <p>This method is used to perform truncating the restriction before
actual processing.
+ *
+ * <p>Signature: {@code void truncateSizedRestriction(<arguments>);}
+ *
+ * <p>This method must satisfy the following constraints:
+ *
+ * <ul>
+ * <li>If one of the arguments is of type {@link OutputReceiver}, then it
will be passed an
+ * output receiver for outputting the splits. All splits must be
output through this
+ * parameter.
+ * <li>If one of its arguments is tagged with the {@link Element}
annotation, then it will be
+ * passed the current element being processed; the argument must be of
type {@code InputT}.
+ * Note that automatic conversion of {@link Row}s and {@link
FieldAccess} parameters are
+ * currently unsupported.
+ * <li>If one of its arguments is tagged with the {@link Restriction}
annotation, then it will
+ * be passed the current restriction being processed; the argument
must be of type {@code
+ * RestrictionT}.
+ * <li>If one of its arguments is tagged with the {@link Timestamp}
annotation, then it will be
+ * passed the timestamp of the current element being processed; the
argument must be of type
+ * {@link Instant}.
Review comment:
```suggestion
* Annotation for the method that truncates the restriction of a <a
* href="https://s.apache.org/splittable-do-fn";>splittable</a> {@link
DoFn} into a bounded one to be
* processed when pipeline starts to drain.
*
* <p>This method is used to perform truncation of the restriction while
it is not actively being processed.
*
* <p>Signature: {@code void truncateRestriction(<arguments>);}
*
* <p>This method must satisfy the following constraints:
*
* <ul>
* <li>If one of the arguments is of type {@link OutputReceiver}, then
it will be passed an
* output receiver for outputting the truncated restrictions. All
truncated restrictions must be output through this
* parameter.
* <li>If one of its arguments is tagged with the {@link Element}
annotation, then it will be
* passed the current element being processed; the argument must be
of type {@code InputT}.
* Note that automatic conversion of {@link Row}s and {@link
FieldAccess} parameters are
* currently unsupported.
* <li>If one of its arguments is tagged with the {@link Restriction}
annotation, then it will
* be passed the current restriction being processed; the argument
must be of type {@code
* RestrictionT}.
* <li>If one of its arguments is tagged with the {@link Timestamp}
annotation, then it will be
* passed the timestamp of the current element being processed; the
argument must be of type
* {@link Instant}.
```
##########
File path:
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
##########
@@ -99,6 +100,10 @@
*/
public abstract void checkDone() throws IllegalStateException;
+ public boolean isBounded() throws NotImplementedException {
Review comment:
We should use the return type `IsBounded`.
##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
##########
@@ -1051,7 +1092,7 @@ public Duration getAllowedTimestampSkew() {
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Experimental(Kind.SPLITTABLE_DO_FN)
- public @interface SplitRestriction {}
+ public @interface TruncateSizedRestriction {}
Review comment:
The DoFn doesn't need to know whether it is sized or not. How about
`TruncateRestriction`?
##########
File path:
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
##########
@@ -574,6 +578,99 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
};
}
break;
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
+ if ((doFnSignature.truncateSizedRestriction() != null
+ && doFnSignature.truncateSizedRestriction().observesWindow())
+ || (doFnSignature.newTracker() != null &&
doFnSignature.newTracker().observesWindow())
+ || (doFnSignature.getSize() != null &&
doFnSignature.getSize().observesWindow())
+ || !sideInputMapping.isEmpty()) {
+ mainInputConsumer =
this::processElementForWindowObservingTruncateSizedRestriction;
+ // OutputT == RestrictionT
+ this.processContext =
+ new WindowObservingProcessBundleContext() {
+ @Override
+ public void outputWithTimestamp(OutputT output, Instant
timestamp) {
+ double size =
+ doFnInvoker.invokeGetSize(
Review comment:
invokeGetSize?
##########
File path:
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
##########
@@ -574,6 +578,99 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
};
}
break;
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
+ if ((doFnSignature.truncateSizedRestriction() != null
+ && doFnSignature.truncateSizedRestriction().observesWindow())
+ || (doFnSignature.newTracker() != null &&
doFnSignature.newTracker().observesWindow())
+ || (doFnSignature.getSize() != null &&
doFnSignature.getSize().observesWindow())
+ || !sideInputMapping.isEmpty()) {
+ mainInputConsumer =
this::processElementForWindowObservingTruncateSizedRestriction;
+ // OutputT == RestrictionT
+ this.processContext =
Review comment:
Since this shares the same ProcessBundleContext as GetSize (except for
the error context string), can we make this a class called
SizedRestrictionWindowObservingProcessBundleContext that extends
WindowObservingProcessBundleContext and is used in both places.
Ditto for the non window observing one as well.
##########
File path:
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
##########
@@ -99,6 +100,10 @@
*/
public abstract void checkDone() throws IllegalStateException;
+ public boolean isBounded() throws NotImplementedException {
+ throw new NotImplementedException("isBounded is not implemented.");
Review comment:
If we want to make this optional, we should follow the `HasProgress`
style and use an interface like `HasBoundedness`
We could also make this a required method and update the changelog stating
that this is a breaking change. This would remove the "default" inference on
the `@UnboundedPerElement`/`@BoundedPerElement` annotations on the DoFn.
Depending on the default ordering we want, we'll want to update update the
restriction trackers in PeriodicSequence and Watch as well (note that Watch
will have to be updated regardless since it has a tracker that converts
unbounded restrictions to bounded ones).
##########
File path:
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
##########
@@ -99,6 +100,10 @@
*/
public abstract void checkDone() throws IllegalStateException;
+ public boolean isBounded() throws NotImplementedException {
+ throw new NotImplementedException("isBounded is not implemented.");
Review comment:
We should add a comment stating the unbounded restriction trackers can
only be used with `@UnboundedPerElement` SDFs. We should add a comment on each
restriction tracker whether it could be unbounded.
##########
File path:
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
##########
@@ -608,17 +705,6 @@ public void accept(WindowedValue input) throws Exception {
pCollectionConsumerRegistry.register(
pTransform.getInputsOrThrow(mainInput), pTransformId, (FnDataReceiver)
mainInputConsumer);
- switch (pTransform.getSpec().getUrn()) {
Review comment:
Thanks
##########
File path:
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
##########
@@ -858,6 +1008,10 @@ public static WindowedSplitResult forRoots(
private void processElementForWindowObservingSizedElementAndRestriction(
WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>,
Double>> elem) {
currentElement = elem.withValue(elem.getValue().getKey().getKey());
+ // No need to process if current restriction is null.
+ if (elem.getValue().getKey().getValue().getKey() == null) {
Review comment:
I don't think this is the right way to go. The output received should
produce 0 outputs instead of null.
##########
File path:
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
##########
@@ -829,6 +917,68 @@ private void
processElementForWindowObservingSplitRestriction(
this.stateAccessor.finalizeState();
}
+ private void processElementForTruncateSizedRestriction(
+ WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>,
Double>> elem) {
+ currentElement = elem.withValue(elem.getValue().getKey().getKey());
+ currentRestriction = elem.getValue().getKey().getValue().getKey();
+ currentWatermarkEstimatorState =
elem.getValue().getKey().getValue().getValue();
+ currentTracker =
Review comment:
It looks like currentTracker isn't being set within both
`processElementForSplitRestriction` methods. Could you fix that in this PR?
##########
File path: sdks/python/apache_beam/runners/worker/bundle_processor.py
##########
@@ -1428,6 +1428,36 @@ def process(self, element_restriction, *args, **kwargs):
return _create_sdf_operation(SplitAndSizeRestrictions, *args)
+@BeamTransformFactory.register_urn(
+ common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn,
+ beam_runner_api_pb2.ParDoPayload)
+def create_truncate_sized_restriction(*args):
+ class TruncateAndSizeRestriction(beam.DoFn):
+ def __init__(self, fn, restriction_provider, watermark_estimator_provider):
+ self.restriction_provider = restriction_provider
+ self.dofn_signature = common.DoFnSignature(fn)
+
+ def process(self, element_restrictin, *args, **kwargs):
Review comment:
element_restrictin -> element_restriction
##########
File path: sdks/python/apache_beam/transforms/core.py
##########
@@ -320,6 +320,22 @@ def split_and_size(self, element, restriction):
for part in self.split(element, restriction):
yield part, self.restriction_size(element, part)
+ def truncate(self, element, restriction):
+ """Truncate the given restriction into finite amount of work when the
+ pipeline starts to drain.
+
+ By default, if the restriction is bounded, it will return the entire
current
+ restriction. If the restriction is unbounded, it will return None.
Review comment:
```suggestion
By default, if the restriction is bounded, it will return the entire
restriction. If the restriction is unbounded, it will not return
anything.
```
##########
File path: sdks/python/apache_beam/transforms/core.py
##########
@@ -320,6 +320,22 @@ def split_and_size(self, element, restriction):
for part in self.split(element, restriction):
yield part, self.restriction_size(element, part)
+ def truncate(self, element, restriction):
+ """Truncate the given restriction into finite amount of work when the
+ pipeline starts to drain.
+
+ By default, if the restriction is bounded, it will return the entire
current
+ restriction. If the restriction is unbounded, it will return None.
+
+ The method throws NotImplementError when RestrictionTracker.is_bounded() is
+ not implemented.
+
+ It's recommended to implement this API if more granularity is required.
+ """
+ restriction_tracker = self.create_tracker(restriction)
+ if restriction_tracker.is_bounded():
+ return restriction
Review comment:
```suggestion
yield restriction
```
##########
File path: sdks/python/apache_beam/io/iobase.py
##########
@@ -1243,6 +1243,17 @@ def try_claim(self, position):
"""
raise NotImplementedError
+ def is_bounded(self):
+ """Identify whether the output produced by the current restriction is
+ bounded.
+
+ The value is important for the default behavior of truncate when the
+ pipeline starts to drain in streaming. If the current restriction is
+ bounded, it will be processed completely by default. If the restriction is
+ unbounded, it will be truncated into null and finish processing
immediately.
+ """
+ raise NotImplementedError
Review comment:
ThreadsafeRestrictionTracker should call the delegate so that users of
it get the same behavior as if they were calling the real restriction tracker
We might want to update the RestrictionTrackerView if we want to expose the
boundedness property during execution.
##########
File path: sdks/python/apache_beam/runners/common.py
##########
@@ -968,6 +973,9 @@ def process_with_sized_restriction(self, windowed_value):
(element, (restriction, estimator_state)), _ = windowed_value.value
watermark_estimator = (
self.do_fn_invoker.invoke_create_watermark_estimator(estimator_state))
+ if restriction is None:
Review comment:
Same issue as in java, we should allow None restrictions and instead
produce 0 outputs.
##########
File path: model/pipeline/src/main/proto/beam_runner_api.proto
##########
@@ -383,6 +383,12 @@ message StandardPTransforms {
//
// Input: KV(KV(element, restriction), size); output: DoFn's output.
PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS = 2 [(beam_urn) =
"beam:transform:sdf_process_sized_element_and_restrictions:v1"];
+
+ // Truncates the restriction of each element/restriction pair and returns
the
+ // finite restriction which is going to be processed when drain starts.
+ // Input: KV(KV(element, restriction), size);
+ // Output: KV(KV(element, restriction), size).
Review comment:
Add a link to the drain design doc that Reuven shared with the Beam
community.
```suggestion
// Truncates the restriction of each element/restriction pair and
returns the
// finite restriction which will be processed when a pipeline is drained
<link>.
//
// Input: KV(KV(element, restriction), size);
// Output: KV(KV(element, restriction), size).
```
##########
File path: sdks/python/apache_beam/transforms/core.py
##########
@@ -320,6 +320,22 @@ def split_and_size(self, element, restriction):
for part in self.split(element, restriction):
yield part, self.restriction_size(element, part)
+ def truncate(self, element, restriction):
Review comment:
Should we also provide truncate_and_size?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
