boyuanzz commented on a change in pull request #10897:
URL: https://github.com/apache/beam/pull/10897#discussion_r411722105
##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/io/Read.java
##########
@@ -375,4 +423,372 @@ public void checkDone() throws IllegalStateException {
}
}
}
+
+ /**
+ * A splittable {@link DoFn} which executes an {@link UnboundedSource}.
+ *
+ * <p>We model the element as the original source and the restriction as a
pair of the sub-source
+ * and its {@link CheckpointMark}. This allows us to split the sub-source
over and over as long as
+ * the checkpoint mark is {@code null} or the {@link NoopCheckpointMark}
since it does not
+ * maintain any state.
+ */
+ @UnboundedPerElement
+ static class UnboundedSourceAsSDFWrapperFn<OutputT, CheckpointT extends
CheckpointMark>
+ extends DoFn<UnboundedSource<OutputT, CheckpointT>,
ValueWithRecordId<OutputT>> {
+ private static final Logger LOG =
LoggerFactory.getLogger(UnboundedSourceAsSDFWrapperFn.class);
+ private static final int DEFAULT_DESIRED_NUM_SPLITS = 20;
+ private static final int DEFAULT_BUNDLE_FINALIZATION_LIMIT_MINS = 10;
+ private final Coder<CheckpointT> restrictionCoder;
+
+ private UnboundedSourceAsSDFWrapperFn(Coder<CheckpointT> restrictionCoder)
{
+ this.restrictionCoder = restrictionCoder;
+ }
+
+ @GetInitialRestriction
+ public KV<UnboundedSource<OutputT, CheckpointT>, CheckpointT>
initialRestriction(
+ @Element UnboundedSource<OutputT, CheckpointT> element) {
+ return KV.of(element, null);
+ }
+
+ @GetSize
+ public double getSize(
Review comment:
Do we assume that `getSize` here is always called for calculating
backlog in streaming? It is not correct for calculating primary size(though
primary size doesn't matter for now).
----------------------------------------------------------------
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:
[email protected]
