pkuzmickas commented on code in PR #39874: URL: https://github.com/apache/beam/pull/39874#discussion_r3904435896
########## 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: Great catch! Thanks. ########## 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: Makes sense, using the name you suggested and modified the description š > 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 We tested several thresholds against different production workflows, but did not find one that performed well enough across them to make a good general default unfortunately at this time :( ########## 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: These helpers are also used by [`LazyFlinkSourceSplitEnumerator`](https://github.com/pkuzmickas/beam/blob/244e38b4c8898715e3701e9356abda83b9c91057/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java#L159-L160) and [`SizeBasedFlinkSourceSplitEnumerator`](https://github.com/pkuzmickas/beam/blob/244e38b4c8898715e3701e9356abda83b9c91057/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java#L129-L132). Moving them into `FlinkSourceSplitEnumerator` would make the lazy and size based implementations depend on the static enumerator for shared source splitting logic. I kept them in a neutral helper for that reason, but Iām happy to use a different approach if you prefer. ########## 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: `FlinkSourceEnumeratorStateSerializer` handles checkpoint-format versioning and upgrades the previous map-based state. For the new version, the checkpoint data (the assignment mode and list of pending splits) still uses the existing Java object serialization, so `FlinkSourceEnumeratorState` needs to implement `Serializable`, as the previous state did. I avoided adding field-by-field encoding to the new serializer for simplicity. Please let me know if you had another approach in mind. -- 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]
