tkaymak commented on code in PR #39971: URL: https://github.com/apache/beam/pull/39971#discussion_r3932009683
########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java: ########## @@ -0,0 +1,277 @@ +/* + * 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.spark.structuredstreaming.io.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.spark.SparkEnv; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.connector.read.streaming.Offset; +import org.apache.spark.storage.BlockManager; +import org.apache.spark.storage.BlockManagerId; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import scala.collection.Iterator; + +/** + * Driver side {@link MicroBatchStream} over a Beam {@link UnboundedSource}. + * + * <p>Offsets are opaque epochs, {@link #latestOffset()} advances by one every trigger. The source + * is split once and the splits are pinned under the checkpoint location, every batch of every run + * plans the same partitions. {@link #commit} purges marks below the committed epoch on a background + * thread, Spark never asks for those again. + * + * <p>Each split prefers the executor rendezvous hashing assigns it, so an executor joining or + * leaving moves only the splits of that executor. Locality is a hint, the reader cache restores a + * split from its durable mark wherever it lands. + */ +public class BeamMicroBatchStream<T> implements MicroBatchStream { + + private static final Logger LOG = LoggerFactory.getLogger(BeamMicroBatchStream.class); + + private final BeamSourceSpec<T> spec; + private final String checkpointLocation; + private final BeamSourceCheckpoint checkpoint; + + private final ExecutorService purger = + Executors.newSingleThreadExecutor( + runnable -> { + Thread thread = new Thread(runnable, "beam-source-mark-purge"); + thread.setDaemon(true); + return thread; + }); + private final AtomicBoolean purgeInFlight = new AtomicBoolean(); + private final AtomicLong purgeRequested = new AtomicLong(); + + private long epoch; + private @Nullable List<UnboundedSource<T, ?>> splits; + + BeamMicroBatchStream(BeamSourceSpec<T> spec, String checkpointLocation) { + this.spec = spec; + this.checkpointLocation = checkpointLocation; + this.checkpoint = + new BeamSourceCheckpoint(checkpointLocation, spec.hadoopConf().value().value()); + } + + @Override + public Offset initialOffset() { + return BeamOffset.ZERO; + } + + @Override + public synchronized Offset latestOffset() { + return new BeamOffset(++epoch); + } + + @Override + public Offset deserializeOffset(String json) { + BeamOffset offset = BeamOffset.fromJson(json); + fastForwardEpoch(offset.epoch()); + return offset; + } + + @Override + public InputPartition[] planInputPartitions(Offset start, Offset end) { + long startEpoch = ((BeamOffset) start).epoch(); + long endEpoch = ((BeamOffset) end).epoch(); + fastForwardEpoch(endEpoch); + List<UnboundedSource<T, ?>> pinned = splits(); + long[] quotas = splitQuotas(spec.maxRecordsPerBatch(), pinned.size()); + List<String> executors = sortedExecutors(); Review Comment: Thank you, will have a look! -- 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]
