[
https://issues.apache.org/jira/browse/BEAM-3851?focusedWorklogId=83270&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-83270
]
ASF GitHub Bot logged work on BEAM-3851:
----------------------------------------
Author: ASF GitHub Bot
Created on: 22/Mar/18 17:56
Start Date: 22/Mar/18 17:56
Worklog Time Spent: 10m
Work Description: XuMingmin closed pull request #4868: [BEAM-3851] Option
to preserve element timestamp while publishing to Kafka.
URL: https://github.com/apache/beam/pull/4868
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaExactlyOnceSink.java
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaExactlyOnceSink.java
index 7345a92f201..9ae69da3834 100644
---
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaExactlyOnceSink.java
+++
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaExactlyOnceSink.java
@@ -64,6 +64,8 @@
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.sdk.values.TimestampedValue.TimestampedValueCoder;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
@@ -173,7 +175,8 @@ static void ensureEOSSupport() {
/**
* Shuffle messages assigning each randomly to a shard.
*/
- private static class Reshard<K, V> extends DoFn<KV<K, V>, KV<Integer, KV<K,
V>>> {
+ private static class Reshard<K, V>
+ extends DoFn<KV<K, V>, KV<Integer, TimestampedValue<KV<K, V>>>> {
private final int numShards;
private transient int shardId;
@@ -190,12 +193,13 @@ public void setup() {
@ProcessElement
public void processElement(ProcessContext ctx) {
shardId = (shardId + 1) % numShards; // round-robin among shards.
- ctx.output(KV.of(shardId, ctx.element()));
+ ctx.output(KV.of(shardId, TimestampedValue.of(ctx.element(),
ctx.timestamp())));
}
}
- private static class Sequencer<K, V>
- extends DoFn<KV<Integer, Iterable<KV<K, V>>>, KV<Integer, KV<Long, KV<K,
V>>>> {
+ private static class Sequencer<K, V> extends DoFn<
+ KV<Integer, Iterable<TimestampedValue<KV<K, V>>>>,
+ KV<Integer, KV<Long, TimestampedValue<KV<K, V>>>>> {
private static final String NEXT_ID = "nextId";
@StateId(NEXT_ID)
@@ -205,7 +209,7 @@ public void processElement(ProcessContext ctx) {
public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState,
ProcessContext ctx) {
long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);
int shard = ctx.element().getKey();
- for (KV<K, V> value : ctx.element().getValue()) {
+ for (TimestampedValue<KV<K, V>> value : ctx.element().getValue()) {
ctx.output(KV.of(shard, KV.of(nextId, value)));
nextId++;
}
@@ -214,7 +218,7 @@ public void processElement(@StateId(NEXT_ID)
ValueState<Long> nextIdState, Proce
}
private static class ExactlyOnceWriter<K, V>
- extends DoFn<KV<Integer, Iterable<KV<Long, KV<K, V>>>>, Void> {
+ extends DoFn<KV<Integer, Iterable<KV<Long, TimestampedValue<KV<K, V>>>>>,
Void> {
private static final String NEXT_ID = "nextId";
private static final String MIN_BUFFERED_ID = "minBufferedId";
@@ -230,7 +234,7 @@ public void processElement(@StateId(NEXT_ID)
ValueState<Long> nextIdState, Proce
@StateId(MIN_BUFFERED_ID)
private final StateSpec<ValueState<Long>> minBufferedIdSpec =
StateSpecs.value();
@StateId(OUT_OF_ORDER_BUFFER)
- private final StateSpec<BagState<KV<Long, KV<K, V>>>> outOfOrderBufferSpec;
+ private final StateSpec<BagState<KV<Long, TimestampedValue<KV<K, V>>>>>
outOfOrderBufferSpec;
// A random id assigned to each shard. Helps with detecting when multiple
jobs are mistakenly
// started with same groupId used for storing state on Kafka side,
including the case where
// a job is restarted with same groupId, but the metadata from previous
run was not cleared.
@@ -248,7 +252,8 @@ public void processElement(@StateId(NEXT_ID)
ValueState<Long> nextIdState, Proce
ExactlyOnceWriter(Write<K, V> spec, Coder<KV<K, V>> elemCoder) {
this.spec = spec;
- this.outOfOrderBufferSpec =
StateSpecs.bag(KvCoder.of(BigEndianLongCoder.of(), elemCoder));
+ this.outOfOrderBufferSpec = StateSpecs.bag(
+ KvCoder.of(BigEndianLongCoder.of(),
TimestampedValueCoder.of(elemCoder)));
}
@Setup
@@ -261,7 +266,7 @@ public void setup() {
public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState,
@StateId(MIN_BUFFERED_ID) ValueState<Long>
minBufferedIdState,
@StateId(OUT_OF_ORDER_BUFFER)
- BagState<KV<Long, KV<K, V>>> oooBufferState,
+ BagState<KV<Long, TimestampedValue<KV<K,
V>>>> oooBufferState,
@StateId(WRITER_ID) ValueState<String>
writerIdState,
ProcessContext ctx)
throws IOException {
@@ -297,10 +302,10 @@ public void processElement(@StateId(NEXT_ID)
ValueState<Long> nextIdState,
// There might be out of order messages buffered in earlier
iterations. These
// will get merged if and when minBufferedId matches nextId.
- Iterator<KV<Long, KV<K, V>>> iter =
ctx.element().getValue().iterator();
+ Iterator<KV<Long, TimestampedValue<KV<K, V>>>> iter =
ctx.element().getValue().iterator();
while (iter.hasNext()) {
- KV<Long, KV<K, V>> kv = iter.next();
+ KV<Long, TimestampedValue<KV<K, V>>> kv = iter.next();
long recordId = kv.getKey();
if (recordId < nextId) {
@@ -339,7 +344,8 @@ public void processElement(@StateId(NEXT_ID)
ValueState<Long> nextIdState,
// Read all of them in to memory and sort them. Reading into memory
// might be problematic in extreme cases. Might need to improve it
in future.
- List<KV<Long, KV<K, V>>> buffered =
Lists.newArrayList(oooBufferState.read());
+ List<KV<Long, TimestampedValue<KV<K, V>>>> buffered =
+ Lists.newArrayList(oooBufferState.read());
buffered.sort(new KV.OrderByKey<>());
LOG.info("{} : merging {} buffered records (min buffered id is
{}).",
@@ -349,8 +355,7 @@ public void processElement(@StateId(NEXT_ID)
ValueState<Long> nextIdState,
minBufferedIdState.clear();
minBufferedId = Long.MAX_VALUE;
- iter =
- Iterators.mergeSorted(
+ iter = Iterators.mergeSorted(
ImmutableList.of(iter, buffered.iterator()), new
KV.OrderByKey<>());
}
}
@@ -428,10 +433,17 @@ void beginTxn() {
ProducerSpEL.beginTransaction(producer);
}
- void sendRecord(KV<K, V> record, Counter sendCounter) {
+ void sendRecord(TimestampedValue<KV<K, V>> record, Counter sendCounter) {
try {
+ Long timestampMillis = spec.getPublishTimestampFunction() != null
+ ?
spec.getPublishTimestampFunction().getTimestamp(record.getValue(),
+
record.getTimestamp()).getMillis()
+ : null;
+
producer.send(
- new ProducerRecord<>(spec.getTopic(), record.getKey(),
record.getValue()));
+ new ProducerRecord<>(
+ spec.getTopic(), null, timestampMillis,
+ record.getValue().getKey(), record.getValue().getValue()));
sendCounter.inc();
} catch (KafkaException e) {
ProducerSpEL.abortTransaction(producer);
diff --git
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index eb292299256..eeb9da95814 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -174,10 +174,15 @@
* .withKeySerializer(LongSerializer.class)
* .withValueSerializer(StringSerializer.class)
*
- * // you can further customize KafkaProducer used to write the records
by adding more
+ * // You can further customize KafkaProducer used to write the records
by adding more
* // settings for ProducerConfig. e.g, to enable compression :
* .updateProducerProperties(ImmutableMap.of("compression.type", "gzip"))
*
+ * // You set publish timestamp for the Kafka records.
+ * .withInputTimestamp() // element timestamp is used while publishing
to Kafka
+ * // or you can also set a custom timestamp with a function.
+ * .withPublishTimestampFunction((elem, elemTs) -> ...)
+ *
* // Optionally enable exactly-once sink (on supported runners). See
JavaDoc for withEOS().
* .withEOS(20, "eos-sink-group-id");
* );
@@ -813,6 +818,8 @@ private KafkaIO() {}
@Nullable abstract Class<? extends Serializer<K>> getKeySerializer();
@Nullable abstract Class<? extends Serializer<V>> getValueSerializer();
+ @Nullable abstract KafkaPublishTimestampFunction<KV<K, V>>
getPublishTimestampFunction();
+
// Configuration for EOS sink
abstract boolean isEOS();
@Nullable abstract String getSinkGroupId();
@@ -830,6 +837,8 @@ private KafkaIO() {}
SerializableFunction<Map<String, Object>, Producer<K, V>> fn);
abstract Builder<K, V> setKeySerializer(Class<? extends Serializer<K>>
serializer);
abstract Builder<K, V> setValueSerializer(Class<? extends Serializer<V>>
serializer);
+ abstract Builder<K, V> setPublishTimestampFunction(
+ KafkaPublishTimestampFunction<KV<K, V>> timestampFunction);
abstract Builder<K, V> setEOS(boolean eosEnabled);
abstract Builder<K, V> setSinkGroupId(String sinkGroupId);
abstract Builder<K, V> setNumShards(int numShards);
@@ -889,6 +898,28 @@ private KafkaIO() {}
return toBuilder().setProducerFactoryFn(producerFactoryFn).build();
}
+ /**
+ * The timestamp for each record being published is set to timestamp of
the element in the
+ * pipeline. This is equivalent to {@code withPublishTimestampFunction((e,
ts) -> ts)}. <br>
+ * NOTE: Kafka's retention policies are based on message timestamps. If
the pipeline
+ * is processing messages from the past, they might be deleted immediately
by Kafka after
+ * being published if the timestamps are older than Kafka cluster's {@code
log.retention.hours}.
+ */
+ public Write<K, V> withInputTimestamp() {
+ return
withPublishTimestampFunction(KafkaPublishTimestampFunction.withElementTimestamp());
+ }
+
+ /**
+ * A function to provide timestamp for records being published. <br>
+ * NOTE: Kafka's retention policies are based on message timestamps. If
the pipeline
+ * is processing messages from the past, they might be deleted immediately
by Kafka after
+ * being published if the timestamps are older than Kafka cluster's {@code
log.retention.hours}.
+ */
+ public Write<K, V> withPublishTimestampFunction(
+ KafkaPublishTimestampFunction<KV<K, V>> timestampFunction) {
+ return
toBuilder().setPublishTimestampFunction(timestampFunction).build();
+ }
+
/**
* Provides exactly-once semantics while writing to Kafka, which enables
applications with
* end-to-end exactly-once guarantees on top of exactly-once semantics
<i>within</i> Beam
diff --git
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaPublishTimestampFunction.java
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaPublishTimestampFunction.java
new file mode 100644
index 00000000000..e0ef6397d23
--- /dev/null
+++
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaPublishTimestampFunction.java
@@ -0,0 +1,45 @@
+/*
+ * 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 java.io.Serializable;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.joda.time.Instant;
+
+/**
+ * An interface for providing custom timestamp for elements written to Kafka.
+ */
+public interface KafkaPublishTimestampFunction<T> extends Serializable {
+
+ /**
+ * Returns timestamp for element being published to Kafka.
+ * See @{@link org.apache.kafka.clients.producer.ProducerRecord}.
+ *
+ * @param element The element being published.
+ * @param elementTimestamp Timestamp of the element from the context
+ * (i.e. @{@link DoFn.ProcessContext#timestamp()}
+ */
+ Instant getTimestamp(T element, Instant elementTimestamp);
+
+ /**
+ * Returns {@link KafkaPublishTimestampFunction} returns element timestamp
from ProcessContext.
+ */
+ static <T> KafkaPublishTimestampFunction<T> withElementTimestamp() {
+ return (element, elementTimestamp) -> elementTimestamp;
+ }
+}
diff --git
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriter.java
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriter.java
index 00b76e503b8..9f2544a96b0 100644
---
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriter.java
+++
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriter.java
@@ -54,8 +54,12 @@ public void processElement(ProcessContext ctx) throws
Exception {
checkForFailures();
KV<K, V> kv = ctx.element();
- producer.send(
- new ProducerRecord<>(spec.getTopic(), kv.getKey(), kv.getValue()), new
SendCallback());
+ Long timestampMillis = spec.getPublishTimestampFunction() != null
+ ? spec.getPublishTimestampFunction().getTimestamp(kv,
ctx.timestamp()).getMillis()
+ : null;
+
+ producer.send(new ProducerRecord<>(
+ spec.getTopic(), null, timestampMillis, kv.getKey(), kv.getValue()),
new SendCallback());
elementsWritten.inc();
}
diff --git
a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 3718c410827..7ae8f1a3ad8 100644
---
a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++
b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -854,13 +854,14 @@ public void testSink() throws Exception {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
+ .withInputTimestamp()
.withProducerFactoryFn(new
ProducerFactoryFn(producerWrapper.producerKey)));
p.run();
completionThread.shutdown();
- verifyProducerRecords(producerWrapper.mockProducer, topic, numElements,
false);
+ verifyProducerRecords(producerWrapper.mockProducer, topic, numElements,
false, true);
}
}
@@ -891,7 +892,7 @@ public void testValuesSink() throws Exception {
completionThread.shutdown();
- verifyProducerRecords(producerWrapper.mockProducer, topic, numElements,
true);
+ verifyProducerRecords(producerWrapper.mockProducer, topic, numElements,
true, false);
}
}
@@ -930,13 +931,14 @@ public void testExactlyOnceSink() {
.withEOS(1, "test")
.withConsumerFactoryFn(new ConsumerFactoryFn(
Lists.newArrayList(topic), 10, 10,
OffsetResetStrategy.EARLIEST))
+ .withPublishTimestampFunction((e, ts) -> ts)
.withProducerFactoryFn(new
ProducerFactoryFn(producerWrapper.producerKey)));
p.run();
completionThread.shutdown();
- verifyProducerRecords(producerWrapper.mockProducer, topic, numElements,
false);
+ verifyProducerRecords(producerWrapper.mockProducer, topic, numElements,
false, true);
}
}
@@ -1198,7 +1200,10 @@ public void testSinkMetrics() throws Exception {
}
private static void verifyProducerRecords(MockProducer<Integer, Long>
mockProducer,
- String topic, int numElements,
boolean keyIsAbsent) {
+ String topic,
+ int numElements,
+ boolean keyIsAbsent,
+ boolean verifyTimestamp) {
// verify that appropriate messages are written to kafka
List<ProducerRecord<Integer, Long>> sent = mockProducer.history();
@@ -1215,6 +1220,9 @@ private static void
verifyProducerRecords(MockProducer<Integer, Long> mockProduc
assertEquals(i, record.key().intValue());
}
assertEquals(i, record.value().longValue());
+ if (verifyTimestamp) {
+ assertEquals(i, record.timestamp().intValue());
+ }
}
}
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
Issue Time Tracking
-------------------
Worklog Id: (was: 83270)
Time Spent: 2h (was: 1h 50m)
> Support element timestamps while publishing to Kafka.
> -----------------------------------------------------
>
> Key: BEAM-3851
> URL: https://issues.apache.org/jira/browse/BEAM-3851
> Project: Beam
> Issue Type: Improvement
> Components: io-java-kafka
> Affects Versions: 2.3.0
> Reporter: Raghu Angadi
> Assignee: Raghu Angadi
> Priority: Major
> Fix For: 2.5.0
>
> Time Spent: 2h
> Remaining Estimate: 0h
>
> KafkaIO sink should support using input element timestamp for the message
> published to Kafka. Otherwise there is no way for user to influence the
> timestamp of the messages in Kafka sink.
--
This message was sent by Atlassian JIRA
(v7.6.3#76005)
