You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by GitBox <gi...@apache.org> on 2020/05/19 02:28:09 UTC
[GitHub] [beam] boyuanzz opened a new pull request #11749: [WIP, DO NOT REVIEW PLEASE] Implement ReadFromKafkaViaSDF
boyuanzz opened a new pull request #11749:
URL: https://github.com/apache/beam/pull/11749
**Please** add a meaningful description for your change here
------------------------
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
- [ ] [**Choose reviewer(s)**](https://beam.apache.org/contribute/#make-your-change) and mention them in a comment (`R: @username`).
- [ ] Format the pull request title like `[BEAM-XXX] Fixes bug in ApproximateQuantiles`, where you replace `BEAM-XXX` with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
- [ ] Update `CHANGES.md` with noteworthy changes.
- [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://beam.apache.org/contribute/#make-reviewers-job-easier).
Post-Commit Tests Status (on master branch)
------------------------------------------------------------------------------------------------
Lang | SDK | Apex | Dataflow | Flink | Gearpump | Samza | Spark
--- | --- | --- | --- | --- | --- | --- | ---
Go | [](https://builds.apache.org/job/beam_PostCommit_Go/lastCompletedBuild/) | --- | --- | [](https://builds.apache.org/job/beam_PostCommit_Go_VR_Flink/lastCompletedBuild/) | --- | --- | [](https://builds.apache.org/job/beam_PostCommit_Go_VR_Spark/lastCompletedBuild/)
Java | [](https://builds.apache.org/job/beam_PostCommit_Java/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Apex/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink_Java11/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Batch/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Gearpump/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Samza/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Spark/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming/lastCompletedBuild/)
Python | [](https://builds.apache.org/job/beam_PostCommit_Python2/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Python35/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Python36/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Python37/lastCompletedBuild/) | --- | [](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow_V2/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Py_ValCont/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PreCommit_Python2_PVR_Flink_Cron/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PostCommit_Python35_VR_Flink/lastCompletedBuild/) | --- | --- | [](https://builds.apache.org/job/beam_PostCommit_Python_VR_Spark/lastCompletedBuild/)
XLang | --- | --- | --- | [](https://builds.apache.org/job/beam_PostCommit_XVR_Flink/lastCompletedBuild/) | --- | --- | [](https://builds.apache.org/job/beam_PostCommit_XVR_Spark/lastCompletedBuild/)
Pre-Commit Tests Status (on master branch)
------------------------------------------------------------------------------------------------
--- |Java | Python | Go | Website
--- | --- | --- | --- | ---
Non-portable | [](https://builds.apache.org/job/beam_PreCommit_Java_Cron/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PreCommit_Python_Cron/lastCompletedBuild/)<br>[](https://builds.apache.org/job/beam_PreCommit_PythonLint_Cron/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PreCommit_Go_Cron/lastCompletedBuild/) | [](https://builds.apache.org/job/beam_PreCommit_Website_Cron/lastCompletedBuild/)
Portable | --- | [](https://builds.apache.org/job/beam_PreCommit_Portable_Python_Cron/lastCompletedBuild/) | --- | ---
See [.test-infra/jenkins/README](https://github.com/apache/beam/blob/master/.test-infra/jenkins/README.md) for trigger phrase, status and link of all Jenkins jobs.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637676917
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r436875673
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
I'm not worried about the size of the `source descriptor` for Kafka but more about the pattern in which people would build and use these objects and about the backwards/forwards compatibility in a pipeline update scenario. I think it would be best if we used a dedicated object not extending the `PTransform` because:
1) it simplifies the migration scenario to a 100% SDF world. `Read` becomes a wrapper around `Create(source descriptor) -> ReadAll` when using SDF otherwise we use the existing `UnboundedSource` implementation and it pulls parameters off of this `source descriptor`
2) we don't couple source descriptor with classes/evolution of pipeline construction APIs
3) seems awkward from a user perspective to have a PCollection<PTransform> -> PTransform -> PCollection<KafkaRecord>
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
Review comment:
I don't think we should expose ReadFromKafkaViaSDF and instead rely on the experiments to toggle which of the two backing implementations the user gets for the existing `KafkaIO.Read` transform. This transform would be like
```
expand() {
if (beam_fn_api) {
Create(source descriptor) -> ReadAll()
} else {
Read(UnboundedKafkaSource)
}
}
```
We could add a separate ReadAll implementation which takes a PCollection of `source descriptors` that exercises the SDF.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637820927
> @boyuanzz Yes, afaik, it's used only for [performance testing](http://metrics.beam.apache.org/d/bnlHKP3Wz/java-io-it-tests-dataflow?panelId=21&fullscreen&orgId=1&from=1588539922272&to=1591131922272) on Dataflow
Thanks! I'm testing it with Dataflow internally. Enabling related java tests required `runner_v2`. I think the best way now to add tests for this transform would be using python x-lang test.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637679153
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] piotr-szuberski commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
piotr-szuberski commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r467756812
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -451,7 +574,9 @@
// Set required defaults
setTopicPartitions(Collections.emptyList());
- setConsumerFactoryFn(Read.KAFKA_CONSUMER_FACTORY_FN);
+ setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN);
+ setMaxNumRecords(Long.MAX_VALUE);
Review comment:
MaxNumRecords is set 2 lines below, this line should be removed.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445225947
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +1082,91 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ }
+ ReadAll<K, V> readTransform =
+ ReadAll.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withManualWatermarkEstimator()
+ .withTimestampPolicyFactory(getTimestampPolicyFactory());
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+ PCollection<KafkaSourceDescription> output =
+ input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(ParDo.of(new GenerateKafkaSourceDescription(this)));
+ try {
+ output.setCoder(KafkaSourceDescription.getCoder(input.getPipeline().getSchemaRegistry()));
Review comment:
It works with `setSchema` but I want to make it explicitly because it's possible that an user writes a DoFn which produces `KafkaSourceDescription`.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445087927
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
Review comment:
```suggestion
* Represents a Kafka source description.
*
* <p>Note that this object should be encoded/decoded with its corresponding {@link #getCoder schema coder}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
Review comment:
```suggestion
* <pre>{@code
* pipeline
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
Review comment:
```suggestion
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
+ @SchemaFieldName("topic")
+ abstract String getTopic();
+
+ @SchemaFieldName("partition")
+ abstract Integer getPartition();
+
+ @SchemaFieldName("start_read_offset")
+ @Nullable
+ abstract Long getStartReadOffset();
+
+ @SchemaFieldName("start_read_time")
+ @Nullable
+ abstract Instant getStartReadTime();
+
+ @SchemaFieldName("bootstrapServers")
Review comment:
Did you mean to make this one snake_case as well?
```suggestion
@SchemaFieldName("bootstrap_servers")
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
Review comment:
```suggestion
* }</pre>
*
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
Review comment:
```suggestion
* query Kafka topics from a BigQuery table and read these topics via {@link ReadAll}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
Review comment:
```suggestion
* Note that the {@code bootstrapServers} can also be populated from the {@link KafkaSourceDescriptor}:
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
Review comment:
```suggestion
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
Review comment:
Should we be defining endReadOffset and endReadTime which are optional as well?
```suggestion
* {@link The initial range for
* a {@link KafkaSourceDescription} is defined by {@code [startOffset, Long.MAX_VALUE)} where {@code startOffset} is defined as:
*
* <ul>
* <li>the {@code startReadOffset} if {@link KafkaSourceDescription#getStartReadOffset} is set.
* <li>the first offset with a greater or equivalent timestamp if {@link KafkaSourceDescription#getStartReadTimestamp} is set.
* <li>the {@code last committed offset + 1} for the {@link
* Consumer#position(TopicPartition) topic partition}.
* </ul>
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
Review comment:
```suggestion
* <h4>Splitting</h4>
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
Review comment:
nit: appropriate appropriate -> appropriate on https://github.com/apache/beam/blob/f98104a22b69972744a13378e17af5f2361fbb3e/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java#L156
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
Review comment:
```suggestion
* <li>{@link ReadAll#isCommitOffsetEnabled()} has the same meaning as {@link
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
Review comment:
```suggestion
* .apply(Create.of(
* KafkaSourceDescription.of(
* new TopicPartition("topic", 1),
* null,
* null,
* ImmutableList.of("broker_1:9092", "broker_2:9092"))
* .apply(KafkaIO.readAll()
* .withKeyDeserializer(LongDeserializer.class).
* .withValueDeserializer(StringDeserializer.class));
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
Review comment:
Do you think the name `KafkaSourceDescriptor` would be more appropriate?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
Review comment:
```suggestion
* <h4>Initial Restriction</h4>
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
Review comment:
```suggestion
* that if the {@code isolation.level} is set to "read_committed" or {@link
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,339 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadAll;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+@UnboundedPerElement
+class ReadFromKafkaDoFn<K, V> extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadAll transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaDoFn.class);
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Valid between bundle start and bundle finish.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, AverageRecordSize> avgRecordSize;
Review comment:
This is going to grow indefinitely, please use an LRU cache with a limited number of elements (e.g. 1000).
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn#restrictionTracker(KafkaSourceDescription, OffsetRange)} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link ReadAll#getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link ReadAll#getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link ReadAll#withProcessingTime()} as {@code extractTimestampFn} and
+ * {@link ReadAll#withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
Review comment:
Move to ReadFromKafkaViaDoFn
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1261,341 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadAll<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadAll.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract ReadAll.Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract ReadAll.Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract ReadAll.Builder<K, V> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract ReadAll.Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract ReadAll.Builder<K, V> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadAll.Builder<K, V> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadAll.Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract ReadAll.Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract ReadAll.Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract ReadAll.Builder<K, V> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn);
+
+ abstract ReadAll.Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadAll.Builder<K, V> setTimestampPolicyFactory(TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadAll<K, V> build();
+ }
- for (String key : updates.keySet()) {
+ public static <K, V> ReadAll<K, V> read() {
+ return new AutoValue_KafkaIO_ReadAll.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the
+ // element
+ // will override the bootstrapServers from the config.
+ public ReadAll<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadAll<K, V> withKeyDeserializerProvider(DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadAll<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadAll<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadAll<K, V> withValueDeserializer(Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadAll<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadAll<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadAll<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadAll<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadAll<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadAll<K, V> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadAll<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadAll<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadAll<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadAll<K, V> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new WallTime(state);
+ });
+ }
+
+ public ReadAll<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadAll<K, V> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
+ // default
+ // value.
+ public ReadAll<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadAll<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadAll<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadAll<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadAllFromRow forExternalBuild() {
+ return new ReadAllFromRow(this);
+ }
+
+ // This transform is used in cross-language case. The input Row should be encoded with an
+ // equivalent schema as KafkaSourceDescription.
+ private static class ReadAllFromRow<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadAll<K, V> readViaSDF;
+
+ ReadAllFromRow(ReadAll read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(Convert.fromRows(KafkaSourceDescription.class))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
Review comment:
We should move this to the external transform builder and remove the `forExternalBuild` method.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1261,341 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadAll<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadAll.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract ReadAll.Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract ReadAll.Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract ReadAll.Builder<K, V> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract ReadAll.Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract ReadAll.Builder<K, V> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadAll.Builder<K, V> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadAll.Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract ReadAll.Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract ReadAll.Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract ReadAll.Builder<K, V> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn);
+
+ abstract ReadAll.Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadAll.Builder<K, V> setTimestampPolicyFactory(TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadAll<K, V> build();
+ }
- for (String key : updates.keySet()) {
+ public static <K, V> ReadAll<K, V> read() {
+ return new AutoValue_KafkaIO_ReadAll.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the
+ // element
+ // will override the bootstrapServers from the config.
+ public ReadAll<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadAll<K, V> withKeyDeserializerProvider(DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadAll<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadAll<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadAll<K, V> withValueDeserializer(Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadAll<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadAll<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadAll<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadAll<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadAll<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadAll<K, V> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadAll<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadAll<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadAll<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadAll<K, V> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new WallTime(state);
+ });
+ }
+
+ public ReadAll<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadAll<K, V> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
+ // default
+ // value.
Review comment:
Should this be javadoc?
```suggestion
// If a transactional producer is used and it's desired to only read records from committed
// transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
// default value.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -729,12 +902,15 @@ public void setValueDeserializer(String valueDeserializer) {
/**
* Provide custom {@link TimestampPolicyFactory} to set event times and watermark for each
- * partition. {@link TimestampPolicyFactory#createTimestampPolicy(TopicPartition, Optional)} is
- * invoked for each partition when the reader starts.
+ * partition when beam_fn_api is disabled. {@link
+ * TimestampPolicyFactory#createTimestampPolicy(TopicPartition, Optional)} is invoked for each
+ * partition when the reader starts.
*
* @see #withLogAppendTime()
* @see #withCreateTime(Duration)
* @see #withProcessingTime()
+ * <p>For the pipeline with beam_fn_api is enabled, you should use {@link
Review comment:
I don't think we need this suggestion anymore.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
Review comment:
```suggestion
* <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} is used to compute the {@code output timestamp} for a given {@link KafkaRecord}. There are three built-in types: {@link
* ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
* ReadAll#withLogAppendTime()}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -729,12 +902,15 @@ public void setValueDeserializer(String valueDeserializer) {
/**
* Provide custom {@link TimestampPolicyFactory} to set event times and watermark for each
- * partition. {@link TimestampPolicyFactory#createTimestampPolicy(TopicPartition, Optional)} is
- * invoked for each partition when the reader starts.
+ * partition when beam_fn_api is disabled. {@link
Review comment:
I think we can keep the original comment here since we support the timestamp policy fn now.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
Review comment:
```suggestion
* <p>For example, to create a {@link ReadAll} with this additional configuration:
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
Review comment:
This is an implementation detail, I'm not sure we want to share it as part of the Javadoc that we want users to read when they look at KafkaIO. It would make sense to have this on the ReadFromKafkaViaDoFn class though.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
Review comment:
```suggestion
* <p>TODO(BEAM-YYY): Add support for initial splitting.
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445236767
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
Review comment:
It seems like `Descriptor` makes more sense.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444354311
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
Review comment:
I believe the default is 30 seconds for Dataflow but I don't thing we should tune this to be Dataflow specific.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445707474
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
Review comment:
`endReadTime` is the same time domain as `startReadTime`.
Having an end would be for batch but could also be useful in streaming pipelines.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439547427
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
I'd suggest to stick with 1 (at least for now). 2 and 3 seems to me could be an error prone since I expect it will require much more work in case if the new servers will be related to different cluster, for example. Wdyt?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442966047
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
Review comment:
Do we really need to poll here, can we check if there is anything left and if there is nothing return a process continuation with a resume delay that represents the poll timeout?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r466683752
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
Review comment:
I would prefer that we leave the documentation here and link it from KafkaIO.ReadSourceDescriptors.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434746269
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java
##########
@@ -0,0 +1,146 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+
+/**
+ * Common utility functions and default configurations for {@link KafkaIO.Read} and {@link
+ * ReadFromKafkaViaSDF}.
+ */
+final class KafkaIOUtils {
+ // A set of config defaults.
Review comment:
Yes, the util is just moved from KafkaIO.java and for code-reuse purpose only.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444646093
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
Review comment:
I don't think it's easy to come up with a default size here. There is no avgSize and avgGap only when initial sizing, which is not really important for streaming.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r466666972
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1198,352 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadSourceDescriptors<K, V>
+ extends PTransform<PCollection<KafkaSourceDescriptor>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadSourceDescriptors.class);
+
+ abstract Map<String, Object> getConsumerConfig();
Review comment:
Most of the documentations are at KafkaIO level.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] piotr-szuberski commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
piotr-szuberski commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-671207252
@boyuanzz This change has broken python postcommits that run kafka cross-language it tests.
Could you take a look on that? I tried to fix it but it seems it would take me a lot of time since I'm not familiar with those SDF changes.
A stacktrace fragment:
```
Caused by: java.util.concurrent.ExecutionException: java.lang.RuntimeException: Error received from SDK harness for instruction 2: java.util.concurrent.ExecutionException: java.lang.RuntimeException: Could not find a way to create AutoValue class class org.apache.beam.sdk.io.kafka.KafkaSourceDescriptor
at java.util.concurrent.CompletableFuture.reportGet(CompletableFuture.java:357)
at java.util.concurrent.CompletableFuture.get(CompletableFuture.java:1908)
at org.apache.beam.sdk.fn.data.CompletableFutureInboundDataClient.awaitCompletion(CompletableFutureInboundDataClient.java:48)
at org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver.awaitCompletion(BeamFnDataInboundObserver.java:91)
at org.apache.beam.fn.harness.BeamFnDataReadRunner.blockTillReadFinishes(BeamFnDataReadRunner.java:342)
at org.apache.beam.fn.harness.data.PTransformFunctionRegistry.lambda$register$0(PTransformFunctionRegistry.java:108)
at org.apache.beam.fn.harness.control.ProcessBundleHandler.processBundle(ProcessBundleHandler.java:302)
at org.apache.beam.fn.harness.control.BeamFnControlClient.delegateOnInstructionRequestType(BeamFnControlClient.java:173)
at org.apache.beam.fn.harness.control.BeamFnControlClient.lambda$processInstructionRequests$0(BeamFnControlClient.java:157)
at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
at java.lang.Thread.run(Thread.java:748)
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445016902
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1261,341 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadAll<K, V>
Review comment:
Thanks for starting the discussion. If taking x-lang usage into consideration, `Read` is not a good choice as input for these DoFn.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439565391
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
In this sense, I agree with option #2 as well.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] iemejia commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
iemejia commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444912597
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1261,341 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadAll<K, V>
Review comment:
I really would like to avoid calling this ReadAll because it is not consistent with the ongoing work where we call ReadAll transforms the ones from `PCollection<Read>`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445242837
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
Review comment:
Are endReadOffset and endReadTime for batch reading case? Is `endReadTime` for process time or for event time?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442574397
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
Review comment:
I misunderstood how nested and inner class gets serialized. I think I should mark `ReadFromKafkaDoFn` as static to avoid serializing outer class. Thanks for pointing it out!
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-638348082
Also, please, take a look on SpotBugs issues.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] TheNeuralBit commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
TheNeuralBit commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r446348275
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +1082,91 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ }
+ ReadAll<K, V> readTransform =
+ ReadAll.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withManualWatermarkEstimator()
+ .withTimestampPolicyFactory(getTimestampPolicyFactory());
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+ PCollection<KafkaSourceDescription> output =
+ input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(ParDo.of(new GenerateKafkaSourceDescription(this)));
+ try {
+ output.setCoder(KafkaSourceDescription.getCoder(input.getPipeline().getSchemaRegistry()));
Review comment:
I don't think you should need `setSchema` either, we should always use the SchemaCoder for a PCollection<T> where T has a Schema registered (unless the PCollection has a specific coder set, or T also has a default coder set):
https://github.com/apache/beam/blob/b83c06d47bd5e2e2905599297981353af234b034/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollection.java#L154-L166
I'm not sure I follow why a DoFn that produces `KafkaSourceDescription` requires any special logic
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] iemejia commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
iemejia commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434666700
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
This seems strangely close to something we lived in the SDF version of HBaseIO. In the first version we did an artificial object called `HBaseQuery` that contained the minimum information we needed to be able to query the Data store in a SDF way, but then other requirements came in and we started to add extra parameters to end up with something that was almost close to the exact 'complete' specification of the Read class so we decided to switch to use a `PCollection<Read>` as input otherwise we will be duplicating code, so we ended up with https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java#L353
Here you can have `PCollection<Read>` as an input and get rid of `KafkaSourceDescription` if you move the missing parameters into normal `Read` and this will have a more consistent user experience for final users. Notice that this `ReadAll` like pattern is also now used in [SolrIO](https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/solr/src/main/java/org/apache/beam/sdk/io/solr/SolrIO.java#L501) and there is an ongoing PR to introduce it for [CassandraIO](https://github.com/apache/beam/pull/10546) so maybe it is a good idea we follow it for consistency.
Notice that in the SolrIO case the change looks even closer to this one because we ended up introducing `ReplicaInfo` (the spiritual equivalent of `TopicPartition`) into normal Read and we guarantee in expansion that this field gets filled if the users don't do it, but if they do well we asume they know what they are doing and we go with it.
Another advantage of having the full specification is that you will be able to read not only from multiple topics but also from different clusters because of the power of having the full `Read` specification,
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r436922850
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
Review comment:
The current does do the the same thing as you mentioned, except the naming(`read()` vs `ReadAll()`).
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
Review comment:
The current implementation does do the the same thing as you mentioned, except the naming(`read()` vs `ReadAll()`).
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r467317077
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
Review comment:
Done, thanks!
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442579140
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
Review comment:
> Why would we close this when the object goes out of scope?
I need to close the connection when there is no more usage. There are 2 option:
1. The connection is only created when it's needed, like when querying the backlog.
2. The connection is created when the tracker is created and keep until the track is out of scope.
I go with 2 to avoid creating the same connection multiple time.
> Would it make sense to have support for a KafkaConsumer pool that expires automatically after X amount of time of not being used?
For Kafka, it's important to close consumer effectively to benefit from Kafka dynamic load balancing mechanism. But we are using manual assignment now. So as long as the connection can be created cheaply, we may not need the pool.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445707474
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +213,154 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadAll} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescription} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadAll} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadAll}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadAll#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadAll#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadAll#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadAll#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadAll#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadAll#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadAll#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadAll#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadAll} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadAll}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadAll#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadAll#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadAll#withExtractOutputTimestampFn(SerializableFunction)} asks for a function which
+ * takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadAll#withProcessingTime()}, {@link ReadAll#withCreateTime()} and {@link
+ * ReadAll#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadAll} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
Review comment:
`endReadTime` is the same time domain as `startReadTime`.
Having an end would be primarily for batch but could also be useful in streaming pipelines.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev removed a comment on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev removed a comment on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637676591
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434890835
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -817,6 +847,24 @@ public void setValueDeserializer(String valueDeserializer) {
return toBuilder().setCommitOffsetsInFinalizeEnabled(true).build();
}
+ /**
+ * The {@link Read} transform will be expanded with {@link ReadFromKafkaViaSDF} transform. While
+ * working with {@link #useSDFTransformInRead()} and {@link
+ * #withTimestampPolicyFactory(TimestampPolicyFactory)} together, only {@link
+ * TimestampPolicyFactory#withCreateTime(Duration)}, {@link
+ * TimestampPolicyFactory#withLogAppendTime()} and {@link
+ * TimestampPolicyFactory#withProcessingTime()} will be populated correctly. For other custom
+ * {@link TimestampPolicy}, the transform will use {@link
+ * TimestampPolicyFactory#withProcessingTime()} by default. It's recommended to use {@link
+ * ReadFromKafkaViaSDF} directly in that case.
+ *
+ * <p>Note that the expansion only happens when tbe pipeline has "beam_fn_api" experiment and
+ * meanwhile "beam_fn_api_use_deprecated_read" is not set.
+ */
+ public Read<K, V> useSDFTransformInRead() {
Review comment:
Discussed with Luke offline. We think it would be better to make `KafkaIO.Read()` expand with SDF transform bu default when `beam_fn_api` is enabled( before introducing this SDF transform, we expand the `Read` with SDFUnboundedWrapper with beam_fn_api).
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442572993
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
Review comment:
I was under impression that `IO' means it will be the root transform. If that's not the case, `readAll` sounds good.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz merged pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz merged pull request #11749:
URL: https://github.com/apache/beam/pull/11749
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-668251029
Based on the [readAll discussion](https://lists.apache.org/thread.html/rcb7538d5a2955d12e6d5f37a24830b929e5d65155d84f6d999801432%40%3Cdev.beam.apache.org%3E), I decided to go with my current approach by renaming `KafkaIO.readAll` to `KafkaIO.readSourceDescriptors`.
@lukecwik and @aromanenko-dev, would you like to do another pass to see whether you have other concerns?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434739659
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
Hi Ismaël, thanks for your review and comments!
Yes I thought about using `KafkaIO.Read` as element and there are some pushbacks from my side:
1. `KafkaIO.Read` is kind of heavy. For me, `KafkaIO.Read` is more like a configuration + element. For me, the element should be something you may only know about it during the pipeline execution time. So I want to isolate element-like into `KafkaSourceDescription`.
2. For the case you mentioned that reading from different clusters, I thought about it and not sure whether it would be a common case for reading from Kafka. So I sent out an email titled with `[Discuss] Build Kafka read transform on top of SplittableDoFn` to our dev mailing list to figure what could be an element in common. So far I didn't hear back from the community for the need of reading from different clusters.
3. We also consider x-lang usage for `ReadFromKafkaViaSDF`, which requires we can encode and decode the element over the wire. So I want to make the element as low weight as possible.
For the concern of increasing needs of element, we want to have `KafkaSourceDescription` easy to be extended, as well as the coder.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r435359001
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +955,110 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ if (!isUseSDFTransform()
+ || !ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
Review comment:
In this case, please, add a comment about that.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444648197
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(expectedOffset - rawRecord.offset());
+ avgRecordSize
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(recordSize);
+ expectedOffset = rawRecord.offset() + 1;
Review comment:
Yes, I think that should be the case. For example, at this moment we get a record with offset 1. Then we expect that the next offset with gap should be 2.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442570364
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ Instant outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize.update(recordSize);
+ avgOffsetGap.update(expectedOffset - rawRecord.offset());
+ expectedOffset = rawRecord.offset() + 1;
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
+ throw anyException;
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize = new KafkaIOUtils.MovingAvg();
Review comment:
> Will avgRecordSize / avgOffsetGap be consistent across multiple topics / partitions.
I don't think there is a guarantee . I even don't think there is a guarantee per topic. The `avgRecordSize` and `avgOffsetGap` are used to calculate size in `GetSize`, which should be an estimated one.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r436179356
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1055,29 +1144,6 @@ public void populateDisplayData(DisplayData.Builder builder) {
private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
- /**
Review comment:
This common part is moved to the KafkaIOUtil.java
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r436187245
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
Review comment:
Yes, it's exposed to the user since the user may want to use this transform separate from `KafkaIO.Read()`, for example, the user can get TopicPartition from BigQuery during runtime.
I'll rename to `read`. Thanks for the naming suggestion!
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r435359001
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +955,110 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ if (!isUseSDFTransform()
+ || !ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
Review comment:
In this case, please, a comment about that.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439531254
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
The question here is, do we want to only have bootstrap servers in element during runtime, or we want to allow the user to set the bootstrap server when constructing the pipeline, and mean while also allow the user to provide new bootstrap server during runtime. For example, if the customer gives the static list of [bootstrap1] during construction time, and also give new bootstrap2 in the runtime, what should be the list we use to connect:
1. just use [bootstrap1]
2. or just use [bootstrap2]
3. or [bootstrap1, bootstrap2]
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434717019
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -817,6 +847,24 @@ public void setValueDeserializer(String valueDeserializer) {
return toBuilder().setCommitOffsetsInFinalizeEnabled(true).build();
}
+ /**
+ * The {@link Read} transform will be expanded with {@link ReadFromKafkaViaSDF} transform. While
+ * working with {@link #useSDFTransformInRead()} and {@link
+ * #withTimestampPolicyFactory(TimestampPolicyFactory)} together, only {@link
+ * TimestampPolicyFactory#withCreateTime(Duration)}, {@link
+ * TimestampPolicyFactory#withLogAppendTime()} and {@link
+ * TimestampPolicyFactory#withProcessingTime()} will be populated correctly. For other custom
+ * {@link TimestampPolicy}, the transform will use {@link
+ * TimestampPolicyFactory#withProcessingTime()} by default. It's recommended to use {@link
+ * ReadFromKafkaViaSDF} directly in that case.
+ *
+ * <p>Note that the expansion only happens when tbe pipeline has "beam_fn_api" experiment and
+ * meanwhile "beam_fn_api_use_deprecated_read" is not set.
+ */
+ public Read<K, V> useSDFTransformInRead() {
Review comment:
Maybe call it just useSDF()? Because it's already known that it's a PTransform used in Read
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java
##########
@@ -0,0 +1,146 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+
+/**
+ * Common utility functions and default configurations for {@link KafkaIO.Read} and {@link
+ * ReadFromKafkaViaSDF}.
+ */
+final class KafkaIOUtils {
+ // A set of config defaults.
Review comment:
I expect that all these constants and methods below were moved here without any changes and just for the sake of code refactoring. If not, please, add some comments on these.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +955,110 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ if (!isUseSDFTransform()
+ || !ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
Review comment:
It looks that we depend on some specific pipeline business logic here. I'd prefer to avoid this if possible.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +955,110 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ if (!isUseSDFTransform()
+ || !ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
Review comment:
The same point as above.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
Review comment:
Do we need to expose it to user? Could it be just `read()` to be consistent with `KafkaIO.Read`?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
I would shorter the name of this class to `ReadWithSDF` since this is clear that it's used to read from Kafka.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
+ return new AutoValue_ReadFromKafkaViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .build();
+ }
+
Review comment:
Do all these configuration methods repeat `KafkaIO.Read` methods? Can we avoid a code duplication with new `ReadFromKafkaViaSDF` transform?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442492359
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java
##########
@@ -0,0 +1,146 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+
+/**
+ * Common utility functions and default configurations for {@link KafkaIO.Read} and {@link
+ * ReadViaSDF}.
+ */
+final class KafkaIOUtils {
+ // A set of config defaults.
+ static final Map<String, Object> DEFAULT_CONSUMER_PROPERTIES =
+ ImmutableMap.of(
+ ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
+ ByteArrayDeserializer.class.getName(),
+ ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
+ ByteArrayDeserializer.class.getName(),
+
+ // Use large receive buffer. Once KAFKA-3135 is fixed, this _may_ not be required.
+ // with default value of of 32K, It takes multiple seconds between successful polls.
+ // All the consumer work is done inside poll(), with smaller send buffer size, it
+ // takes many polls before a 1MB chunk from the server is fully read. In my testing
+ // about half of the time select() inside kafka consumer waited for 20-30ms, though
+ // the server had lots of data in tcp send buffers on its side. Compared to default,
+ // this setting increased throughput by many fold (3-4x).
+ ConsumerConfig.RECEIVE_BUFFER_CONFIG,
+ 512 * 1024,
+
+ // default to latest offset when we are not resuming.
+ ConsumerConfig.AUTO_OFFSET_RESET_CONFIG,
+ "latest",
+ // disable auto commit of offsets. we don't require group_id. could be enabled by user.
+ ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG,
+ false);
+
+ // A set of properties that are not required or don't make sense for our consumer.
+ static final Map<String, String> IGNORED_CONSUMER_PROPERTIES =
Review comment:
nit:
IGNORED_CONSUMER_PROPERTIES -> DISALLOWED_CONSUMER_PROPERTIES
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +946,123 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
+
+ return input.getPipeline().apply(transform);
+ } else {
+ // If extractOutputTimestampFn is not set, use processing time by default.
+ SerializableFunction<KafkaRecord<K, V>, Instant> timestampFn;
+ if (getExtractOutputTimestampFn() != null) {
+ timestampFn = getExtractOutputTimestampFn();
+ } else {
+ timestampFn = ReadViaSDF.ExtractOutputTimestampFns.useProcessingTime();
+ }
+ ReadViaSDF<K, V> readTransform =
+ ReadViaSDF.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withExtractOutputTimestampFn(timestampFn);
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(
+ ParDo.of(
+ new GenerateKafkaSourceDescription(
+ readTransform.getKafkaSourceDescriptionSchema())))
+ .setCoder(RowCoder.of(readTransform.getKafkaSourceDescriptionSchema()))
+ .apply(readTransform)
+ .setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
+ }
+ }
+
+ /**
+ * A DoFn which generates {@link Row} with {@link KafkaSourceDescriptionSchemas#getSchema()}
+ * based on the configuration of {@link Read}.
+ */
+ @VisibleForTesting
+ class GenerateKafkaSourceDescription extends DoFn<byte[], Row> {
+ GenerateKafkaSourceDescription(Schema schema) {
+ this.kafkaSourceDescriptionSchema = schema;
+ }
+
+ private final Schema kafkaSourceDescriptionSchema;
+
+ private final Map<String, Object> consumerConfig = Read.this.getConsumerConfig();
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = Read.this.getConsumerFactoryFn();
- if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
- transform =
- unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ private final List<String> topics = Read.this.getTopics();
+
+ private final List<TopicPartition> topicPartitions = Read.this.getTopicPartitions();
+
+ private final Instant startReadTime = Read.this.getStartReadTime();
+
+ @VisibleForTesting
+ Map<String, Object> getConsumerConfig() {
+ return consumerConfig;
}
- return input.getPipeline().apply(transform);
+ @VisibleForTesting
+ List<String> getTopics() {
+ return topics;
+ }
+
+ @ProcessElement
+ public void processElement(OutputReceiver<Row> receiver) {
+ List<TopicPartition> partitions = new ArrayList<>(topicPartitions);
+ if (partitions.isEmpty()) {
+ try (Consumer<?, ?> consumer = consumerFactoryFn.apply(consumerConfig)) {
+ for (String topic : topics) {
+ for (PartitionInfo p : consumer.partitionsFor(topic)) {
+ partitions.add(new TopicPartition(p.topic(), p.partition()));
+ }
+ }
+ }
+ }
+ partitions.stream()
Review comment:
nit: this is much harder to read then a for loop. Also the performance of stream() is poor relative to the for loop as well.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +946,123 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
+
+ return input.getPipeline().apply(transform);
+ } else {
+ // If extractOutputTimestampFn is not set, use processing time by default.
+ SerializableFunction<KafkaRecord<K, V>, Instant> timestampFn;
+ if (getExtractOutputTimestampFn() != null) {
+ timestampFn = getExtractOutputTimestampFn();
+ } else {
+ timestampFn = ReadViaSDF.ExtractOutputTimestampFns.useProcessingTime();
+ }
+ ReadViaSDF<K, V> readTransform =
+ ReadViaSDF.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withExtractOutputTimestampFn(timestampFn);
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(
+ ParDo.of(
+ new GenerateKafkaSourceDescription(
+ readTransform.getKafkaSourceDescriptionSchema())))
+ .setCoder(RowCoder.of(readTransform.getKafkaSourceDescriptionSchema()))
+ .apply(readTransform)
+ .setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
+ }
+ }
+
+ /**
+ * A DoFn which generates {@link Row} with {@link KafkaSourceDescriptionSchemas#getSchema()}
+ * based on the configuration of {@link Read}.
+ */
+ @VisibleForTesting
+ class GenerateKafkaSourceDescription extends DoFn<byte[], Row> {
+ GenerateKafkaSourceDescription(Schema schema) {
+ this.kafkaSourceDescriptionSchema = schema;
+ }
+
+ private final Schema kafkaSourceDescriptionSchema;
+
+ private final Map<String, Object> consumerConfig = Read.this.getConsumerConfig();
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = Read.this.getConsumerFactoryFn();
- if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
- transform =
- unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ private final List<String> topics = Read.this.getTopics();
+
+ private final List<TopicPartition> topicPartitions = Read.this.getTopicPartitions();
+
+ private final Instant startReadTime = Read.this.getStartReadTime();
+
+ @VisibleForTesting
+ Map<String, Object> getConsumerConfig() {
+ return consumerConfig;
}
- return input.getPipeline().apply(transform);
+ @VisibleForTesting
+ List<String> getTopics() {
+ return topics;
+ }
+
+ @ProcessElement
+ public void processElement(OutputReceiver<Row> receiver) {
+ List<TopicPartition> partitions = new ArrayList<>(topicPartitions);
+ if (partitions.isEmpty()) {
+ try (Consumer<?, ?> consumer = consumerFactoryFn.apply(consumerConfig)) {
+ for (String topic : topics) {
+ for (PartitionInfo p : consumer.partitionsFor(topic)) {
+ partitions.add(new TopicPartition(p.topic(), p.partition()));
+ }
+ }
+ }
+ }
+ partitions.stream()
+ .forEach(
+ topicPartition -> {
+ FieldValueBuilder descriptiorBuilder =
Review comment:
descriptiorBuilder -> descriptorBuilder
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
Review comment:
I believe this is the default implementation as well but I don't see it in the documentation so not sure if this is needed or not.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ Instant outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize.update(recordSize);
+ avgOffsetGap.update(expectedOffset - rawRecord.offset());
+ expectedOffset = rawRecord.offset() + 1;
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
+ throw anyException;
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize = new KafkaIOUtils.MovingAvg();
Review comment:
Will `avgRecordSize` / `avgOffsetGap` be consistent across multiple topics / partitions.
If we need to track them at the topic/partition level then they will need to part of the restriction.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -795,6 +828,12 @@ public void setValueDeserializer(String valueDeserializer) {
return withWatermarkFn2(unwrapKafkaAndThen(watermarkFn));
}
+ /** A function to the compute output timestamp from a {@link KafkaRecord}. */
+ public Read<K, V> withExtractOutputTimestampFn(
Review comment:
How is this different from withTimestampFn2?
We could fix the "translation" logic by storing each property the user sets and clear other properties that would be affected. Then in the expand step we can "translate" timestampFn2 to a WatermarkPolicy for the UnboundedSource version and use it directly in the SDF version.
Setting the top level properties allow us to say that this property is supported when used as an SDF.
Also, what prevents us from supporting TimestampPolicy? We should be able to call it and give it the three pieces of information it requests (message backlog / backlog check time / current kafka record).
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -681,11 +696,13 @@ public void setValueDeserializer(String valueDeserializer) {
}
/**
- * Sets {@link TimestampPolicy} to {@link TimestampPolicyFactory.LogAppendTimePolicy}. The
- * policy assigns Kafka's log append time (server side ingestion time) to each record. The
- * watermark for each Kafka partition is the timestamp of the last record read. If a partition
- * is idle, the watermark advances to couple of seconds behind wall time. Every record consumed
- * from Kafka is expected to have its timestamp type set to 'LOG_APPEND_TIME'.
+ * Sets {@link TimestampPolicy} to {@link TimestampPolicyFactory.LogAppendTimePolicy} which is
+ * used when beam_fn_api is disabled, and sets {@code extractOutputTimestampFn} as {@link
+ * ReadViaSDF.ExtractOutputTimestampFns#withLogAppendTime()}, which is used when beam_fn_api is
+ * enabled. The policy assigns Kafka's log append time (server side ingestion time) to each
+ * record. The watermark for each Kafka partition is the timestamp of the last record read. If a
+ * partition is idle, the watermark advances to couple of seconds behind wall time. Every record
+ * consumed from Kafka is expected to have its timestamp type set to 'LOG_APPEND_TIME'.
Review comment:
It looks like we overwrite the setExtractOutputTimestampFn regardless of whether the experiment is enabled or not which doesn't align with what the comment is telling us.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +946,123 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
+
+ return input.getPipeline().apply(transform);
+ } else {
+ // If extractOutputTimestampFn is not set, use processing time by default.
+ SerializableFunction<KafkaRecord<K, V>, Instant> timestampFn;
+ if (getExtractOutputTimestampFn() != null) {
+ timestampFn = getExtractOutputTimestampFn();
+ } else {
+ timestampFn = ReadViaSDF.ExtractOutputTimestampFns.useProcessingTime();
+ }
+ ReadViaSDF<K, V> readTransform =
+ ReadViaSDF.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withExtractOutputTimestampFn(timestampFn);
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(
+ ParDo.of(
+ new GenerateKafkaSourceDescription(
+ readTransform.getKafkaSourceDescriptionSchema())))
+ .setCoder(RowCoder.of(readTransform.getKafkaSourceDescriptionSchema()))
+ .apply(readTransform)
+ .setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
+ }
+ }
+
+ /**
+ * A DoFn which generates {@link Row} with {@link KafkaSourceDescriptionSchemas#getSchema()}
+ * based on the configuration of {@link Read}.
+ */
+ @VisibleForTesting
+ class GenerateKafkaSourceDescription extends DoFn<byte[], Row> {
+ GenerateKafkaSourceDescription(Schema schema) {
+ this.kafkaSourceDescriptionSchema = schema;
+ }
+
+ private final Schema kafkaSourceDescriptionSchema;
+
+ private final Map<String, Object> consumerConfig = Read.this.getConsumerConfig();
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = Read.this.getConsumerFactoryFn();
- if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
- transform =
- unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ private final List<String> topics = Read.this.getTopics();
+
+ private final List<TopicPartition> topicPartitions = Read.this.getTopicPartitions();
+
+ private final Instant startReadTime = Read.this.getStartReadTime();
+
+ @VisibleForTesting
+ Map<String, Object> getConsumerConfig() {
Review comment:
instead of adding methods, make the member variable `@VisibleForTesting`
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +946,123 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
+
+ return input.getPipeline().apply(transform);
+ } else {
+ // If extractOutputTimestampFn is not set, use processing time by default.
+ SerializableFunction<KafkaRecord<K, V>, Instant> timestampFn;
+ if (getExtractOutputTimestampFn() != null) {
+ timestampFn = getExtractOutputTimestampFn();
+ } else {
+ timestampFn = ReadViaSDF.ExtractOutputTimestampFns.useProcessingTime();
+ }
+ ReadViaSDF<K, V> readTransform =
+ ReadViaSDF.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withExtractOutputTimestampFn(timestampFn);
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(
+ ParDo.of(
+ new GenerateKafkaSourceDescription(
+ readTransform.getKafkaSourceDescriptionSchema())))
+ .setCoder(RowCoder.of(readTransform.getKafkaSourceDescriptionSchema()))
+ .apply(readTransform)
+ .setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
+ }
+ }
+
+ /**
+ * A DoFn which generates {@link Row} with {@link KafkaSourceDescriptionSchemas#getSchema()}
+ * based on the configuration of {@link Read}.
+ */
+ @VisibleForTesting
+ class GenerateKafkaSourceDescription extends DoFn<byte[], Row> {
+ GenerateKafkaSourceDescription(Schema schema) {
+ this.kafkaSourceDescriptionSchema = schema;
+ }
+
+ private final Schema kafkaSourceDescriptionSchema;
+
+ private final Map<String, Object> consumerConfig = Read.this.getConsumerConfig();
Review comment:
You should make the class static to prevent pulling in the parent transform during serialization. This would require passing forward arguments from the Read transform to the constructor. This will prevent accidentally capturing things we don't want captured.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
Review comment:
Shouldn't this be a construction time error instead of a runtime warning?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ Instant outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize.update(recordSize);
+ avgOffsetGap.update(expectedOffset - rawRecord.offset());
+ expectedOffset = rawRecord.offset() + 1;
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
+ throw anyException;
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize = new KafkaIOUtils.MovingAvg();
+ avgOffsetGap = new KafkaIOUtils.MovingAvg();
+ consumerSpEL = new ConsumerSpEL();
+ keyDeserializerInstance = keyDeserializerProvider.getDeserializer(consumerConfig, true);
+ valueDeserializerInstance = valueDeserializerProvider.getDeserializer(consumerConfig, false);
+ }
+
+ @Teardown
+ public void teardown() throws Exception {
+ try {
+ Closeables.close(keyDeserializerInstance, true);
+ Closeables.close(valueDeserializerInstance, true);
+ avgRecordSize = null;
Review comment:
these will go out of scope so no need to setup them to null in teardown.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
Review comment:
I would not make this class public and instead add future variants such as readAll to KafkaIO keeping this as an implementation detail.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
Review comment:
```suggestion
* For more details about the concept of {@code SplittableDoFn}, please refer to the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a>
* and <a href="https://s.apache.org/beam-fn-api">design doc</a>.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
Review comment:
Why not invoke `initialRestriction` and pass in the necessary arguments instead of duplicating the setup work.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
Review comment:
Why store these as member variables if we plan to pull them from the parent enclosed class?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
Review comment:
ReadFromKafkaViaSDF?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
Review comment:
Why would we close this when the object goes out of scope?
Would it make sense to have support for a KafkaConsumer pool that expires automatically after X amount of time of not being used?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r438949999
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
> I'm agree with @iemejia in the way that `HbaseIO` is a good example when we can end up with reimplementing new `Read` in the dedicated initially more light class, like `KafkaSourceDescription`, in the end.
>
> 1. Yes, it's quite heavy but we don't expect too many elements like this in the PCollection, right? So it should not be a problem.
> 2. Of course, reading from several clusters is not a main feature but, afair, we had several requests on that from users.
> 3. Agree, this is a good point. Do we have any principal objections to use `Read` for that?
Hi Alexey, I want to elaborate more details on having bootstrapServer dynamically. I'm wondering what should we do if the user sets bootstrapServer in the consumer config but also emits different bootstrapServer dynamically. In this case, should we override the consumer config with the dynamic one or we just update the consumer config with the new one?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-641001191
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442951203
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -681,11 +696,13 @@ public void setValueDeserializer(String valueDeserializer) {
}
/**
- * Sets {@link TimestampPolicy} to {@link TimestampPolicyFactory.LogAppendTimePolicy}. The
- * policy assigns Kafka's log append time (server side ingestion time) to each record. The
- * watermark for each Kafka partition is the timestamp of the last record read. If a partition
- * is idle, the watermark advances to couple of seconds behind wall time. Every record consumed
- * from Kafka is expected to have its timestamp type set to 'LOG_APPEND_TIME'.
+ * Sets {@link TimestampPolicy} to {@link TimestampPolicyFactory.LogAppendTimePolicy} which is
+ * used when beam_fn_api is disabled, and sets {@code extractOutputTimestampFn} as {@link
+ * ReadViaSDF.ExtractOutputTimestampFns#withLogAppendTime()}, which is used when beam_fn_api is
+ * enabled. The policy assigns Kafka's log append time (server side ingestion time) to each
+ * record. The watermark for each Kafka partition is the timestamp of the last record read. If a
+ * partition is idle, the watermark advances to couple of seconds behind wall time. Every record
+ * consumed from Kafka is expected to have its timestamp type set to 'LOG_APPEND_TIME'.
Review comment:
The ` which is used when beam_fn_api is disabled` means `TimestampPolicy` is used when beam_fn_api is disabled. I'll say it explicitly.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445235681
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1261,341 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadAll<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadAll.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract ReadAll.Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract ReadAll.Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract ReadAll.Builder<K, V> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract ReadAll.Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract ReadAll.Builder<K, V> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadAll.Builder<K, V> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadAll.Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract ReadAll.Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract ReadAll.Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract ReadAll.Builder<K, V> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn);
+
+ abstract ReadAll.Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadAll.Builder<K, V> setTimestampPolicyFactory(TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadAll<K, V> build();
+ }
- for (String key : updates.keySet()) {
+ public static <K, V> ReadAll<K, V> read() {
+ return new AutoValue_KafkaIO_ReadAll.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the
+ // element
+ // will override the bootstrapServers from the config.
+ public ReadAll<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadAll<K, V> withKeyDeserializerProvider(DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadAll<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadAll<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadAll<K, V> withValueDeserializer(Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadAll<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadAll<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadAll<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadAll<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadAll<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadAll<K, V> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadAll<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadAll<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadAll<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ReadAll.ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadAll<K, V> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new WallTime(state);
+ });
+ }
+
+ public ReadAll<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadAll<K, V> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
+ // default
+ // value.
+ public ReadAll<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadAll<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadAll<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadAll<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadAllFromRow forExternalBuild() {
+ return new ReadAllFromRow(this);
+ }
+
+ // This transform is used in cross-language case. The input Row should be encoded with an
+ // equivalent schema as KafkaSourceDescription.
+ private static class ReadAllFromRow<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadAll<K, V> readViaSDF;
+
+ ReadAllFromRow(ReadAll read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(Convert.fromRows(KafkaSourceDescription.class))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
Review comment:
I plan to have a separate PR to introduce external transform builder for ReadAll(). The `buildExternal` will be like:
`return build().forExternalBuild()`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r443052453
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ Instant outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize.update(recordSize);
+ avgOffsetGap.update(expectedOffset - rawRecord.offset());
+ expectedOffset = rawRecord.offset() + 1;
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
+ throw anyException;
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize = new KafkaIOUtils.MovingAvg();
Review comment:
Yes, currently `KafkaUnboundedReader` tracks `avgRecordSize` /`avgOffsetGap` per TopicPartition, but I don't think these values are serialized when checkpoint happens. We can bound `avgRecordSize` /`avgOffsetGap` to a TopicPartition by always creating a new `MovingAvg` when `processElement` is started.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r435354750
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
I'm agree with @iemejia in the way that `HbaseIO` is a good example when we can end up with reimplementing new `Read` in the dedicated initially more light class, like `KafkaSourceDescription`, in the end.
1. Yes, it's quite heavy but we don't expect too many elements like this in the PCollection, right? So it should not be a problem.
2. Of course, reading from several clusters is not a main feature but, afair, we had several requests on that from users.
3. Agree, this is a good point. Do we have any principal objections to use `Read` for that?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442573079
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
Review comment:
Yes, it should be an error.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r443937048
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
Review comment:
I forgot the default progress interval. Is it 5s or 5 mins?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r466673488
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1198,352 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadSourceDescriptors<K, V>
+ extends PTransform<PCollection<KafkaSourceDescriptor>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadSourceDescriptors.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract ReadSourceDescriptors.Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract ReadSourceDescriptors.Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setCommitOffsetEnabled(
+ boolean commitOffsetEnabled);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadSourceDescriptors<K, V> build();
+ }
+
+ public static <K, V> ReadSourceDescriptors<K, V> read() {
+ return new AutoValue_KafkaIO_ReadSourceDescriptors.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the
+ // element
+ // will override the bootstrapServers from the config.
+ public ReadSourceDescriptors<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new WallTime(state);
+ });
+ }
+
+ public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
+ // default value.
+ public ReadSourceDescriptors<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadSourceDescriptors<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ // TODO(BEAM-10320): Create external build transform for ReadSourceDescriptors().
Review comment:
Yes you are right. Thanks!
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r436123330
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
Sorry for the late.
> 3. Agree, this is a good point. Do we have any principal objections to use Read for that?
Yes, there are some concerns around that. As I mentioned in point1, `KafkaIO.Read` contains lots of configurations which are supported to determined during runtime. For these part, they shouldn't be the element logically. Also `KafkaIO.Read` is also used by the user who doesn't migrate to fnapi yet, which may last for a long time. It will not be easy for us to maintain or develop new things based on `KafkaIO.Read` when considering compatibility.
I think we should figure out the major needs from KafkaIO beam user, especially what they want to be dynamic during pipeline runtime. Alexey, do you know what can be a good place to collect such idea?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439566281
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
Thanks! I'll ping this thread once the PR is ready for another round of review.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439565391
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
Agree with option #2 as well.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444434591
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -295,21 +301,32 @@
/**
* Creates an uninitialized {@link Read} {@link PTransform}. Before use, basic Kafka configuration
* should set with {@link Read#withBootstrapServers(String)} and {@link Read#withTopics(List)}.
- * Other optional settings include key and value {@link Deserializer}s, custom timestamp and
+ * Other optional settings include key and value {@link Deserializer}s, custom timestamp,
* watermark functions.
*/
public static <K, V> Read<K, V> read() {
return new AutoValue_KafkaIO_Read.Builder<K, V>()
.setTopics(new ArrayList<>())
.setTopicPartitions(new ArrayList<>())
- .setConsumerFactoryFn(Read.KAFKA_CONSUMER_FACTORY_FN)
- .setConsumerConfig(Read.DEFAULT_CONSUMER_PROPERTIES)
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
.setMaxNumRecords(Long.MAX_VALUE)
.setCommitOffsetsInFinalizeEnabled(false)
.setTimestampPolicyFactory(TimestampPolicyFactory.withProcessingTime())
.build();
}
+ /**
+ * Creates an uninitialized {@link ReadViaSDF} {@link PTransform}. Different from {@link Read},
+ * setting up {@code topics} and {@code bootstrapServers} is not required during construction
+ * time. But the {@code bootstrapServers} still can be configured {@link
+ * ReadViaSDF#withBootstrapServers(String)}. Please refer to {@link ReadViaSDF} for more details.
+ */
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> readAll() {
+ return ReadViaSDF.<K, V, WatermarkEstimatorT>read();
Review comment:
The `WatermarkEstimatorT` is used when defining `createWatermarkEstimatorFn` and when `@NewWatermarkEstimator` is called. I understand that we can always use `WatermarkEstimator` as the type, I thought it would be better to make the type explicitly,
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637806567
@boyuanzz Yes, afaik, it's used only for [performance testing](http://metrics.beam.apache.org/d/bnlHKP3Wz/java-io-it-tests-dataflow?panelId=21&fullscreen&orgId=1&from=1588539922272&to=1591131922272) on Dataflow
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r436177315
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -949,45 +1077,6 @@ public void setValueDeserializer(String valueDeserializer) {
final SerializableFunction<KV<KeyT, ValueT>, OutT> fn) {
return record -> fn.apply(record.getKV());
}
- ///////////////////////////////////////////////////////////////////////////////////////
Review comment:
This common part is moved to the KafkaIOUtil.java.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev removed a comment on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev removed a comment on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637676917
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-638439090
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] TheNeuralBit commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
TheNeuralBit commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445181264
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
+ @SchemaFieldName("topic")
+ abstract String getTopic();
+
+ @SchemaFieldName("partition")
+ abstract Integer getPartition();
+
+ @SchemaFieldName("start_read_offset")
+ @Nullable
+ abstract Long getStartReadOffset();
+
+ @SchemaFieldName("start_read_time")
+ @Nullable
+ abstract Instant getStartReadTime();
+
+ @SchemaFieldName("bootstrapServers")
+ @Nullable
+ abstract List<String> getBootStrapServers();
+
+ private TopicPartition topicPartition = null;
+
+ public TopicPartition getTopicPartition() {
Review comment:
This will get pulled into the generated schema which I don't think is your intention. You should change the name so it's not a getter, or add `@SchemaIgnore`
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
+ @SchemaFieldName("topic")
+ abstract String getTopic();
+
+ @SchemaFieldName("partition")
+ abstract Integer getPartition();
+
+ @SchemaFieldName("start_read_offset")
+ @Nullable
+ abstract Long getStartReadOffset();
+
+ @SchemaFieldName("start_read_time")
+ @Nullable
+ abstract Instant getStartReadTime();
+
+ @SchemaFieldName("bootstrapServers")
+ @Nullable
+ abstract List<String> getBootStrapServers();
+
+ private TopicPartition topicPartition = null;
+
+ public TopicPartition getTopicPartition() {
+ if (topicPartition == null) {
+ topicPartition = new TopicPartition(getTopic(), getPartition());
+ }
+ return topicPartition;
+ }
+
+ public static KafkaSourceDescription of(
+ TopicPartition topicPartition,
+ Long startReadOffset,
+ Instant startReadTime,
+ List<String> bootstrapServers) {
+ return new AutoValue_KafkaSourceDescription(
+ topicPartition.topic(),
+ topicPartition.partition(),
+ startReadOffset,
+ startReadTime,
+ bootstrapServers);
+ }
+
+ public static Coder<KafkaSourceDescription> getCoder(SchemaRegistry registry)
+ throws NoSuchSchemaException {
+ return SchemaCoder.of(
+ registry.getSchema(KafkaSourceDescription.class),
+ TypeDescriptor.of(KafkaSourceDescription.class),
+ registry.getToRowFunction(KafkaSourceDescription.class),
+ registry.getFromRowFunction(KafkaSourceDescription.class));
+ }
Review comment:
I don't think you you should need this function. By registering a schema for KafkaSourceDescription Beam we will default to using SchemaCoder
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
+ @SchemaFieldName("topic")
+ abstract String getTopic();
+
+ @SchemaFieldName("partition")
+ abstract Integer getPartition();
+
+ @SchemaFieldName("start_read_offset")
+ @Nullable
+ abstract Long getStartReadOffset();
+
+ @SchemaFieldName("start_read_time")
+ @Nullable
+ abstract Instant getStartReadTime();
+
+ @SchemaFieldName("bootstrapServers")
Review comment:
+1
The default in the inferred schema should be camel-case with the first letter lower-case so this would be a no-op as written (same with topic and partition, but there's value in making them explicit if you want).
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +1082,91 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ }
+ ReadAll<K, V> readTransform =
+ ReadAll.<K, V>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withManualWatermarkEstimator()
+ .withTimestampPolicyFactory(getTimestampPolicyFactory());
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
+ PCollection<KafkaSourceDescription> output =
+ input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(ParDo.of(new GenerateKafkaSourceDescription(this)));
+ try {
+ output.setCoder(KafkaSourceDescription.getCoder(input.getPipeline().getSchemaRegistry()));
Review comment:
This should happen automatically because there's a schema registered for `KafkaSourceDescription`. Was that not working?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] piotr-szuberski commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
piotr-szuberski commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-671207415
Run Python 3.8 PostCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434729096
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
I would shorter the name of this class to `ReadWithSDF` or `ReadViaSDF` since this is clear that it's used to read from Kafka.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-649059330
Run Python2_PVR_Flink PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r467165240
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting value bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize value objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withValueDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
}
+ /**
+ * A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
+ * supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.
+ */
public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
+ */
public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
Map<String, Object> configUpdates) {
Map<String, Object> config =
KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
return toBuilder().setConsumerConfig(config).build();
}
+ /**
+ * A function to calculate output timestamp for a given {@link KafkaRecord}. The default value
+ * is {@link #withProcessingTime()}.
+ */
public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
return toBuilder().setExtractOutputTimestampFn(fn).build();
}
+ /**
+ * A function to create a {@link WatermarkEstimator}. The default value is {@link
+ * MonotonicallyIncreasing}.
+ */
public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
}
+ /** Use the log append time as the output timestamp. */
public ReadSourceDescriptors<K, V> withLogAppendTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
}
+ /** Use the processing time as the output timestamp. */
public ReadSourceDescriptors<K, V> withProcessingTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
}
+ /** Use the creation time of {@link KafkaRecord} as the output timestamp. */
public ReadSourceDescriptors<K, V> withCreateTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
}
+ /** Use the {@link WallTime} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new WallTime(state);
});
}
+ /** Use the {@link MonotonicallyIncreasing} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new MonotonicallyIncreasing(state);
});
}
+ /** Use the {@link Manual} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new Manual(state);
});
}
- // If a transactional producer is used and it's desired to only read records from committed
- // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
- // default value.
+ /**
+ * Sets "isolation_level" to "read_committed" in Kafka consumer configuration. This is ensures
+ * that the consumer does not read uncommitted messages. Kafka version 0.11 introduced
+ * transactional writes. Applications requiring end-to-end exactly-once semantics should only
+ * read committed messages. See JavaDoc for {@link KafkaConsumer} for more description.
Review comment:
```suggestion
* Sets "isolation_level" to "read_committed" in Kafka consumer configuration. This ensures
* that the consumer does not read uncommitted messages. Kafka version 0.11 introduced
* transactional writes. Applications requiring end-to-end exactly-once semantics should only
* read committed messages. See JavaDoc for {@link KafkaConsumer} for more description.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1288,9 +1294,12 @@ public void populateDisplayData(DisplayData.Builder builder) {
.withMonotonicallyIncreasingWatermarkEstimator();
}
- // Note that if the bootstrapServers is set here but also populated with the element, the
- // element
- // will override the bootstrapServers from the config.
+ /**
+ * Sets the bootstrap servers for the Kafka consumer. If the bootstrapServers is set here but
+ * also populated with the {@link KafkaSourceDescriptor}, the {@link
+ * KafkaSourceDescriptor#getBootStrapServers()} will override the bootstrapServers from the
+ * config.
Review comment:
```suggestion
* Sets the bootstrap servers to use for the Kafka consumer if unspecified via
* KafkaSourceDescriptor#getBootStrapServers()}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1439,9 +1542,18 @@ public void processElement(
}
}
+ /**
+ * Set the {@link TimestampPolicyFactory}. If the {@link TimestampPolicyFactory} is given, the
+ * output timestamp will be computed by the {@link
+ * TimestampPolicyFactory#createTimestampPolicy(TopicPartition, Optional)} and the {@link
+ * Manual} is used as the watermark estimator.
Review comment:
```suggestion
* Set the {@link TimestampPolicyFactory}. If the {@link TimestampPolicyFactory} is given, the
* output timestamp will be computed by the {@link
* TimestampPolicyFactory#createTimestampPolicy(TopicPartition, Optional)} and {@link
* Manual} is used as the watermark estimator.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting value bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize value objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withValueDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
}
+ /**
+ * A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
+ * supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.
+ */
public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
+ */
public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
Map<String, Object> configUpdates) {
Map<String, Object> config =
KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
return toBuilder().setConsumerConfig(config).build();
}
+ /**
+ * A function to calculate output timestamp for a given {@link KafkaRecord}. The default value
+ * is {@link #withProcessingTime()}.
+ */
public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
return toBuilder().setExtractOutputTimestampFn(fn).build();
}
+ /**
+ * A function to create a {@link WatermarkEstimator}. The default value is {@link
+ * MonotonicallyIncreasing}.
+ */
public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
}
+ /** Use the log append time as the output timestamp. */
public ReadSourceDescriptors<K, V> withLogAppendTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
}
+ /** Use the processing time as the output timestamp. */
public ReadSourceDescriptors<K, V> withProcessingTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
}
+ /** Use the creation time of {@link KafkaRecord} as the output timestamp. */
public ReadSourceDescriptors<K, V> withCreateTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
}
+ /** Use the {@link WallTime} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new WallTime(state);
});
}
+ /** Use the {@link MonotonicallyIncreasing} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new MonotonicallyIncreasing(state);
});
}
+ /** Use the {@link Manual} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new Manual(state);
});
}
- // If a transactional producer is used and it's desired to only read records from committed
- // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
- // default value.
+ /**
+ * Sets "isolation_level" to "read_committed" in Kafka consumer configuration. This is ensures
+ * that the consumer does not read uncommitted messages. Kafka version 0.11 introduced
+ * transactional writes. Applications requiring end-to-end exactly-once semantics should only
+ * read committed messages. See JavaDoc for {@link KafkaConsumer} for more description.
+ */
public ReadSourceDescriptors<K, V> withReadCommitted() {
return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
}
+ /**
+ * Enable committing record offset. If {@link #withReadCommitted()} or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set together with {@link #commitOffsets()},
+ * {@link #commitOffsets()} will be ignored.
+ */
public ReadSourceDescriptors<K, V> commitOffsets() {
return toBuilder().setCommitOffsetEnabled(true).build();
}
+ /**
+ * Set additional configuration for the backend offset consumer. It may be required for a
+ * secured Kafka cluster, especially when you see similar WARN log message 'exception while
+ * fetching latest offset for partition {}. will be retried'.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, offset consumer inherits the configuration from main consumer, with an
+ * auto-generated {@link ConsumerConfig#GROUP_ID_CONFIG}. This may not work in a secured Kafka
+ * which requires more configurations.
Review comment:
```suggestion
* Set additional configuration for the offset consumer. It may be required for a
* secured Kafka cluster, especially when you see similar WARN log message {@code exception while
* fetching latest offset for partition {}. will be retried}.
*
* <p>In {@link ReadFromKafkaDoFn}, there are two consumers running in the backend:
* <ol>
* <li>the main consumer which reads data from kafka.
* <li>the secondary offset consumer which is used to estimate the backlog by fetching the latest offset.
* </ol>
*
* <p>By default, offset consumer inherits the configuration from main consumer, with an
* auto-generated {@link ConsumerConfig#GROUP_ID_CONFIG}. This may not work in a secured Kafka
* which requires additional configuration.
*
* <p>See {@link #withConsumerConfigUpdates} for configuring the main consumer.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
Review comment:
```suggestion
* <p>Use this method to override the coder inference performed within {@link
* #withKeyDeserializer(Class)}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting value bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize value objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withValueDeserializer(Class)}.
Review comment:
```suggestion
* <p>Use this method to override the coder inference performed within {@link
* #withValueDeserializer(Class)}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting value bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize value objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withValueDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
}
+ /**
+ * A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
+ * supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.
+ */
public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
Review comment:
```suggestion
* Updates configuration for the main consumer. This method merges updates from the provided map with
* with any prior updates using {@link
* KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES} as the starting configuration.
*
* <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend:
* <ol>
* <li>the main consumer which reads data from kafka.
* <li>the secondary offset consumer which is used to estimate the backlog by fetching the latest offset.
* </ol>
*
* <p>See {@link #withConsumerConfigOverrides} for overriding the configuration instead of updating it.
*
* <p>See {@link #withOffsetConsumerConfigOverrides} for configuring the secondary offset consumer.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
Review comment:
This comment seems incorrect, it looks like it should be a clone of the key one above but stating that we are deserializing the value.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting value bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize value objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withValueDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
}
+ /**
+ * A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
+ * supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.
+ */
public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
+ */
public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
Map<String, Object> configUpdates) {
Map<String, Object> config =
KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
return toBuilder().setConsumerConfig(config).build();
}
+ /**
+ * A function to calculate output timestamp for a given {@link KafkaRecord}. The default value
+ * is {@link #withProcessingTime()}.
+ */
public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
return toBuilder().setExtractOutputTimestampFn(fn).build();
}
+ /**
+ * A function to create a {@link WatermarkEstimator}. The default value is {@link
+ * MonotonicallyIncreasing}.
+ */
public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
}
+ /** Use the log append time as the output timestamp. */
public ReadSourceDescriptors<K, V> withLogAppendTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
}
+ /** Use the processing time as the output timestamp. */
public ReadSourceDescriptors<K, V> withProcessingTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
}
+ /** Use the creation time of {@link KafkaRecord} as the output timestamp. */
public ReadSourceDescriptors<K, V> withCreateTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
}
+ /** Use the {@link WallTime} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new WallTime(state);
});
}
+ /** Use the {@link MonotonicallyIncreasing} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new MonotonicallyIncreasing(state);
});
}
+ /** Use the {@link Manual} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new Manual(state);
});
}
- // If a transactional producer is used and it's desired to only read records from committed
- // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
- // default value.
+ /**
+ * Sets "isolation_level" to "read_committed" in Kafka consumer configuration. This is ensures
+ * that the consumer does not read uncommitted messages. Kafka version 0.11 introduced
+ * transactional writes. Applications requiring end-to-end exactly-once semantics should only
+ * read committed messages. See JavaDoc for {@link KafkaConsumer} for more description.
+ */
public ReadSourceDescriptors<K, V> withReadCommitted() {
return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
}
+ /**
+ * Enable committing record offset. If {@link #withReadCommitted()} or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set together with {@link #commitOffsets()},
+ * {@link #commitOffsets()} will be ignored.
+ */
public ReadSourceDescriptors<K, V> commitOffsets() {
return toBuilder().setCommitOffsetEnabled(true).build();
}
+ /**
+ * Set additional configuration for the backend offset consumer. It may be required for a
+ * secured Kafka cluster, especially when you see similar WARN log message 'exception while
+ * fetching latest offset for partition {}. will be retried'.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, offset consumer inherits the configuration from main consumer, with an
+ * auto-generated {@link ConsumerConfig#GROUP_ID_CONFIG}. This may not work in a secured Kafka
+ * which requires more configurations.
+ */
public ReadSourceDescriptors<K, V> withOffsetConsumerConfigOverrides(
Map<String, Object> offsetConsumerConfig) {
return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
+ */
public ReadSourceDescriptors<K, V> withConsumerConfigOverrides(
Map<String, Object> consumerConfig) {
return toBuilder().setConsumerConfig(consumerConfig).build();
}
- // TODO(BEAM-10320): Create external build transform for ReadSourceDescriptors().
ReadAllFromRow forExternalBuild() {
return new ReadAllFromRow(this);
}
- // This transform is used in cross-language case. The input Row should be encoded with an
- // equivalent schema as KafkaSourceDescriptor.
+ /**
+ * A transform that is used in cross-language case. The input Row should be encoded with an
+ * equivalent schema as {@link KafkaSourceDescriptor}.
Review comment:
```suggestion
* A transform that is used in cross-language case. The input {@link Row} should be encoded with an
* equivalent schema as {@link KafkaSourceDescriptor}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1306,112 +1315,206 @@ public void populateDisplayData(DisplayData.Builder builder) {
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
+ *
+ * <p>In addition, Beam also needs a {@link Coder} to serialize and deserialize key objects at
+ * runtime. KafkaIO tries to infer a coder for the key based on the {@link Deserializer} class,
+ * however in case that fails, you can use {@link #withKeyDeserializerAndCoder(Class, Coder)} to
+ * provide the key coder explicitly.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializer(
Class<? extends Deserializer<K>> keyDeserializer) {
return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializer(
Class<? extends Deserializer<V>> valueDeserializer) {
return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting key bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize key objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withKeyDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
}
+ /**
+ * Sets a Kafka {@link Deserializer} for interpreting value bytes read from Kafka along with a
+ * {@link Coder} for helping the Beam runner materialize value objects at runtime if necessary.
+ *
+ * <p>Use this method only if your pipeline doesn't work with plain {@link
+ * #withValueDeserializer(Class)}.
+ */
public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
}
+ /**
+ * A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
+ * supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.
+ */
public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
+ */
public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
Map<String, Object> configUpdates) {
Map<String, Object> config =
KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
return toBuilder().setConsumerConfig(config).build();
}
+ /**
+ * A function to calculate output timestamp for a given {@link KafkaRecord}. The default value
+ * is {@link #withProcessingTime()}.
+ */
public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
return toBuilder().setExtractOutputTimestampFn(fn).build();
}
+ /**
+ * A function to create a {@link WatermarkEstimator}. The default value is {@link
+ * MonotonicallyIncreasing}.
+ */
public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
}
+ /** Use the log append time as the output timestamp. */
public ReadSourceDescriptors<K, V> withLogAppendTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
}
+ /** Use the processing time as the output timestamp. */
public ReadSourceDescriptors<K, V> withProcessingTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
}
+ /** Use the creation time of {@link KafkaRecord} as the output timestamp. */
public ReadSourceDescriptors<K, V> withCreateTime() {
return withExtractOutputTimestampFn(
ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
}
+ /** Use the {@link WallTime} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new WallTime(state);
});
}
+ /** Use the {@link MonotonicallyIncreasing} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new MonotonicallyIncreasing(state);
});
}
+ /** Use the {@link Manual} as the watermark estimator. */
public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
return withCreatWatermarkEstimatorFn(
state -> {
return new Manual(state);
});
}
- // If a transactional producer is used and it's desired to only read records from committed
- // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
- // default value.
+ /**
+ * Sets "isolation_level" to "read_committed" in Kafka consumer configuration. This is ensures
+ * that the consumer does not read uncommitted messages. Kafka version 0.11 introduced
+ * transactional writes. Applications requiring end-to-end exactly-once semantics should only
+ * read committed messages. See JavaDoc for {@link KafkaConsumer} for more description.
+ */
public ReadSourceDescriptors<K, V> withReadCommitted() {
return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
}
+ /**
+ * Enable committing record offset. If {@link #withReadCommitted()} or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set together with {@link #commitOffsets()},
+ * {@link #commitOffsets()} will be ignored.
+ */
public ReadSourceDescriptors<K, V> commitOffsets() {
return toBuilder().setCommitOffsetEnabled(true).build();
}
+ /**
+ * Set additional configuration for the backend offset consumer. It may be required for a
+ * secured Kafka cluster, especially when you see similar WARN log message 'exception while
+ * fetching latest offset for partition {}. will be retried'.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, offset consumer inherits the configuration from main consumer, with an
+ * auto-generated {@link ConsumerConfig#GROUP_ID_CONFIG}. This may not work in a secured Kafka
+ * which requires more configurations.
+ */
public ReadSourceDescriptors<K, V> withOffsetConsumerConfigOverrides(
Map<String, Object> offsetConsumerConfig) {
return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
}
+ /**
+ * Update configuration for the backend main consumer. Note that the default consumer properties
+ * will not be completely overridden. This method only updates the value which has the same key.
+ *
+ * <p>In {@link ReadFromKafkaDoFn}, there're two consumers running in the backend actually:<br>
+ * 1. the main consumer, which reads data from kafka;<br>
+ * 2. the secondary offset consumer, which is used to estimate backlog, by fetching latest
+ * offset;<br>
+ *
+ * <p>By default, main consumer uses the configuration from {@link
+ * KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
Review comment:
```suggestion
* Replaces the configuration for the main consumer.
*
* <p>In {@link ReadFromKafkaDoFn}, there are two consumers running in the backend:
* <ol>
* <li>the main consumer which reads data from kafka.
* <li>the secondary offset consumer which is used to estimate the backlog by fetching the latest offset.
* </ol>
*
* <p>By default, main consumer uses the configuration from {@link
* KafkaIOUtils#DEFAULT_CONSUMER_PROPERTIES}.
*
* See {@link #withConsumerConfigUpdates} for updating the configuration instead of overriding it.
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442572993
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
Review comment:
I was under impression that `IO` means it will be the root transform. If that's not the case, `readAll` sounds good.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] iemejia commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
iemejia commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434666700
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
This seems strangely close to something we lived in the SDF version of HBaseIO. In the first version we did an artificial object called `HBaseQuery` that contained the minimum information we needed to be able to query the Data store in a SDF way, but then other requirements came in and we started to add extra parameters to end up with something that was almost close to the exact 'complete' specification of the Read class so we decided to switch to use a `PCollection<Read>` as input otherwise we will be duplicating code, so we ended up with https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java#L353
Here you can have `PCollection<Read>` as an input and get rid of `KafkaSourceDescription` and this will have a more consistent user experience for final users. Notice that this `ReadAll` like pattern is also now used now in [SolrIO](https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/solr/src/main/java/org/apache/beam/sdk/io/solr/SolrIO.java#L501) and there is an ongoing PR to introduce it for [CassandraIO](https://github.com/apache/beam/pull/10546) so maybe it is a good idea we follow it for consistency.
Notice that in the SolrIO case the change looks even closer to this one because we ended up introducing `ReplicaInfo` (the spiritual equivalent of `TopicPartition`) into normal Read and we guarantee in expansion that this field gets filled if the users don't do it, but if they do well we asume they know what they are doing and we go with it.
Another advantage of having the full specification is that you will be able to read not only from multiple topics but also from different clusters because of the power of having the full `Read` specification,
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444360263
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,73 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.values.Row;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with {@link Schemas#getSchema()} as a {@link Row} when crossing the wire.
+ */
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
+ abstract TopicPartition getTopicPartition();
+
+ @Nullable
+ abstract Long getStartReadOffset();
+
+ @Nullable
+ abstract Instant getStartReadTime();
+
+ @Nullable
+ abstract List<String> getBootStrapServers();
+
+ public static KafkaSourceDescription of(
+ TopicPartition topicPartition,
+ Long startReadOffset,
+ Instant startReadTime,
+ List<String> bootstrapServers) {
+ return new AutoValue_KafkaSourceDescription(
+ topicPartition, startReadOffset, startReadTime, bootstrapServers);
+ }
+
+ static class Schemas {
Review comment:
nit: Schemas -> Schema
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1002,7 +1053,8 @@ public void populateDisplayData(DisplayData.Builder builder) {
DisplayData.item("topicPartitions", Joiner.on(",").join(topicPartitions))
.withLabel("Topic Partition/s"));
}
- Set<String> ignoredConsumerPropertiesKeys = IGNORED_CONSUMER_PROPERTIES.keySet();
+ Set<String> ignoredConsumerPropertiesKeys =
Review comment:
nit: ignoredConsumerPropertiesKeys -> disallowedConsumerPropertiesKeys
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -295,21 +301,32 @@
/**
* Creates an uninitialized {@link Read} {@link PTransform}. Before use, basic Kafka configuration
* should set with {@link Read#withBootstrapServers(String)} and {@link Read#withTopics(List)}.
- * Other optional settings include key and value {@link Deserializer}s, custom timestamp and
+ * Other optional settings include key and value {@link Deserializer}s, custom timestamp,
* watermark functions.
*/
public static <K, V> Read<K, V> read() {
return new AutoValue_KafkaIO_Read.Builder<K, V>()
.setTopics(new ArrayList<>())
.setTopicPartitions(new ArrayList<>())
- .setConsumerFactoryFn(Read.KAFKA_CONSUMER_FACTORY_FN)
- .setConsumerConfig(Read.DEFAULT_CONSUMER_PROPERTIES)
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
.setMaxNumRecords(Long.MAX_VALUE)
.setCommitOffsetsInFinalizeEnabled(false)
.setTimestampPolicyFactory(TimestampPolicyFactory.withProcessingTime())
.build();
}
+ /**
+ * Creates an uninitialized {@link ReadViaSDF} {@link PTransform}. Different from {@link Read},
+ * setting up {@code topics} and {@code bootstrapServers} is not required during construction
+ * time. But the {@code bootstrapServers} still can be configured {@link
+ * ReadViaSDF#withBootstrapServers(String)}. Please refer to {@link ReadViaSDF} for more details.
+ */
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> readAll() {
+ return ReadViaSDF.<K, V, WatermarkEstimatorT>read();
Review comment:
I don't think it is important to the user to know which WatermarkEstimatorT that they are getting as the type variable doesn't provide any value to the user.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
Review comment:
Why not numOfRecords * DEFAULT_MESSAGE_SIZE?
Where DEFAULT_MESSAGE_SIZE=100 or 1000?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +926,89 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ } else {
+ ReadViaSDF<K, V, Manual> readTransform =
+ ReadViaSDF.<K, V, Manual>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withManualWatermarkEstimator()
+ .withTimestampPolicyFactory(getTimestampPolicyFactory());
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
- if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
- transform =
- unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(ParDo.of(new GenerateKafkaSourceDescription(this)))
+ .setCoder(SerializableCoder.of(KafkaSourceDescription.class))
+ .apply(readTransform)
+ .setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
Review comment:
read transforms expand method already sets the output coder appropriately so this is not necessary
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
Review comment:
```suggestion
// When there are no records available for the current TopicPartition, self-checkpoint
// and move to process the next element.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
Review comment:
This method should be able to be GrowableOffsetRangeTracker. This will allow you to avoid the HasProgress cast in getSize
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,73 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.values.Row;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with {@link Schemas#getSchema()} as a {@link Row} when crossing the wire.
Review comment:
Shouldn't we instead create a SchemaProviderRegistrar and the appropriate SchemaProvider instead of asking users to encode/decode with `Schemas#getSchema`?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +926,89 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ } else {
+ ReadViaSDF<K, V, Manual> readTransform =
+ ReadViaSDF.<K, V, Manual>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withManualWatermarkEstimator()
+ .withTimestampPolicyFactory(getTimestampPolicyFactory());
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
- if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
- transform =
- unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(ParDo.of(new GenerateKafkaSourceDescription(this)))
+ .setCoder(SerializableCoder.of(KafkaSourceDescription.class))
Review comment:
SchemaCoder?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +926,89 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ } else {
Review comment:
nit: We shouldn't need the additional indentation of the `else` part so we can structure this as a guard statement:
```
if (x) {
// handle special case
return;
}
// do default
return;
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
Review comment:
We should be using the SchemaProvider `fromRow` function here.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
Review comment:
Might want to have a comment on the class specifically stating that we are using a row here since the xlang representation works by using records encoded in the raw row format `KafkaSourceDescriptor#Schemas#getSchema`
nit: ReadViaSDFExternally -> ReadViaSDFFromRow
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
Review comment:
```suggestion
// Valid between bundle start and bundle finish.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
Review comment:
ReadFromKafkaViaSDF?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
Review comment:
nit: memorizedBacklog -> memoizedBacklog
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(expectedOffset - rawRecord.offset());
+ avgRecordSize
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(recordSize);
+ expectedOffset = rawRecord.offset() + 1;
+ Instant outputTimestamp;
+ // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+ // WatermarkEstimator should be a Manual one.
+ if (timestampPolicy != null) {
+ checkState(watermarkEstimator instanceof ManualWatermarkEstimator);
+ TimestampPolicyContext context =
+ new TimestampPolicyContext(
+ (long) ((HasProgress) tracker).getProgress().getWorkRemaining(),
+ Instant.now());
+ outputTimestamp = timestampPolicy.getTimestampForRecord(context, kafkaRecord);
+ ((ManualWatermarkEstimator) watermarkEstimator)
+ .setWatermark(timestampPolicy.getWatermark(context));
+ } else {
+ outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ }
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
Review comment:
I'm guessing that you don't actually need this and are using it for debugging.
Any reason why higher level logging that the process element call failed wouldn't be enough?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
Review comment:
nit: numOfRecords -> numRecords
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
Review comment:
It would be nice if instead we had one map/object that was able to handle both the offset gap and the record size.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
Review comment:
I know I suggested this but am wary now. I think it won't be as performant but checking to see if there are any messages seems like the safer bet (effectively the logic you had before).
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
Review comment:
```suggestion
// If there is no known work, resume with max timeout and move to the next element.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(expectedOffset - rawRecord.offset());
+ avgRecordSize
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(recordSize);
+ expectedOffset = rawRecord.offset() + 1;
Review comment:
If we update the expectedOffset to be currentOffset + 1, won't the expectedOffset - nextOffset typically be 0 when there are no gaps?
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(expectedOffset - rawRecord.offset());
+ avgRecordSize
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(recordSize);
+ expectedOffset = rawRecord.offset() + 1;
+ Instant outputTimestamp;
+ // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+ // WatermarkEstimator should be a Manual one.
+ if (timestampPolicy != null) {
+ checkState(watermarkEstimator instanceof ManualWatermarkEstimator);
+ TimestampPolicyContext context =
+ new TimestampPolicyContext(
+ (long) ((HasProgress) tracker).getProgress().getWorkRemaining(),
+ Instant.now());
Review comment:
The backlog check time doesn't accurately reflect when the value was memoized. Not sure if this will cause an issue.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442962149
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
Review comment:
Relying on finalize leads to having the JVM perform closing during garbage collection which means you can't control the timing of it and it could be a long time.
Lets stick with this for now and maybe there will be a better suggestion later if we think it will be a problem.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] iemejia commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
iemejia commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434666700
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
This seems strangely close to something we lived in the SDF version of HBaseIO. In the first version we did an artificial object called `HBaseQuery` that contained the minimum information we needed to be able to query the Data store in a SDF way, but then other requirements came in and we started to add extra parameters to end up with something that was almost close to the exact 'complete' specification of the Read class so we decided to read from a PCollection<Read> otherwise we will be duplicating code, so we ended up with https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java#L353
Here you can have `PCollection<Read>` as an input and get rid of `KafkaSourceDescription` and this will have a more consistent user experience for final users. Notice that this `ReadAll` like pattern is also now used now in [SolrIO](https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/solr/src/main/java/org/apache/beam/sdk/io/solr/SolrIO.java#L501) and there is an ongoing PR to introduce it for [CassandraIO](https://github.com/apache/beam/pull/10546) so maybe it is a good idea we follow it for consistency.
Notice that in the SolrIO case the change looks even closer to this one because we ended up introducing `ReplicaInfo` (the spiritual equivalent of `TopicPartition`) into normal Read and we guarantee in expansion that this field gets filled if the users don't do it, but if they do well we asume they know what they are doing and we go with it.
Another advantage of having the full specification is that you will be able to read not only from multiple topics but also from different clusters because of the power of having the full `Read` specification,
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444380726
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
Review comment:
I would be helpful to add a comment why this number was chosen.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-671463121
> @boyuanzz This change has broken python postcommits that run kafka cross-language it tests.
>
> Could you take a look on that? I tried to fix it but it seems it would take me a lot of time since I'm not familiar with those SDF changes - it's possible that it's something straightforward. I tried to remove --experiments=beam_fn_api to run cross-language tests with deprecated transforms without SDF, but without luck.
>
> A stacktrace fragment:
>
> ```
> Caused by: java.util.concurrent.ExecutionException: java.lang.RuntimeException: Error received from SDK harness for instruction 2: java.util.concurrent.ExecutionException: java.lang.RuntimeException: Could not find a way to create AutoValue class class org.apache.beam.sdk.io.kafka.KafkaSourceDescriptor
> at java.util.concurrent.CompletableFuture.reportGet(CompletableFuture.java:357)
> at java.util.concurrent.CompletableFuture.get(CompletableFuture.java:1908)
> at org.apache.beam.sdk.fn.data.CompletableFutureInboundDataClient.awaitCompletion(CompletableFutureInboundDataClient.java:48)
> at org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver.awaitCompletion(BeamFnDataInboundObserver.java:91)
> at org.apache.beam.fn.harness.BeamFnDataReadRunner.blockTillReadFinishes(BeamFnDataReadRunner.java:342)
> at org.apache.beam.fn.harness.data.PTransformFunctionRegistry.lambda$register$0(PTransformFunctionRegistry.java:108)
> at org.apache.beam.fn.harness.control.ProcessBundleHandler.processBundle(ProcessBundleHandler.java:302)
> at org.apache.beam.fn.harness.control.BeamFnControlClient.delegateOnInstructionRequestType(BeamFnControlClient.java:173)
> at org.apache.beam.fn.harness.control.BeamFnControlClient.lambda$processInstructionRequests$0(BeamFnControlClient.java:157)
> at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
> at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
> at java.lang.Thread.run(Thread.java:748)
> ```
Looking.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442957543
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
Review comment:
If you take a look at FileIO, you can see we have PTransforms like ReadMatches, MatchAll, ...
I think we can do the same with KafkaIO where we add ReadTopics<PCollection<String>, PCollection<KafkaRecord>> and these are wrappers over the ReadViaSDF implementation with a transform that converts String -> Row in this example.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444354311
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
Review comment:
I believe the default is 30 seconds for Dataflow but I don't thing we should tune this to be Dataflow specific and more about making this relative to the average cost it takes to figure this out. 5 seconds doesn't sound too bad.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r466680059
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +209,102 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadSourceDescriptors} is the {@link PTransform} that takes a PCollection of {@link
Review comment:
The order is explained in the javadoc of ReadFromKafkaDoFn, [Initial Restriction] section.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r464667015
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +926,89 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ // The Read will be expanded into SDF transform when "beam_fn_api" is enabled and
+ // "beam_fn_api_use_deprecated_read" is not enabled.
+ if (!ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
+ || ExperimentalOptions.hasExperiment(
+ input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
+ // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
+ Unbounded<KafkaRecord<K, V>> unbounded =
+ org.apache.beam.sdk.io.Read.from(
+ toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+
+ PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+
+ if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
+ transform =
+ unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ }
- PTransform<PBegin, PCollection<KafkaRecord<K, V>>> transform = unbounded;
+ return input.getPipeline().apply(transform);
+ } else {
+ ReadViaSDF<K, V, Manual> readTransform =
+ ReadViaSDF.<K, V, Manual>read()
+ .withConsumerConfigOverrides(getConsumerConfig())
+ .withOffsetConsumerConfigOverrides(getOffsetConsumerConfig())
+ .withConsumerFactoryFn(getConsumerFactoryFn())
+ .withKeyDeserializerProvider(getKeyDeserializerProvider())
+ .withValueDeserializerProvider(getValueDeserializerProvider())
+ .withManualWatermarkEstimator()
+ .withTimestampPolicyFactory(getTimestampPolicyFactory());
+ if (isCommitOffsetsInFinalizeEnabled()) {
+ readTransform = readTransform.commitOffsets();
+ }
- if (getMaxNumRecords() < Long.MAX_VALUE || getMaxReadTime() != null) {
- transform =
- unbounded.withMaxReadTime(getMaxReadTime()).withMaxNumRecords(getMaxNumRecords());
+ return input
+ .getPipeline()
+ .apply(Impulse.create())
+ .apply(ParDo.of(new GenerateKafkaSourceDescription(this)))
+ .setCoder(SerializableCoder.of(KafkaSourceDescription.class))
+ .apply(readTransform)
+ .setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
Review comment:
The coder needs to be set when in x-lang case. It seems like there is something not correct when x-lang expand the transform.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r443059352
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ Instant outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize.update(recordSize);
+ avgOffsetGap.update(expectedOffset - rawRecord.offset());
+ expectedOffset = rawRecord.offset() + 1;
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
+ throw anyException;
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize = new KafkaIOUtils.MovingAvg();
Review comment:
That makes more sense but relies on knowing the lifecycle of when things get called. What if GetSize is invoked before ProcessElement instead of after. This will lead to implementation assumptions which will prevent SDF evolution.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434766297
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
+ return new AutoValue_ReadFromKafkaViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .build();
+ }
+
Review comment:
As mentioned in the javadoc of `ReadFromKafkaViaSDF`, all consumer-related configurations are the same to `KafkaIO.Read`. `ReadFromKafkaViaSDF` has 2 special configurations:
1. `setExtractOutputTimestampFn` to set a function to extract output timestamp from `KafkaRecord`
2. `setCommitOffsetEnabled` to enable committing offset manually.
The duplications are all for building the transform with common configurations. I was thinking about refactoring `KafkaIO` into a factory like pattern, which can build a `KafkaIO.Read` or `ReadFromKafkaViaSDF`. This refactor is not backward compatible for the API user, which means the user need to rewrite the pipeline if they want to work with the new SDK. My purpose here is to make the code changes for the use as less as possible. But if that's a minor issue, refactor is the way to go.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637676403
@boyuanzz Thanks! I'll take a look asap.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439562951
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
Since we want to empower the user to have bootstrapServer on the fly, I think we we have 2 options:
1. We remove the `bootstrapServers` from construction time configuration and `KafkaSourceDescription` will be the source of truth.
2. We still allow setting default `bootstrapServers` at construction time, and treat `KafkaSourceDescription.bootstrapServers` as override.
I'm leaning to #2 as long as it's easy for the user to understand the meaning of override.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r443040694
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
Review comment:
Yep the default implementation is doing nothing: https://github.com/apache/beam/blob/591de3473144de54beef0932131025e2a4d8504b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java#L300-L307
But I still want to keep the implementation there to make it explicitly.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442575623
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
Review comment:
`initialRestriction` still gives the infinite range. I should invoke `restrictionTracker().getProgress()`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637922023
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434745710
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -906,19 +955,110 @@ public void setValueDeserializer(String valueDeserializer) {
Coder<K> keyCoder = getKeyCoder(coderRegistry);
Coder<V> valueCoder = getValueCoder(coderRegistry);
- // Handles unbounded source to bounded conversion if maxNumRecords or maxReadTime is set.
- Unbounded<KafkaRecord<K, V>> unbounded =
- org.apache.beam.sdk.io.Read.from(
- toBuilder().setKeyCoder(keyCoder).setValueCoder(valueCoder).build().makeSource());
+ if (!isUseSDFTransform()
+ || !ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
Review comment:
SDF is only supported over `beam_fn_api`. We shouldn't expand the `KafkaIO.Read` with SDF when the `beam_fn_api` is not enbaled, or `beam_fn_api_use_deprecated_read ` is enabled.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] piotr-szuberski removed a comment on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
piotr-szuberski removed a comment on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-671207415
Run Python 3.8 PostCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442524939
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract ReadFromKafkaViaSDF<K, V> build();
+ }
+
+ public static <K, V> ReadFromKafkaViaSDF<K, V> create() {
Review comment:
I would suggest making ReadViaSDF package private and add a ReadAll PTransform to KafkaIO.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444564026
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
Review comment:
It should be `ReadAll`.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] piotr-szuberski edited a comment on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
piotr-szuberski edited a comment on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-671207252
@boyuanzz This change has broken python postcommits that run kafka cross-language it tests.
Could you take a look on that? I tried to fix it but it seems it would take me a lot of time since I'm not familiar with those SDF changes - it's possible that it's something straightforward. I tried to remove --experiments=beam_fn_api to run cross-language tests with deprecated transforms without SDF, but without luck.
A stacktrace fragment:
```
Caused by: java.util.concurrent.ExecutionException: java.lang.RuntimeException: Error received from SDK harness for instruction 2: java.util.concurrent.ExecutionException: java.lang.RuntimeException: Could not find a way to create AutoValue class class org.apache.beam.sdk.io.kafka.KafkaSourceDescriptor
at java.util.concurrent.CompletableFuture.reportGet(CompletableFuture.java:357)
at java.util.concurrent.CompletableFuture.get(CompletableFuture.java:1908)
at org.apache.beam.sdk.fn.data.CompletableFutureInboundDataClient.awaitCompletion(CompletableFutureInboundDataClient.java:48)
at org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver.awaitCompletion(BeamFnDataInboundObserver.java:91)
at org.apache.beam.fn.harness.BeamFnDataReadRunner.blockTillReadFinishes(BeamFnDataReadRunner.java:342)
at org.apache.beam.fn.harness.data.PTransformFunctionRegistry.lambda$register$0(PTransformFunctionRegistry.java:108)
at org.apache.beam.fn.harness.control.ProcessBundleHandler.processBundle(ProcessBundleHandler.java:302)
at org.apache.beam.fn.harness.control.BeamFnControlClient.delegateOnInstructionRequestType(BeamFnControlClient.java:173)
at org.apache.beam.fn.harness.control.BeamFnControlClient.lambda$processInstructionRequests$0(BeamFnControlClient.java:157)
at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
at java.lang.Thread.run(Thread.java:748)
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637676761
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r466022189
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +209,102 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadSourceDescriptors} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescriptor} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadSourceDescriptors} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from a BigQuery table and read these topics via {@link ReadSourceDescriptors}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadSourceDescriptors#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadSourceDescriptors#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadSourceDescriptors#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadSourceDescriptors#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadSourceDescriptors#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadSourceDescriptors#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadSourceDescriptors#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadSourceDescriptors#isCommitOffsetEnabled()} has the same meaning as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadSourceDescriptors} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescriptor.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * Note that the {@code bootstrapServers} can also be populated from the {@link
+ * KafkaSourceDescriptor}:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescriptor.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadSourceDescriptors}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadSourceDescriptors#commitOffsets()} enables committing offset after processing the
+ * record. Note that if the {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadSourceDescriptors#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadSourceDescriptors#withExtractOutputTimestampFn(SerializableFunction)} is used to
+ * compute the {@code output timestamp} for a given {@link KafkaRecord}. There are three built-in
+ * types: {@link ReadSourceDescriptors#withProcessingTime()}, {@link
+ * ReadSourceDescriptors#withCreateTime()} and {@link ReadSourceDescriptors#withLogAppendTime()}.
Review comment:
use an unorderered list for the options: `<ul><li></ul>`
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +209,102 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadSourceDescriptors} is the {@link PTransform} that takes a PCollection of {@link
Review comment:
We should state the order of preference for how we start reading (offset/timestamp/last commit offset) in this block somewhere.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescriptor }. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn#restrictionTracker(KafkaSourceDescriptor, OffsetRange)} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescriptor, OffsetRange).}
Review comment:
```suggestion
* <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescriptor, OffsetRange)}.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1198,352 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadSourceDescriptors<K, V>
+ extends PTransform<PCollection<KafkaSourceDescriptor>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadSourceDescriptors.class);
+
+ abstract Map<String, Object> getConsumerConfig();
Review comment:
It is unfortunate that Read doesn't have the documentation for these methods that we could copy.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1198,352 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
Review comment:
In the javadoc, we should state:
* what the defaults are if nothing is configured
* all the withYYY methods, we should state that the KafkaSourceDescriptor.x takes precendence and if unset then we default to using this method.
* all the KafkaSourceDescriptor methods should state what they override in the ReadSourceDescriptors transform
* that the watermark is controlled by one of withCreateWatermarkEstimatorFn / withWallTimeWatermarkEstimator / ...
* that the output timestamp is controlled by one of withProcessingTime / withLogAppendTime / ...
some of this Javadoc should go on KafkaIO while others makes more sense to place on the methods.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
Review comment:
Most of my comments above about Javadoc can be copied directly from the comments you have made here.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
Review comment:
```suggestion
* <p>The initial range for a {@link KafkaSourceDescriptor} is defined by {@code [startOffset,
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -198,6 +209,102 @@
* ...
* }</pre>
*
+ * <h2>Read from Kafka as a {@link DoFn}</h2>
+ *
+ * {@link ReadSourceDescriptors} is the {@link PTransform} that takes a PCollection of {@link
+ * KafkaSourceDescriptor} as input and outputs a PCollection of {@link KafkaRecord}. The core
+ * implementation is based on {@code SplittableDoFn}. For more details about the concept of {@code
+ * SplittableDoFn}, please refer to the <a
+ * href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadSourceDescriptors} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from a BigQuery table and read these topics via {@link ReadSourceDescriptors}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadSourceDescriptors#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadSourceDescriptors#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadSourceDescriptors#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadSourceDescriptors#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadSourceDescriptors#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadSourceDescriptors#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadSourceDescriptors#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadSourceDescriptors#isCommitOffsetEnabled()} has the same meaning as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadSourceDescriptors} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescriptor.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * Note that the {@code bootstrapServers} can also be populated from the {@link
+ * KafkaSourceDescriptor}:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescriptor.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadSourceDescriptors}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadSourceDescriptors#commitOffsets()} enables committing offset after processing the
+ * record. Note that if the {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadSourceDescriptors#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadSourceDescriptors#withExtractOutputTimestampFn(SerializableFunction)} is used to
+ * compute the {@code output timestamp} for a given {@link KafkaRecord}. There are three built-in
Review comment:
```suggestion
* compute the {@code output timestamp} for a given {@link KafkaRecord} and controls the watermark advancement. There are three built-in
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescriptor }. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
Review comment:
```suggestion
* KafkaSourceDescriptor}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1198,352 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadSourceDescriptors<K, V>
+ extends PTransform<PCollection<KafkaSourceDescriptor>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadSourceDescriptors.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract ReadSourceDescriptors.Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract ReadSourceDescriptors.Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setCommitOffsetEnabled(
+ boolean commitOffsetEnabled);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadSourceDescriptors<K, V> build();
+ }
+
+ public static <K, V> ReadSourceDescriptors<K, V> read() {
+ return new AutoValue_KafkaIO_ReadSourceDescriptors.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the
+ // element
+ // will override the bootstrapServers from the config.
+ public ReadSourceDescriptors<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new WallTime(state);
+ });
+ }
+
+ public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
+ // default value.
+ public ReadSourceDescriptors<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadSourceDescriptors<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ // TODO(BEAM-10320): Create external build transform for ReadSourceDescriptors().
Review comment:
This looks like it is done.
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
Review comment:
```suggestion
* checkpoint the current {@link KafkaSourceDescriptor} and move to process the next element. These
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -1051,33 +1198,352 @@ public void populateDisplayData(DisplayData.Builder builder) {
}
}
- ////////////////////////////////////////////////////////////////////////////////////////////////
-
- private static final Logger LOG = LoggerFactory.getLogger(KafkaIO.class);
-
/**
- * Returns a new config map which is merge of current config and updates. Verifies the updates do
- * not includes ignored properties.
+ * A {@link PTransform} to read from Kafka. See {@link KafkaIO} for more information on usage and
+ * configuration.
*/
- private static Map<String, Object> updateKafkaProperties(
- Map<String, Object> currentConfig,
- Map<String, String> ignoredProperties,
- Map<String, Object> updates) {
+ @Experimental(Kind.PORTABILITY)
+ @AutoValue
+ public abstract static class ReadSourceDescriptors<K, V>
+ extends PTransform<PCollection<KafkaSourceDescriptor>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadSourceDescriptors.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract ReadSourceDescriptors.Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract ReadSourceDescriptors.Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setCommitOffsetEnabled(
+ boolean commitOffsetEnabled);
+
+ abstract ReadSourceDescriptors.Builder<K, V> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadSourceDescriptors<K, V> build();
+ }
+
+ public static <K, V> ReadSourceDescriptors<K, V> read() {
+ return new AutoValue_KafkaIO_ReadSourceDescriptors.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the
+ // element
+ // will override the bootstrapServers from the config.
+ public ReadSourceDescriptors<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadSourceDescriptors<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimator<Instant>> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(
+ ReadSourceDescriptors.ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadSourceDescriptors<K, V> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new WallTime(state);
+ });
+ }
+
+ public ReadSourceDescriptors<K, V> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadSourceDescriptors<K, V> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the
+ // default value.
+ public ReadSourceDescriptors<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadSourceDescriptors<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadSourceDescriptors<K, V> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ // TODO(BEAM-10320): Create external build transform for ReadSourceDescriptors().
+ ReadAllFromRow forExternalBuild() {
+ return new ReadAllFromRow(this);
+ }
+
+ // This transform is used in cross-language case. The input Row should be encoded with an
+ // equivalent schema as KafkaSourceDescriptor.
+ private static class ReadAllFromRow<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadSourceDescriptors<K, V> readViaSDF;
+
+ ReadAllFromRow(ReadSourceDescriptors read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(Convert.fromRows(KafkaSourceDescriptor.class))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadSourceDescriptors<K, V> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
- for (String key : updates.keySet()) {
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescriptor> input) {
checkArgument(
- !ignoredProperties.containsKey(key),
- "No need to configure '%s'. %s",
- key,
- ignoredProperties.get(key));
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadSourceDescriptors can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescriptor during runtime. Otherwise, the pipeline will fail.");
Review comment:
```suggestion
"The bootstrapServers is not set. It must be populated through the KafkaSourceDescriptor during runtime otherwise the pipeline will fail.");
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescriptor }. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn#restrictionTracker(KafkaSourceDescriptor, OffsetRange)} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescriptor, OffsetRange).}
+ * A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * <p>The {@link WatermarkEstimator} is created by {@link
+ * ReadSourceDescriptors#getCreateWatermarkEstimatorFn()}. The estimated watermark is computed by
+ * this {@link WatermarkEstimator} based on output timestamps computed by {@link
+ * ReadSourceDescriptors#getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link ReadSourceDescriptors#withProcessingTime()} as {@code
Review comment:
```suggestion
* configuration is using {@link ReadSourceDescriptors#withProcessingTime()} as the {@code
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescriptor }. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn#restrictionTracker(KafkaSourceDescriptor, OffsetRange)} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescriptor, OffsetRange).}
+ * A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * <p>The {@link WatermarkEstimator} is created by {@link
+ * ReadSourceDescriptors#getCreateWatermarkEstimatorFn()}. The estimated watermark is computed by
+ * this {@link WatermarkEstimator} based on output timestamps computed by {@link
+ * ReadSourceDescriptors#getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link ReadSourceDescriptors#withProcessingTime()} as {@code
+ * extractTimestampFn} and {@link
+ * ReadSourceDescriptors#withMonotonicallyIncreasingWatermarkEstimator()} as {@link
+ * WatermarkEstimator}.
+ */
+@UnboundedPerElement
+class ReadFromKafkaDoFn<K, V> extends DoFn<KafkaSourceDescriptor, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadSourceDescriptors transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaDoFn.class);
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Valid between bundle start and bundle finish.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient LoadingCache<TopicPartition, AverageRecordSize> avgRecordSize;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memoizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memoizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memoizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescriptor kafkaSourceDescriptor) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescriptor);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescriptor.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescriptor.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescriptor.getStartReadOffset();
+ } else if (kafkaSourceDescriptor.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescriptor.getTopicPartition(),
+ kafkaSourceDescriptor.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescriptor.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimator<Instant> newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescriptor kafkaSourceDescriptor, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numRecords =
+ restrictionTracker(kafkaSourceDescriptor, offsetRange).getProgress().getWorkRemaining();
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (!avgRecordSize.asMap().containsKey(kafkaSourceDescriptor.getTopicPartition())) {
+ return numRecords;
+ }
+ return avgRecordSize.get(kafkaSourceDescriptor.getTopicPartition()).getTotalSize(numRecords);
+ }
+
+ @NewTracker
+ public GrowableOffsetRangeTracker restrictionTracker(
+ @Element KafkaSourceDescriptor kafkaSourceDescriptor, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescriptor);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescriptor.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescriptor.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescriptor kafkaSourceDescriptor,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescriptor);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescriptor.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescriptor.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescriptor.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there are no records available for the current TopicPartition, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize
+ .getUnchecked(kafkaSourceDescriptor.getTopicPartition())
+ .update(recordSize, rawRecord.offset() - expectedOffset);
+ expectedOffset = rawRecord.offset() + 1;
+ Instant outputTimestamp;
+ // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+ // WatermarkEstimator should be a Manual one.
Review comment:
```suggestion
// WatermarkEstimator should be a manual one.
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescriptor }. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn#restrictionTracker(KafkaSourceDescriptor, OffsetRange)} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescriptor, OffsetRange).}
+ * A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * <p>The {@link WatermarkEstimator} is created by {@link
+ * ReadSourceDescriptors#getCreateWatermarkEstimatorFn()}. The estimated watermark is computed by
+ * this {@link WatermarkEstimator} based on output timestamps computed by {@link
+ * ReadSourceDescriptors#getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link ReadSourceDescriptors#withProcessingTime()} as {@code
+ * extractTimestampFn} and {@link
+ * ReadSourceDescriptors#withMonotonicallyIncreasingWatermarkEstimator()} as {@link
+ * WatermarkEstimator}.
+ */
+@UnboundedPerElement
+class ReadFromKafkaDoFn<K, V> extends DoFn<KafkaSourceDescriptor, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadSourceDescriptors transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadFromKafkaDoFn.class);
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimator<Instant>>
+ createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Valid between bundle start and bundle finish.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient LoadingCache<TopicPartition, AverageRecordSize> avgRecordSize;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memoizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memoizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memoizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescriptor kafkaSourceDescriptor) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescriptor);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescriptor.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescriptor.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescriptor.getStartReadOffset();
+ } else if (kafkaSourceDescriptor.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescriptor.getTopicPartition(),
+ kafkaSourceDescriptor.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescriptor.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimator<Instant> newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescriptor kafkaSourceDescriptor, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numRecords =
+ restrictionTracker(kafkaSourceDescriptor, offsetRange).getProgress().getWorkRemaining();
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (!avgRecordSize.asMap().containsKey(kafkaSourceDescriptor.getTopicPartition())) {
+ return numRecords;
+ }
+ return avgRecordSize.get(kafkaSourceDescriptor.getTopicPartition()).getTotalSize(numRecords);
+ }
+
+ @NewTracker
+ public GrowableOffsetRangeTracker restrictionTracker(
+ @Element KafkaSourceDescriptor kafkaSourceDescriptor, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescriptor);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescriptor.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescriptor.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescriptor kafkaSourceDescriptor,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescriptor);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescriptor.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescriptor.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescriptor.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there are no records available for the current TopicPartition, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize
+ .getUnchecked(kafkaSourceDescriptor.getTopicPartition())
+ .update(recordSize, rawRecord.offset() - expectedOffset);
+ expectedOffset = rawRecord.offset() + 1;
+ Instant outputTimestamp;
+ // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+ // WatermarkEstimator should be a Manual one.
+ if (timestampPolicy != null) {
+ checkState(watermarkEstimator instanceof ManualWatermarkEstimator);
+ TimestampPolicyContext context =
+ new TimestampPolicyContext(
+ (long) ((HasProgress) tracker).getProgress().getWorkRemaining(), Instant.now());
+ outputTimestamp = timestampPolicy.getTimestampForRecord(context, kafkaRecord);
+ ((ManualWatermarkEstimator) watermarkEstimator)
+ .setWatermark(timestampPolicy.getWatermark(context));
+ } else {
+ outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ }
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize =
+ CacheBuilder.newBuilder()
+ .maximumSize(1000L)
+ .build(
+ new CacheLoader<TopicPartition, AverageRecordSize>() {
+ @Override
+ public AverageRecordSize load(TopicPartition topicPartition) throws Exception {
+ return new AverageRecordSize();
+ }
+ });
+ consumerSpEL = new ConsumerSpEL();
+ keyDeserializerInstance = keyDeserializerProvider.getDeserializer(consumerConfig, true);
+ valueDeserializerInstance = valueDeserializerProvider.getDeserializer(consumerConfig, false);
+ }
+
+ @Teardown
+ public void teardown() throws Exception {
+ try {
+ Closeables.close(keyDeserializerInstance, true);
+ Closeables.close(valueDeserializerInstance, true);
+ } catch (Exception anyException) {
+ LOG.warn("Fail to close resource during finishing bundle: {}", anyException.getMessage());
Review comment:
slf4j has special logic to format exceptions.
```suggestion
LOG.warn("Fail to close resource during finishing bundle.", anyException);
```
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
##########
@@ -0,0 +1,403 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescriptor} and outputs {@link KafkaRecord}.
+ * By default, a {@link MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ *
+ * <p>{@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescriptor}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initial Restriction</h4>
+ *
+ * <p>The initial range for a {@link KafkaSourceDescriptor } is defined by {@code [startOffset,
+ * Long.MAX_VALUE)} where {@code startOffset} is defined as:
+ *
+ * <ul>
+ * <li>the {@code startReadOffset} if {@link KafkaSourceDescriptor#getStartReadOffset} is set.
+ * <li>the first offset with a greater or equivalent timestamp if {@link
+ * KafkaSourceDescriptor#getStartReadTime()} is set.
+ * <li>the {@code last committed offset + 1} for the {@link Consumer#position(TopicPartition)
+ * topic partition}.
+ * </ul>
+ *
+ * <h4>Splitting</h4>
+ *
+ * <p>TODO(BEAM-10319): Add support for initial splitting.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescriptor } and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescriptor }. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn#restrictionTracker(KafkaSourceDescriptor, OffsetRange)} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescriptor, OffsetRange).}
+ * A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * <p>The {@link WatermarkEstimator} is created by {@link
+ * ReadSourceDescriptors#getCreateWatermarkEstimatorFn()}. The estimated watermark is computed by
+ * this {@link WatermarkEstimator} based on output timestamps computed by {@link
+ * ReadSourceDescriptors#getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link ReadSourceDescriptors#withProcessingTime()} as {@code
+ * extractTimestampFn} and {@link
+ * ReadSourceDescriptors#withMonotonicallyIncreasingWatermarkEstimator()} as {@link
Review comment:
```suggestion
* ReadSourceDescriptors#withMonotonicallyIncreasingWatermarkEstimator()} as the {@link
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r445226123
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSourceDescription.java
##########
@@ -0,0 +1,92 @@
+/*
+ * 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.sdk.io.kafka;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.kafka.common.TopicPartition;
+import org.joda.time.Instant;
+
+/**
+ * An AutoValue object which represents a Kafka source description. Note that this object should be
+ * encoded/decoded with equivalent {@link Schema} as a {@link Row} when crossing the wire.
+ */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class KafkaSourceDescription implements Serializable {
+ @SchemaFieldName("topic")
+ abstract String getTopic();
+
+ @SchemaFieldName("partition")
+ abstract Integer getPartition();
+
+ @SchemaFieldName("start_read_offset")
+ @Nullable
+ abstract Long getStartReadOffset();
+
+ @SchemaFieldName("start_read_time")
+ @Nullable
+ abstract Instant getStartReadTime();
+
+ @SchemaFieldName("bootstrapServers")
+ @Nullable
+ abstract List<String> getBootStrapServers();
+
+ private TopicPartition topicPartition = null;
+
+ public TopicPartition getTopicPartition() {
Review comment:
Thanks!
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637686542
> @boyuanzz Thanks! I'll take a look asap.
Thanks! Also I'm curious is there any runner run KafkaIOIT over fnapi? It seems like the KafkaIOIT is only used in performance test?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637719079
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r443059453
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
Review comment:
I think we should remove it and update the DoFn documentation stating the default clearly.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637676591
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442980355
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
Review comment:
It seems like there is no difference between poll ->resume and resume -> poll? For example, a bad situation could be there is only one `TopicPartition` and there is always no available records. In this case, both poll ->resume and resume -> poll will poll the records with a time interval of api timeout + residual reschedule time.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442963985
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
+ throws Exception {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ double numOfRecords = 0.0;
+ if (offsetRange.getTo() != Long.MAX_VALUE) {
+ numOfRecords = (new OffsetRangeTracker(offsetRange)).getProgress().getWorkRemaining();
+ } else {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetEstimator =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "size-" + topicPartition.toString(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ topicPartition);
+ numOfRecords =
+ (new GrowableOffsetRangeTracker(offsetRange.getFrom(), offsetEstimator))
+ .getProgress()
+ .getWorkRemaining();
+ }
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap != null) {
+ numOfRecords = numOfRecords / (1 + avgOffsetGap.get());
+ }
+ return (avgRecordSize == null ? 1 : avgRecordSize.get()) * numOfRecords;
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element Row kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange,
+ OutputReceiver<OffsetRange> receiver)
+ throws Exception {
+ receiver.output(offsetRange);
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element Row kafkaSourceDescription, @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ if (restriction.getTo() == Long.MAX_VALUE) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + topicPartition, offsetConsumerConfig, updatedConsumerConfig)),
+ topicPartition);
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element Row kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(consumer, ImmutableList.of(topicPartition));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(topicPartition, startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(kafkaPollTimeout.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ Instant outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgRecordSize.update(recordSize);
+ avgOffsetGap.update(expectedOffset - rawRecord.offset());
+ expectedOffset = rawRecord.offset() + 1;
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
+ throw anyException;
+ }
+ }
+ }
+
+ @GetRestrictionCoder
+ public Coder<OffsetRange> restrictionCoder() {
+ return new OffsetRange.Coder();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ // Start to track record size and offset gap per bundle.
+ avgRecordSize = new KafkaIOUtils.MovingAvg();
Review comment:
The UnboundedSource implementation tracked it for a specific TopicPartition and it didn't apply to multiple topics. When we are processing multiple restrictions we'll conflate this information from all of them.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
lukecwik commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442962378
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link Row} IN {@link
+ * KafkaSourceDescriptionSchemas} which represents Kafka source description as input and outputs a
+ * PCollection of {@link KafkaRecord}. The core implementation is based on {@code SplittableDoFn}.
+ * For more details about the concept of {@code SplittableDoFn}, please refer to the beam blog post:
+ * https://beam.apache.org/blog/splittable-do-fn/ and design doc:https://s.apache.org/beam-fn-api.
+ * The major difference from {@link KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source
+ * descriptions(e.g., {@link KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()},
+ * {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead,
+ * the pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from Kafka source description in {@link Row}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * Row}, and the restriction is an {@link OffsetRange} which represents record offset. A {@link
+ * GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with {@code
+ * Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(Row)} creates an initial range for a input element
+ * {@link Row}. The end of range will be initialized as {@code Long.MAX_VALUE}. For the start of the
+ * range:
+ *
+ * <ul>
+ * <li>If {@code start_read_offset} in {@link Row} is set, use this offset as start.
+ * <li>If {@code start_read_time} in {@link Row} is set, seek the start offset based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link Row} and move to process the next element. These deferred elements
+ * will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link Row}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in the
+ * {@link GrowableOffsetRangeTracker} as the {@link GrowableOffsetRangeTracker.RangeEndEstimator} to
+ * poll the latest offset. Please refer to {@link ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for
+ * details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(Row, OffsetRange).} A {@link
+ * KafkaIOUtils.MovingAvg} is used to track the average size of kafka records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link Row}.
+ */
+@AutoValue
+public abstract class ReadViaSDF<K, V>
+ extends PTransform<PCollection<Row>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ abstract Schema getKafkaSourceDescriptionSchema();
+
+ abstract Builder<K, V> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V> {
+ abstract Builder<K, V> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V> setOffsetConsumerConfig(Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V> setKeyDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setValueDeserializerProvider(DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V> setKafkaSourceDescriptionSchema(Schema schema);
+
+ abstract ReadViaSDF<K, V> build();
+ }
+
+ static class KafkaSourceDescriptionSchemas {
+ static final String TOPIC = "topic";
+ static final String PARTITION = "partition";
+ static final String START_READ_OFFSET = "start_read_offset";
+ static final String START_READ_TIME = "start_read_time";
+ static final String BOOTSTRAP_SERVERS = "bootstrap_servers";
+
+ static Schema getSchema() {
+ return Schema.builder()
+ .addStringField(TOPIC)
+ .addInt32Field(PARTITION)
+ .addNullableField(START_READ_OFFSET, FieldType.INT32)
+ .addNullableField(START_READ_TIME, FieldType.INT64)
+ .addNullableField(BOOTSTRAP_SERVERS, FieldType.array(FieldType.STRING))
+ .build();
+ }
+ }
+
+ public static <K, V> ReadViaSDF<K, V> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime())
+ .setCommitOffsetEnabled(false)
+ .setKafkaSourceDescriptionSchema(KafkaSourceDescriptionSchemas.getSchema())
+ .build();
+ }
+
+ public ReadViaSDF<K, V> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializer(Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigUpdates(Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(
+ getConsumerConfig(), KafkaIOUtils.IGNORED_CONSUMER_PROPERTIES, configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V> withConsumerConfigOverrides(Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<Row> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input.apply(ParDo.of(new ReadFromKafkaDoFn())).setCoder(outputCoder);
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ LOG.warn("Offset committed is not supported yet. Ignore the value.");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link Row} in {@link KafkaSourceDescriptionSchemas} which
+ * represents a Kafka source description and outputs {@link KafkaRecord}. By default, a {@link
+ * MonotonicallyIncreasing} watermark estimator is used to track watermark.
+ */
+ @VisibleForTesting
+ class ReadFromKafkaDoFn extends DoFn<Row, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn() {}
+
+ private final Map<String, Object> consumerConfig = ReadViaSDF.this.getConsumerConfig();
+
+ private final Map<String, Object> offsetConsumerConfig =
+ ReadViaSDF.this.getOffsetConsumerConfig();
+
+ private final DeserializerProvider keyDeserializerProvider =
+ ReadViaSDF.this.getKeyDeserializerProvider();
+ private final DeserializerProvider valueDeserializerProvider =
+ ReadViaSDF.this.getValueDeserializerProvider();
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn = ReadViaSDF.this.getConsumerFactoryFn();
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn =
+ ReadViaSDF.this.getExtractOutputTimestampFn();
+
+ private final Duration kafkaPollTimeout = Duration.millis(1000);
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient KafkaIOUtils.MovingAvg avgRecordSize = null;
+ private transient KafkaIOUtils.MovingAvg avgOffsetGap = null;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element Row kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ kafkaSourceDescription.getString(KafkaSourceDescriptionSchemas.TOPIC),
+ kafkaSourceDescription.getInt32(KafkaSourceDescriptionSchemas.PARTITION));
+ consumerSpEL.evaluateAssign(offsetConsumer, ImmutableList.of(topicPartition));
+ long startOffset;
+ if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET)
+ != null) {
+ startOffset =
+ kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_OFFSET);
+ } else if (kafkaSourceDescription.getInt64(KafkaSourceDescriptionSchemas.START_READ_TIME)
+ != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ topicPartition,
+ Instant.ofEpochMilli(kafkaSourceDescription.getInt64("start_read_time")));
+ } else {
+ startOffset = offsetConsumer.position(topicPartition);
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public MonotonicallyIncreasing newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return new MonotonicallyIncreasing(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(@Element Row kafkaSourceDescription, @Restriction OffsetRange offsetRange)
Review comment:
My bad, totally meant to say `restrictionTracker()`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] iemejia commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
iemejia commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434666700
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
This seems strangely close to something we lived in the SDF version of HBaseIO. In the first version we did an artificial object called `HBaseQuery` that contained the minimum information we needed to be able to query the Data store in a SDF way, but then other requirements came in and we started to add extra parameters to end up with something that was almost close to the exact 'complete' specification of the Read class so we decided to switch to use a `PCollection<Read>` as input otherwise we will be duplicating code, so we ended up with https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java#L353
Here you can have `PCollection<Read>` as an input and get rid of `KafkaSourceDescription` and this will have a more consistent user experience for final users. Notice that this `ReadAll` like pattern is also now used in [SolrIO](https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/solr/src/main/java/org/apache/beam/sdk/io/solr/SolrIO.java#L501) and there is an ongoing PR to introduce it for [CassandraIO](https://github.com/apache/beam/pull/10546) so maybe it is a good idea we follow it for consistency.
Notice that in the SolrIO case the change looks even closer to this one because we ended up introducing `ReplicaInfo` (the spiritual equivalent of `TopicPartition`) into normal Read and we guarantee in expansion that this field gets filled if the users don't do it, but if they do well we asume they know what they are doing and we go with it.
Another advantage of having the full specification is that you will be able to read not only from multiple topics but also from different clusters because of the power of having the full `Read` specification,
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on pull request #11749:
URL: https://github.com/apache/beam/pull/11749#issuecomment-637806753
Run Java PreCommit
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] iemejia commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
iemejia commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r434666700
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
This seems strangely close to something we lived in the SDF version of HBaseIO. In the first version we did an artificial object called `HBaseQuery` that contained the minimum information we needed to be able to query the Data store in a SDF way, but then other requirements came in and we started to add extra parameters to end up with something that was almost close to the exact 'complete' specification of the Read class so we decided to switch to use a `PCollection<Read>` as input otherwise we will be duplicating code, so we ended up with https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java#L353
Here you can have `PCollection<Read>` as an input and get rid of `KafkaSourceDescription` if you move the missing parameters like `TopicPartition` into normal `Read` and this will have a more consistent user experience for final users. Notice that this `ReadAll` like pattern is also now used in [SolrIO](https://github.com/apache/beam/blob/f6ef9032f521180f1cc26959d9d6ab86dd37a13c/sdks/java/io/solr/src/main/java/org/apache/beam/sdk/io/solr/SolrIO.java#L501) and there is an ongoing PR to introduce it for [CassandraIO](https://github.com/apache/beam/pull/10546) so maybe it is a good idea we follow it for consistency.
Notice that in the SolrIO case the change looks even closer to this one because we ended up introducing `ReplicaInfo` (the spiritual equivalent of `TopicPartition`) into normal Read and we guarantee in expansion that this field gets filled if the users don't do it, but if they do well we asume they know what they are doing and we go with it.
Another advantage of having the full specification is that you will be able to read not only from multiple topics but also from different clusters because of the power of having the full `Read` specification,
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] aromanenko-dev commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
aromanenko-dev commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r439509913
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaViaSDF.java
##########
@@ -0,0 +1,697 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.Element;
+import org.apache.beam.sdk.transforms.DoFn.GetRestrictionCoder;
+import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the beam blog post: https://beam.apache.org/blog/splittable-do-fn/ and design
+ * doc:https://s.apache.org/beam-fn-api. The major difference from {@link KafkaIO.Read} is, {@link
+ * ReadFromKafkaViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadFromKafkaViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadFromKafkaViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyCoder()} is the same as {@link
+ * KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueCoder()} is the same as {@link
+ * KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadFromKafkaViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadFromKafkaViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadFromKafkaViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadFromKafkaViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#commitOffsets()} enables committing offset after processing the
+ * record. Note that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadFromKafkaViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadFromKafkaViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a
+ * function which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is
+ * used to produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadFromKafkaViaSDF#withProcessingTime()}, {@link ReadFromKafkaViaSDF#withCreateTime()} and
+ * {@link ReadFromKafkaViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadFromKafkaViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("my_topic", 1))))
+ * .apply(ReadFromKafkaViaSDF.create()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@link KafkaSourceDescription#getStartOffset()} is set, use this offset as start.
+ * <li>If {@link KafkaSourceDescription#getStartReadTime()} is set, seek the start offset based on
+ * this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or {@link OffsetRangeTracker}
+ * per {@link KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer}
+ * is used in the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The estimated watermark is computed by {@link MonotonicallyIncreasing} based on output timestamps
+ * per {@link KafkaSourceDescription}.
+ */
+@AutoValue
+public abstract class ReadFromKafkaViaSDF<K, V>
Review comment:
@boyuanzz It's quite tricky question. Well, initially we use the static list of bootstrap servers for message and offset consumers. So we expect them equal and I think it should be the same with SDF.
Message consumer is also used to fetch the topic(s) partitions for initial split. With SDF we don't need that.
What is not clear for me, is by "user emits different bootstrapServer dynamically" you mean that it will be set in a runtime? Is it going to be changed during a runtime or set only once?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444654119
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(expectedOffset - rawRecord.offset());
+ avgRecordSize
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(recordSize);
+ expectedOffset = rawRecord.offset() + 1;
+ Instant outputTimestamp;
+ // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+ // WatermarkEstimator should be a Manual one.
+ if (timestampPolicy != null) {
+ checkState(watermarkEstimator instanceof ManualWatermarkEstimator);
+ TimestampPolicyContext context =
+ new TimestampPolicyContext(
+ (long) ((HasProgress) tracker).getProgress().getWorkRemaining(),
+ Instant.now());
+ outputTimestamp = timestampPolicy.getTimestampForRecord(context, kafkaRecord);
+ ((ManualWatermarkEstimator) watermarkEstimator)
+ .setWatermark(timestampPolicy.getWatermark(context));
+ } else {
+ outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+ }
+ receiver.outputWithTimestamp(kafkaRecord, outputTimestamp);
+ }
+ }
+ } catch (Exception anyException) {
+ LOG.error("{}: Exception while reading from Kafka", this, anyException);
Review comment:
Nope, I just forgot that we have higher level logging.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r442566478
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
##########
@@ -795,6 +828,12 @@ public void setValueDeserializer(String valueDeserializer) {
return withWatermarkFn2(unwrapKafkaAndThen(watermarkFn));
}
+ /** A function to the compute output timestamp from a {@link KafkaRecord}. */
+ public Read<K, V> withExtractOutputTimestampFn(
Review comment:
> How is this different from withTimestampFn2?
`withTimestampFn2` has been deprecated in `KafkaIO.Read`. The major concern of reusing `withTimestampFn2` is, it will means differently under SDF and UnboundedSource, which causes confusion.
> Setting the top level properties allow us to say that this property is supported when used as an SDF.
This is what I want to do by having `withExtractOutputTimestampFn `
> Also, what prevents us from supporting TimestampPolicy? We should be able to call it and give it the three pieces of information it requests (message backlog / backlog check time / current kafka record).
The difficulty is the message backlog / backlog check time is not memorized per (element, restriction). With SDF framework, the backlog is retrieved by called `RestricitonTracker.getProgress()`, we cannot call it per element in order to extract timestamp.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [beam] boyuanzz commented on a change in pull request #11749: [BEAM-9977] Implement ReadFromKafkaViaSDF
Posted by GitBox <gi...@apache.org>.
boyuanzz commented on a change in pull request #11749:
URL: https://github.com/apache/beam/pull/11749#discussion_r444648850
##########
File path: sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadViaSDF.java
##########
@@ -0,0 +1,861 @@
+/*
+ * 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.sdk.io.kafka;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
+import org.apache.beam.sdk.io.kafka.KafkaSourceDescription.Schemas;
+import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.splittabledofn.GrowableOffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
+import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Closeables;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} that takes a PCollection of {@link KafkaSourceDescription} as input and
+ * outputs a PCollection of {@link KafkaRecord}. The core implementation is based on {@code
+ * SplittableDoFn}. For more details about the concept of {@code SplittableDoFn}, please refer to
+ * the <a href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
+ * href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
+ * KafkaIO.Read} is, {@link ReadViaSDF} doesn't require source descriptions(e.g., {@link
+ * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
+ * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
+ * pipeline can populate these source descriptions during runtime. For example, the pipeline can
+ * query Kafka topics from BigQuery table and read these topics via {@link ReadViaSDF}.
+ *
+ * <h3>Common Kafka Consumer Configurations</h3>
+ *
+ * <p>Most Kafka consumer configurations are similar to {@link KafkaIO.Read}:
+ *
+ * <ul>
+ * <li>{@link ReadViaSDF#getConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getConsumerFactoryFn()} is the same as {@link
+ * KafkaIO.Read#getConsumerFactoryFn()}.
+ * <li>{@link ReadViaSDF#getOffsetConsumerConfig()} is the same as {@link
+ * KafkaIO.Read#getOffsetConsumerConfig()}.
+ * <li>{@link ReadViaSDF#getKeyCoder()} is the same as {@link KafkaIO.Read#getKeyCoder()}.
+ * <li>{@link ReadViaSDF#getValueCoder()} is the same as {@link KafkaIO.Read#getValueCoder()}.
+ * <li>{@link ReadViaSDF#getKeyDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getKeyDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#getValueDeserializerProvider()} is the same as {@link
+ * KafkaIO.Read#getValueDeserializerProvider()}.
+ * <li>{@link ReadViaSDF#isCommitOffsetEnabled()} means the same as {@link
+ * KafkaIO.Read#isCommitOffsetsInFinalizeEnabled()}.
+ * </ul>
+ *
+ * <p>For example, to create a basic {@link ReadViaSDF} transform:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(KafkaSourceDescription.of(new TopicPartition("topic", 1)))
+ * .apply(KafkaIO.readAll()
+ * .withBootstrapServers("broker_1:9092,broker_2:9092")
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * Note that the {@code bootstrapServers} can also be populated from {@link KafkaSourceDescription}:
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class));
+ *
+ * }</pre>
+ *
+ * <h3>Configurations of {@link ReadViaSDF}</h3>
+ *
+ * <p>Except configurations of Kafka Consumer, there are some other configurations which are related
+ * to processing records.
+ *
+ * <p>{@link ReadViaSDF#commitOffsets()} enables committing offset after processing the record. Note
+ * that if {@code isolation.level} is set to "read_committed" or {@link
+ * ConsumerConfig#ENABLE_AUTO_COMMIT_CONFIG} is set in the consumer config, the {@link
+ * ReadViaSDF#commitOffsets()} will be ignored.
+ *
+ * <p>{@link ReadViaSDF#withExtractOutputTimestampFn(SerializableFunction)} asks for a function
+ * which takes a {@link KafkaRecord} as input and outputs outputTimestamp. This function is used to
+ * produce output timestamp per {@link KafkaRecord}. There are three built-in types: {@link
+ * ReadViaSDF#withProcessingTime()}, {@link ReadViaSDF#withCreateTime()} and {@link
+ * ReadViaSDF#withLogAppendTime()}.
+ *
+ * <p>For example, to create a {@link ReadViaSDF} with these configurations:
+ *
+ * <pre>{@code
+ * pipeline
+ * .apply(Create.of(
+ * KafkaSourceDescription.of(
+ * new TopicPartition("topic", 1),
+ * null,
+ * null,
+ * ImmutableList.of("broker_1:9092", "broker_2:9092"))
+ * .apply(KafkaIO.readAll()
+ * .withKeyDeserializer(LongDeserializer.class).
+ * .withValueDeserializer(StringDeserializer.class)
+ * .withProcessingTime()
+ * .commitOffsets());
+ *
+ * }</pre>
+ *
+ * <h3>Read from {@link KafkaSourceDescription}</h3>
+ *
+ * {@link ReadFromKafkaDoFn} implements the logic of reading from Kafka. The element is a {@link
+ * KafkaSourceDescription}, and the restriction is an {@link OffsetRange} which represents record
+ * offset. A {@link GrowableOffsetRangeTracker} is used to track an {@link OffsetRange} ended with
+ * {@code Long.MAX_VALUE}. For a finite range, a {@link OffsetRangeTracker} is created.
+ *
+ * <h4>Initialize Restriction</h4>
+ *
+ * {@link ReadFromKafkaDoFn#initialRestriction(KafkaSourceDescription)} creates an initial range for
+ * a input element {@link KafkaSourceDescription}. The end of range will be initialized as {@code
+ * Long.MAX_VALUE}. For the start of the range:
+ *
+ * <ul>
+ * <li>If {@code startReadOffset} in {@link KafkaSourceDescription} is set, use this offset as
+ * start.
+ * <li>If {@code startReadTime} in {@link KafkaSourceDescription} is set, seek the start offset
+ * based on this time.
+ * <li>Otherwise, the last committed offset + 1 will be returned by {@link
+ * Consumer#position(TopicPartition)} as the start.
+ * </ul>
+ *
+ * <h4>Initial Split</h4>
+ *
+ * <p>There is no initial split for now.
+ *
+ * <h4>Checkpoint and Resume Processing</h4>
+ *
+ * <p>There are 2 types of checkpoint here: self-checkpoint which invokes by the DoFn and
+ * system-checkpoint which is issued by the runner via {@link
+ * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest}. Every time the
+ * consumer gets empty response from {@link Consumer#poll(long)}, {@link ReadFromKafkaDoFn} will
+ * checkpoint at current {@link KafkaSourceDescription} and move to process the next element. These
+ * deferred elements will be resumed by the runner as soon as possible.
+ *
+ * <h4>Progress and Size</h4>
+ *
+ * <p>The progress is provided by {@link GrowableOffsetRangeTracker} or per {@link
+ * KafkaSourceDescription}. For an infinite {@link OffsetRange}, a Kafka {@link Consumer} is used in
+ * the {@link GrowableOffsetRangeTracker} as the {@link
+ * GrowableOffsetRangeTracker.RangeEndEstimator} to poll the latest offset. Please refer to {@link
+ * ReadFromKafkaDoFn.KafkaLatestOffsetEstimator} for details.
+ *
+ * <p>The size is computed by {@link ReadFromKafkaDoFn#getSize(KafkaSourceDescription,
+ * OffsetRange).} A {@link KafkaIOUtils.MovingAvg} is used to track the average size of kafka
+ * records.
+ *
+ * <h4>Track Watermark</h4>
+ *
+ * The {@link WatermarkEstimator} is created by {@link #getCreateWatermarkEstimatorFn()}. The
+ * estimated watermark is computed by this {@link WatermarkEstimator} based on output timestamps
+ * computed by {@link #getExtractOutputTimestampFn()} (SerializableFunction)}. The default
+ * configuration is using {@link #withProcessingTime()} as {@code extractTimestampFn} and {@link
+ * #withMonotonicallyIncreasingWatermarkEstimator()} as {@link WatermarkEstimator}.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoValue
+abstract class ReadViaSDF<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<KafkaSourceDescription>, PCollection<KafkaRecord<K, V>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadViaSDF.class);
+
+ abstract Map<String, Object> getConsumerConfig();
+
+ @Nullable
+ abstract Map<String, Object> getOffsetConsumerConfig();
+
+ @Nullable
+ abstract DeserializerProvider getKeyDeserializerProvider();
+
+ @Nullable
+ abstract DeserializerProvider getValueDeserializerProvider();
+
+ @Nullable
+ abstract Coder<K> getKeyCoder();
+
+ @Nullable
+ abstract Coder<V> getValueCoder();
+
+ abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ getConsumerFactoryFn();
+
+ @Nullable
+ abstract SerializableFunction<KafkaRecord<K, V>, Instant> getExtractOutputTimestampFn();
+
+ @Nullable
+ abstract SerializableFunction<Instant, WatermarkEstimatorT> getCreateWatermarkEstimatorFn();
+
+ abstract boolean isCommitOffsetEnabled();
+
+ @Nullable
+ abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
+
+ abstract Builder<K, V, WatermarkEstimatorT> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>> {
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerConfig(Map<String, Object> config);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setOffsetConsumerConfig(
+ Map<String, Object> offsetConsumerConfig);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueDeserializerProvider(
+ DeserializerProvider deserializerProvider);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setKeyCoder(Coder<K> keyCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setValueCoder(Coder<V> valueCoder);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCreateWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setCommitOffsetEnabled(boolean commitOffsetEnabled);
+
+ abstract Builder<K, V, WatermarkEstimatorT> setTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> policy);
+
+ abstract ReadViaSDF<K, V, WatermarkEstimatorT> build();
+ }
+
+ public static <K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ ReadViaSDF<K, V, WatermarkEstimatorT> read() {
+ return new AutoValue_ReadViaSDF.Builder<K, V, WatermarkEstimatorT>()
+ .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
+ .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
+ .setCommitOffsetEnabled(false)
+ .build()
+ .withProcessingTime()
+ .withMonotonicallyIncreasingWatermarkEstimator();
+ }
+
+ // Note that if the bootstrapServers is set here but also populated with the element, the element
+ // will override the bootstrapServers from the config.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withBootstrapServers(String bootstrapServers) {
+ return withConsumerConfigUpdates(
+ ImmutableMap.of(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerProvider(
+ DeserializerProvider<K> deserializerProvider) {
+ return toBuilder().setKeyDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerProvider(
+ DeserializerProvider<V> deserializerProvider) {
+ return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializer(
+ Class<? extends Deserializer<K>> keyDeserializer) {
+ return withKeyDeserializerProvider(LocalDeserializerProvider.of(keyDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializer(
+ Class<? extends Deserializer<V>> valueDeserializer) {
+ return withValueDeserializerProvider(LocalDeserializerProvider.of(valueDeserializer));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withKeyDeserializerAndCoder(
+ Class<? extends Deserializer<K>> keyDeserializer, Coder<K> keyCoder) {
+ return withKeyDeserializer(keyDeserializer).toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withValueDeserializerAndCoder(
+ Class<? extends Deserializer<V>> valueDeserializer, Coder<V> valueCoder) {
+ return withValueDeserializer(valueDeserializer).toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigUpdates(
+ Map<String, Object> configUpdates) {
+ Map<String, Object> config =
+ KafkaIOUtils.updateKafkaProperties(getConsumerConfig(), configUpdates);
+ return toBuilder().setConsumerConfig(config).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withExtractOutputTimestampFn(
+ SerializableFunction<KafkaRecord<K, V>, Instant> fn) {
+ return toBuilder().setExtractOutputTimestampFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreatWatermarkEstimatorFn(
+ SerializableFunction<Instant, WatermarkEstimatorT> fn) {
+ return toBuilder().setCreateWatermarkEstimatorFn(fn).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withLogAppendTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useLogAppendTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withProcessingTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useProcessingTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withCreateTime() {
+ return withExtractOutputTimestampFn(ExtractOutputTimestampFns.useCreateTime());
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withWallTimeWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new WallTime(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withMonotonicallyIncreasingWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new MonotonicallyIncreasing(state);
+ });
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withManualWatermarkEstimator() {
+ return withCreatWatermarkEstimatorFn(
+ state -> {
+ return (WatermarkEstimatorT) new Manual(state);
+ });
+ }
+
+ // If a transactional producer is used and it's desired to only read records from committed
+ // transaction, it's recommended to set read_committed. Otherwise, read_uncommitted is the default
+ // value.
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withReadCommitted() {
+ return withConsumerConfigUpdates(ImmutableMap.of("isolation.level", "read_committed"));
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> commitOffsets() {
+ return toBuilder().setCommitOffsetEnabled(true).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withOffsetConsumerConfigOverrides(
+ Map<String, Object> offsetConsumerConfig) {
+ return toBuilder().setOffsetConsumerConfig(offsetConsumerConfig).build();
+ }
+
+ public ReadViaSDF<K, V, WatermarkEstimatorT> withConsumerConfigOverrides(
+ Map<String, Object> consumerConfig) {
+ return toBuilder().setConsumerConfig(consumerConfig).build();
+ }
+
+ ReadViaSDFExternally forExternalBuild() {
+ return new ReadViaSDFExternally(this);
+ }
+
+ private static class ReadViaSDFExternally<
+ K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends PTransform<PCollection<Row>, PCollection<KV<K, V>>> {
+
+ private final ReadViaSDF<K, V, WatermarkEstimatorT> readViaSDF;
+
+ ReadViaSDFExternally(ReadViaSDF read) {
+ readViaSDF = read;
+ }
+
+ @Override
+ public PCollection<KV<K, V>> expand(PCollection<Row> input) {
+ return input
+ .apply(
+ ParDo.of(
+ new DoFn<Row, KafkaSourceDescription>() {
+ @ProcessElement
+ public void processElement(
+ @Element Row row, OutputReceiver<KafkaSourceDescription> outputReceiver) {
+ TopicPartition topicPartition =
+ new TopicPartition(
+ row.getString(Schemas.TOPIC), row.getInt32(Schemas.PARTITION));
+ Instant startReadTime =
+ row.getInt64(Schemas.START_READ_TIME) != null
+ ? Instant.ofEpochMilli(row.getInt64(Schemas.START_READ_TIME))
+ : null;
+ outputReceiver.output(
+ KafkaSourceDescription.of(
+ topicPartition,
+ row.getInt64(Schemas.START_READ_OFFSET),
+ startReadTime,
+ new ArrayList<>(row.getArray(Schemas.BOOTSTRAP_SERVERS))));
+ }
+ }))
+ .apply(readViaSDF)
+ .apply(
+ ParDo.of(
+ new DoFn<KafkaRecord<K, V>, KV<K, V>>() {
+ @ProcessElement
+ public void processElement(
+ @Element KafkaRecord element, OutputReceiver<KV<K, V>> outputReceiver) {
+ outputReceiver.output(element.getKV());
+ }
+ }))
+ .setCoder(KvCoder.<K, V>of(readViaSDF.getKeyCoder(), readViaSDF.getValueCoder()));
+ }
+ }
+
+ ReadViaSDF<K, V, WatermarkEstimatorT> withTimestampPolicyFactory(
+ TimestampPolicyFactory<K, V> timestampPolicyFactory) {
+ return toBuilder().setTimestampPolicyFactory(timestampPolicyFactory).build();
+ }
+
+ @Override
+ public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaSourceDescription> input) {
+ checkArgument(
+ ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
+ "The ReadFromKafkaViaSDF can only used when beam_fn_api is enabled.");
+
+ checkArgument(getKeyDeserializerProvider() != null, "withKeyDeserializer() is required");
+ checkArgument(getValueDeserializerProvider() != null, "withValueDeserializer() is required");
+
+ ConsumerSpEL consumerSpEL = new ConsumerSpEL();
+ if (!consumerSpEL.hasOffsetsForTimes()) {
+ LOG.warn(
+ "Kafka client version {} is too old. Versions before 0.10.1.0 are deprecated and "
+ + "may not be supported in next release of Apache Beam. "
+ + "Please upgrade your Kafka client version.",
+ AppInfoParser.getVersion());
+ }
+
+ if (isCommitOffsetEnabled()) {
+ if (configuredKafkaCommit()) {
+ LOG.info(
+ "Either read_committed or auto_commit is set together with commitOffsetEnabled but you "
+ + "only need one of them. The commitOffsetEnabled is going to be ignored");
+ }
+ }
+
+ if (getConsumerConfig().get(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG) == null) {
+ LOG.warn(
+ "The bootstrapServers is not set. Then it must be populated through KafkaSourceDescription during runtime. Otherwise, the pipeline will fail.");
+ }
+
+ CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
+ Coder<K> keyCoder = getKeyCoder(coderRegistry);
+ Coder<V> valueCoder = getValueCoder(coderRegistry);
+ Coder<KafkaRecord<K, V>> outputCoder = KafkaRecordCoder.of(keyCoder, valueCoder);
+ PCollection<KafkaRecord<K, V>> output =
+ input
+ .apply(ParDo.of(new ReadFromKafkaDoFn<K, V, WatermarkEstimatorT>(this)))
+ .setCoder(outputCoder);
+ // TODO(BEAM-10123): Add CommitOffsetTransform to expansion.
+ if (isCommitOffsetEnabled() && !configuredKafkaCommit()) {
+ throw new IllegalStateException("Offset committed is not supported yet");
+ }
+ return output;
+ }
+
+ private Coder<K> getKeyCoder(CoderRegistry coderRegistry) {
+ return (getKeyCoder() != null)
+ ? getKeyCoder()
+ : getKeyDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private Coder<V> getValueCoder(CoderRegistry coderRegistry) {
+ return (getValueCoder() != null)
+ ? getValueCoder()
+ : getValueDeserializerProvider().getCoder(coderRegistry);
+ }
+
+ private boolean configuredKafkaCommit() {
+ return getConsumerConfig().get("isolation.level") == "read_committed"
+ || Boolean.TRUE.equals(getConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG));
+ }
+
+ static class ExtractOutputTimestampFns<K, V> {
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useProcessingTime() {
+ return record -> Instant.now();
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useCreateTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.CREATE_TIME,
+ "Kafka record's timestamp is not 'CREATE_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+
+ public static <K, V> SerializableFunction<KafkaRecord<K, V>, Instant> useLogAppendTime() {
+ return record -> {
+ checkArgument(
+ record.getTimestampType() == KafkaTimestampType.LOG_APPEND_TIME,
+ "Kafka record's timestamp is not 'LOG_APPEND_TIME' "
+ + "(topic: %s, partition %s, offset %s, timestamp type '%s')",
+ record.getTopic(),
+ record.getPartition(),
+ record.getOffset(),
+ record.getTimestampType());
+ return new Instant(record.getTimestamp());
+ };
+ }
+ }
+
+ /**
+ * A SplittableDoFn which reads from {@link KafkaSourceDescription} and outputs {@link
+ * KafkaRecord}. By default, a {@link MonotonicallyIncreasing} watermark estimator is used to
+ * track watermark.
+ */
+ @VisibleForTesting
+ @UnboundedPerElement
+ static class ReadFromKafkaDoFn<K, V, WatermarkEstimatorT extends WatermarkEstimator<Instant>>
+ extends DoFn<KafkaSourceDescription, KafkaRecord<K, V>> {
+
+ ReadFromKafkaDoFn(ReadViaSDF transform) {
+ this.consumerConfig = transform.getConsumerConfig();
+ this.offsetConsumerConfig = transform.getOffsetConsumerConfig();
+ this.keyDeserializerProvider = transform.getKeyDeserializerProvider();
+ this.valueDeserializerProvider = transform.getValueDeserializerProvider();
+ this.consumerFactoryFn = transform.getConsumerFactoryFn();
+ this.extractOutputTimestampFn = transform.getExtractOutputTimestampFn();
+ this.createWatermarkEstimatorFn = transform.getCreateWatermarkEstimatorFn();
+ this.timestampPolicyFactory = transform.getTimestampPolicyFactory();
+ }
+
+ private final Map<String, Object> offsetConsumerConfig;
+
+ private final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
+ consumerFactoryFn;
+ private final SerializableFunction<KafkaRecord<K, V>, Instant> extractOutputTimestampFn;
+ private final SerializableFunction<Instant, WatermarkEstimatorT> createWatermarkEstimatorFn;
+ private final TimestampPolicyFactory<K, V> timestampPolicyFactory;
+
+ // Variables that are initialized when bundle is started and closed when FinishBundle is called.
+ private transient ConsumerSpEL consumerSpEL = null;
+ private transient Deserializer<K> keyDeserializerInstance = null;
+ private transient Deserializer<V> valueDeserializerInstance = null;
+
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgRecordSize;
+ private transient HashMap<TopicPartition, KafkaIOUtils.MovingAvg> avgOffsetGap;
+
+ private static final Duration KAFKA_POLL_TIMEOUT = Duration.millis(1000);
+
+ @VisibleForTesting final DeserializerProvider keyDeserializerProvider;
+ @VisibleForTesting final DeserializerProvider valueDeserializerProvider;
+ @VisibleForTesting final Map<String, Object> consumerConfig;
+
+ /**
+ * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
+ * fetch backlog.
+ */
+ private static class KafkaLatestOffsetEstimator
+ implements GrowableOffsetRangeTracker.RangeEndEstimator {
+
+ private final Consumer<byte[], byte[]> offsetConsumer;
+ private final TopicPartition topicPartition;
+ private final ConsumerSpEL consumerSpEL;
+ private final Supplier<Long> memorizedBacklog;
+
+ KafkaLatestOffsetEstimator(
+ Consumer<byte[], byte[]> offsetConsumer, TopicPartition topicPartition) {
+ this.offsetConsumer = offsetConsumer;
+ this.topicPartition = topicPartition;
+ this.consumerSpEL = new ConsumerSpEL();
+ this.consumerSpEL.evaluateAssign(
+ this.offsetConsumer, ImmutableList.of(this.topicPartition));
+ memorizedBacklog =
+ Suppliers.memoizeWithExpiration(
+ () -> {
+ consumerSpEL.evaluateSeek2End(offsetConsumer, topicPartition);
+ return offsetConsumer.position(topicPartition);
+ },
+ 5,
+ TimeUnit.SECONDS);
+ }
+
+ @Override
+ protected void finalize() {
+ try {
+ Closeables.close(offsetConsumer, true);
+ } catch (Exception anyException) {
+ LOG.warn("Failed to close offset consumer for {}", topicPartition);
+ }
+ }
+
+ @Override
+ public long estimate() {
+ return memorizedBacklog.get();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange initialRestriction(@Element KafkaSourceDescription kafkaSourceDescription) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ try (Consumer<byte[], byte[]> offsetConsumer =
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "initialOffset", offsetConsumerConfig, updatedConsumerConfig))) {
+ consumerSpEL.evaluateAssign(
+ offsetConsumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset;
+ if (kafkaSourceDescription.getStartReadOffset() != null) {
+ startOffset = kafkaSourceDescription.getStartReadOffset();
+ } else if (kafkaSourceDescription.getStartReadTime() != null) {
+ startOffset =
+ consumerSpEL.offsetForTime(
+ offsetConsumer,
+ kafkaSourceDescription.getTopicPartition(),
+ kafkaSourceDescription.getStartReadTime());
+ } else {
+ startOffset = offsetConsumer.position(kafkaSourceDescription.getTopicPartition());
+ }
+ return new OffsetRange(startOffset, Long.MAX_VALUE);
+ }
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public Instant getInitialWatermarkEstimatorState(@Timestamp Instant currentElementTimestamp) {
+ return currentElementTimestamp;
+ }
+
+ @NewWatermarkEstimator
+ public WatermarkEstimatorT newWatermarkEstimator(
+ @WatermarkEstimatorState Instant watermarkEstimatorState) {
+ return createWatermarkEstimatorFn.apply(watermarkEstimatorState);
+ }
+
+ @GetSize
+ public double getSize(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange offsetRange)
+ throws Exception {
+ double numOfRecords =
+ ((HasProgress) restrictionTracker(kafkaSourceDescription, offsetRange))
+ .getProgress()
+ .getWorkRemaining();
+
+ // Before processing elements, we don't have a good estimated size of records and offset gap.
+ if (avgOffsetGap.containsKey(kafkaSourceDescription.getTopicPartition())) {
+ numOfRecords =
+ numOfRecords / (1 + avgOffsetGap.get(kafkaSourceDescription.getTopicPartition()).get());
+ }
+ return (!avgRecordSize.containsKey(kafkaSourceDescription.getTopicPartition())
+ ? 1
+ : avgRecordSize.get(kafkaSourceDescription.getTopicPartition()).get())
+ * numOfRecords;
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> restrictionTracker(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ @Restriction OffsetRange restriction) {
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ KafkaLatestOffsetEstimator offsetPoller =
+ new KafkaLatestOffsetEstimator(
+ consumerFactoryFn.apply(
+ KafkaIOUtils.getOffsetConsumerConfig(
+ "tracker-" + kafkaSourceDescription.getTopicPartition(),
+ offsetConsumerConfig,
+ updatedConsumerConfig)),
+ kafkaSourceDescription.getTopicPartition());
+ return new GrowableOffsetRangeTracker(restriction.getFrom(), offsetPoller);
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element KafkaSourceDescription kafkaSourceDescription,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ WatermarkEstimator watermarkEstimator,
+ OutputReceiver<KafkaRecord<K, V>> receiver) {
+ // If there is no future work, resume with max timeout and move to the next element.
+ if (((HasProgress) tracker).getProgress().getWorkRemaining() <= 0.0) {
+ return ProcessContinuation.resume().withResumeDelay(KAFKA_POLL_TIMEOUT);
+ }
+ Map<String, Object> updatedConsumerConfig =
+ overrideBootstrapServersConfig(consumerConfig, kafkaSourceDescription);
+ // If there is a timestampPolicyFactory, create the TimestampPolicy for current
+ // TopicPartition.
+ TimestampPolicy timestampPolicy = null;
+ if (timestampPolicyFactory != null) {
+ timestampPolicy =
+ timestampPolicyFactory.createTimestampPolicy(
+ kafkaSourceDescription.getTopicPartition(),
+ Optional.ofNullable(watermarkEstimator.currentWatermark()));
+ }
+ try (Consumer<byte[], byte[]> consumer = consumerFactoryFn.apply(updatedConsumerConfig)) {
+ consumerSpEL.evaluateAssign(
+ consumer, ImmutableList.of(kafkaSourceDescription.getTopicPartition()));
+ long startOffset = tracker.currentRestriction().getFrom();
+ long expectedOffset = startOffset;
+ consumer.seek(kafkaSourceDescription.getTopicPartition(), startOffset);
+ ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+
+ try {
+ while (true) {
+ rawRecords = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+ // When there is no records from the current TopicPartition temporarily, self-checkpoint
+ // and move to process the next element.
+ if (rawRecords.isEmpty()) {
+ return ProcessContinuation.resume();
+ }
+ for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ if (!tracker.tryClaim(rawRecord.offset())) {
+ return ProcessContinuation.stop();
+ }
+ KafkaRecord<K, V> kafkaRecord =
+ new KafkaRecord<>(
+ rawRecord.topic(),
+ rawRecord.partition(),
+ rawRecord.offset(),
+ consumerSpEL.getRecordTimestamp(rawRecord),
+ consumerSpEL.getRecordTimestampType(rawRecord),
+ ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+ keyDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.key()),
+ valueDeserializerInstance.deserialize(rawRecord.topic(), rawRecord.value()));
+ int recordSize =
+ (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+ avgOffsetGap
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(expectedOffset - rawRecord.offset());
+ avgRecordSize
+ .computeIfAbsent(kafkaSourceDescription.getTopicPartition(), k -> new MovingAvg())
+ .update(recordSize);
+ expectedOffset = rawRecord.offset() + 1;
+ Instant outputTimestamp;
+ // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+ // WatermarkEstimator should be a Manual one.
+ if (timestampPolicy != null) {
+ checkState(watermarkEstimator instanceof ManualWatermarkEstimator);
+ TimestampPolicyContext context =
+ new TimestampPolicyContext(
+ (long) ((HasProgress) tracker).getProgress().getWorkRemaining(),
+ Instant.now());
Review comment:
If the timeout of the supplier is 5s, it wouldn't have a significant impact.
----------------------------------------------------------------
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:
users@infra.apache.org