Abacn commented on code in PR #39874:
URL: https://github.com/apache/beam/pull/39874#discussion_r3897819445
##########
runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java:
##########
@@ -18,23 +18,437 @@
package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.stream.Collectors;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.runners.flink.translation.utils.SerdeUtils;
import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.TestBoundedCountingSource;
import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.TestCountingSource;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.io.FileBasedSource;
+import org.apache.beam.sdk.io.FileBasedSource.FileBasedReader;
+import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.Source;
+import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.values.KV;
+import org.apache.flink.api.connector.source.SplitEnumerator;
import
org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext;
+import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.junit.Test;
-/** Unit tests for {@link FlinkSourceSplitEnumerator}. */
+/** Unit tests for the Flink source split enumerators. */
public class FlinkSourceSplitEnumeratorTest {
+ private static final long MEBIBYTE = 1024L * 1024L;
+ private static final long AUTO_THRESHOLD_MB = 6144L;
+ private static final int SOURCE_PARALLELISM = 2;
+ private static final int REQUESTED_SPLITS = 4;
+
+ @Test
Review Comment:
Thanks for adding thorough test coverage! However, adding 600+ lines of test
code introduces maintenance overhead and duplication. We can cut ~300+ lines
while keeping 100% of the effective test coverage by simplifying:
- Consolidate threshold decision tests: testSmallBoundedSource...,
testLargeBoundedSource..., testUnknownEstimates...,
testDefaultConfiguration..., testEmptyBoundedSource..., and
testConfigurationCanForce... all test selectAssignmentMode. These 6 methods can
be collapsed into a single parameterized or table-driven test using
assertAssignmentMode.
- Deduplicate restore tests:
testStaticRestoreReturnsPendingSplitsToOriginalOwnersAtSameParallelism is
redundant with
testRestoreRepartitionsStaticSplitsForNewParallelismWithoutEstimating (rescale
already covers the same-parallelism case).
- Remove out-of-scope test & mock:
testStaticAssignmentRespectsFileInputSplitMaxSize and TestFileBasedSource test
fileInputSplitMaxSizeMB, which is an existing option not introduced by this PR.
- Deduplicate early reader registration tests:
testSignalsNoMoreSplitsToEarlyReaderWithoutAssignment and
testSizeBasedSelectionAssignsStaticSplitsToEarlyReaders test the same
registration race condition.
- Trim custom mock classes: Removing TestFileBasedSource and simplifying
TestEstimatedSizeBoundedSource (e.g. dropping failEstimation/createFailing) to
eliminate mock boilerplate.
##########
runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitUtils.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.io.FileBasedSource;
+import org.apache.beam.sdk.io.Source;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+
+/** Shared Beam source sizing and splitting helpers. */
+final class FlinkSourceSplitUtils {
Review Comment:
It's only used by FlinkSourceSplitEnumerator. Consider just put helper
methods there.
In general XxxUtils classes are discouraged for specific helper methods
(admittedly it exists everywhere in code base)
##########
runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitAssignmentMode.java:
##########
@@ -15,6 +15,10 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
Review Comment:
THis thin enum can be put as an internal class of FlinkSourceEnumeratorState
##########
runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorState.java:
##########
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/** Checkpoint state shared by the lazy and static source split assignment
strategies. */
+public final class FlinkSourceEnumeratorState<T> implements Serializable {
Review Comment:
Since there is a dedicate FlinkSourceEnumeratorStateSerializer, why still
need implements java.io.Serializable? If it's needed it suggests the Serializer
isn't effective in some serialization paths
##########
runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java:
##########
@@ -380,6 +380,17 @@ public Long create(PipelineOptions options) {
void setFileInputSplitMaxSizeMB(Long fileInputSplitMaxSizeMB);
+ @Description(
+ "Minimum estimated input size in MiB per source reader for lazy split
assignment of "
+ + "bounded sources. The default of 0 always uses lazy assignment. A
positive value "
+ + "selects static round-robin assignment for sources estimated below
the threshold "
+ + "and lazy assignment for sources at or above it. Any negative
value always uses "
+ + "static assignment.")
+ @Default.Long(0)
+ Long getLazySourceSplitAssignmentMinSizeMbPerReader();
Review Comment:
A few comments
- We can leave it opt-in for now, however ideally the default option should
be good enough for most common use cases, provide a reasonably performant
outcome. I understand now it's expected to set this to some positive values to
benefit from it.
- The naming is pretty long. In addition, I don't see "lazySplit" as a term
anywhere in documentation. Not sure how it ended up with this naming.
Consider "getSourceStaticSplitThresholdMb" and just note (1) positive sizes
below this threshold uses static split (2) it's per reader in the description
string (3) note in description this is for DataStream mode only.
##########
runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java:
##########
@@ -0,0 +1,208 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.flink.api.connector.source.SplitEnumerator;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Selects static or lazy assignment by estimated bounded-source size. */
+final class SizeBasedFlinkSourceSplitEnumerator<T>
+ implements SplitEnumerator<FlinkSourceSplit<T>,
FlinkSourceEnumeratorState<T>> {
+ private static final Logger LOG =
+ LoggerFactory.getLogger(SizeBasedFlinkSourceSplitEnumerator.class);
+
+ private final SplitEnumeratorContext<FlinkSourceSplit<T>> context;
+ private final BoundedSource<T> boundedSource;
+ private final PipelineOptions pipelineOptions;
+ private final int numSplits;
+ private final Map<Integer, Optional<String>> pendingSplitRequests;
+ private final List<ReturnedSplits<T>> returnedSplits;
+
+ private @Nullable SplitEnumerator<FlinkSourceSplit<T>,
FlinkSourceEnumeratorState<T>> delegate;
+
+ SizeBasedFlinkSourceSplitEnumerator(
+ SplitEnumeratorContext<FlinkSourceSplit<T>> context,
+ BoundedSource<T> boundedSource,
+ PipelineOptions pipelineOptions,
+ int numSplits) {
+ this.context = context;
+ this.boundedSource = boundedSource;
+ this.pipelineOptions = pipelineOptions;
+ this.numSplits = numSplits;
+ this.pendingSplitRequests = new LinkedHashMap<>();
+ this.returnedSplits = new ArrayList<>();
+ }
+
+ @Override
+ public void start() {
+ context.callAsync(
+ this::selectAndSplit,
+ (initialState, error) -> {
+ if (error != null) {
+ throw new RuntimeException("Failed to select a source split
assignment mode.", error);
+ }
+
+ SplitEnumerator<FlinkSourceSplit<T>, FlinkSourceEnumeratorState<T>>
selectedDelegate =
+ createDelegate(initialState);
+ delegate = selectedDelegate;
+ returnedSplits.forEach(
+ returned -> selectedDelegate.addSplitsBack(returned.splits,
returned.subtaskId));
+ returnedSplits.clear();
+ selectedDelegate.start();
+ pendingSplitRequests.forEach(
+ (subtaskId, hostname) ->
+ selectedDelegate.handleSplitRequest(subtaskId,
hostname.orElse(null)));
+ pendingSplitRequests.clear();
+ });
+ }
+
+ @Override
+ public void handleSplitRequest(int subtaskId, @Nullable String
requesterHostname) {
+ if (delegate == null) {
+ pendingSplitRequests.put(subtaskId,
Optional.ofNullable(requesterHostname));
+ } else {
+ delegate.handleSplitRequest(subtaskId, requesterHostname);
+ }
+ }
+
+ @Override
+ public void addSplitsBack(List<FlinkSourceSplit<T>> splits, int subtaskId) {
+ if (delegate == null) {
+ returnedSplits.add(new ReturnedSplits<>(new ArrayList<>(splits),
subtaskId));
+ } else {
+ delegate.addSplitsBack(splits, subtaskId);
+ }
+ }
+
+ @Override
+ public void addReader(int subtaskId) {
+ if (delegate != null) {
+ delegate.addReader(subtaskId);
+ }
+ }
+
+ @Override
+ public FlinkSourceEnumeratorState<T> snapshotState(long checkpointId) throws
Exception {
+ if (delegate == null) {
+ return new FlinkSourceEnumeratorState<>(
+ FlinkSourceSplitAssignmentMode.UNDECIDED, new ArrayList<>());
+ }
+ return delegate.snapshotState(checkpointId);
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (delegate != null) {
+ delegate.close();
+ }
+ }
+
+ private FlinkSourceEnumeratorState<T> selectAndSplit() throws Exception {
+ long estimatedSizeBytes =
+ FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource,
pipelineOptions);
+ FlinkSourceSplitAssignmentMode selectedMode =
selectAssignmentMode(estimatedSizeBytes);
+ ArrayList<FlinkSourceSplit<T>> splits =
+ FlinkSourceSplitUtils.splitBoundedSource(
+ boundedSource, pipelineOptions, numSplits, estimatedSizeBytes);
+ LOG.info(
+ "Split bounded source {} into {} splits using {} assignment",
+ boundedSource,
+ splits.size(),
+ selectedMode);
+ return new FlinkSourceEnumeratorState<>(selectedMode, splits);
+ }
+
+ private FlinkSourceSplitAssignmentMode selectAssignmentMode(long
estimatedSizeBytes) {
+ long thresholdMb =
+ pipelineOptions
+ .as(FlinkPipelineOptions.class)
+ .getLazySourceSplitAssignmentMinSizeMbPerReader();
+ if (thresholdMb <= 0) {
+ throw new IllegalArgumentException(
+ "Size-based source assignment requires a positive threshold, but
received "
+ + thresholdMb
+ + ".");
+ }
+ if (estimatedSizeBytes < 0 || estimatedSizeBytes == Long.MAX_VALUE) {
+ LOG.info(
+ "Estimated size of bounded source {} is unknown. Using lazy split
assignment.",
+ boundedSource);
+ return FlinkSourceSplitAssignmentMode.LAZY;
+ }
+
+ int sourceParallelism = context.currentParallelism();
+ if (sourceParallelism <= 0) {
+ throw new IllegalStateException(
+ "Source parallelism must be positive, but was " + sourceParallelism
+ ".");
+ }
+ long estimatedBytesPerReader = estimatedSizeBytes / sourceParallelism;
+ long thresholdBytes = FlinkSourceSplitUtils.mebibytesToBytes(thresholdMb);
+ FlinkSourceSplitAssignmentMode selectedMode =
Review Comment:
Reader may return estimatedBytes as 0 if it does not support estimation.
This is common in both Beam builtin and user IO Connectors. To make the option
truely opt-in I understand the condition should read `(thresholdBytes <=0 ||
estimatedBytesPerReader >= thresholdBytes)`
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
