Skip to content

Commit f5dd621

Browse files
ferenc-csakyferenc-csaky
committed
feat: Add Adaptive Kafka watermark generator
1 parent 1a63811 commit f5dd621

7 files changed

Lines changed: 459 additions & 12 deletions

File tree

‎connectors/kafka-safe-connector/src/main/java/com/datasqrl/flinkrunner/connector/kafka/KafkaRecordTimestampWatermarkStrategy.java‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
/*
2+
* Copyright © 2026 DataSQRL (contact@datasqrl.com)
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.datasqrl.flinkrunner.connector.kafka;
17+
18+
import java.io.Serial;
19+
import java.time.Duration;
20+
import java.util.Arrays;
21+
import org.apache.flink.annotation.Internal;
22+
import org.apache.flink.api.common.eventtime.Watermark;
23+
import org.apache.flink.api.common.eventtime.WatermarkGenerator;
24+
import org.apache.flink.api.common.eventtime.WatermarkGeneratorSupplier;
25+
import org.apache.flink.api.common.eventtime.WatermarkOutput;
26+
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
27+
import org.apache.flink.table.data.RowData;
28+
29+
/** Adaptive watermark strategy based on Kafka record timestamps. */
30+
@Internal
31+
public final class KafkaRecordTimestampWatermarkStrategy implements WatermarkStrategy<RowData> {
32+
33+
@Serial private static final long serialVersionUID = 1L;
34+
35+
static final int MIN_RECORDS = 250;
36+
static final int SAMPLE_SIZE = 4096;
37+
static final long MIN_OUT_OF_ORDERNESS_MILLIS = 50L;
38+
static final long MAX_OUT_OF_ORDERNESS_MILLIS = Duration.ofDays(1).toMillis();
39+
40+
private static final double OUT_OF_ORDERNESS_QUANTILE = 0.95D;
41+
42+
public static final KafkaRecordTimestampWatermarkStrategy INSTANCE =
43+
new KafkaRecordTimestampWatermarkStrategy();
44+
45+
private KafkaRecordTimestampWatermarkStrategy() {}
46+
47+
@Override
48+
public WatermarkGenerator<RowData> createWatermarkGenerator(
49+
WatermarkGeneratorSupplier.Context context) {
50+
51+
return new AdaptiveKafkaRecordTimestampWatermarkGenerator();
52+
}
53+
54+
/**
55+
* Watermark generator that derives event time from Kafka record timestamps and adapts its
56+
* out-of-orderness delay from recently observed records.
57+
*
58+
* <p>The generator watches how far records tend to arrive behind the newest Kafka timestamp seen
59+
* so far. It keeps a rolling sample of those delays and uses the 95th percentile as the safety
60+
* margin before advancing the watermark. This lets the source tolerate typical out-of-order
61+
* records without waiting for rare extreme delays.
62+
*
63+
* <p>The {@code eventTimestamp} argument is supplied by Flink from the Kafka source record
64+
* timestamp.
65+
*/
66+
private static final class AdaptiveKafkaRecordTimestampWatermarkGenerator
67+
implements WatermarkGenerator<RowData> {
68+
69+
/**
70+
* Circular buffer of observed lateness values in milliseconds.
71+
*
72+
* <p>For a record with timestamp {@code t}, lateness is {@code max(0, maxTimestamp - t)} using
73+
* the maximum timestamp seen before that record. In-order records therefore contribute {@code
74+
* 0}, while older records contribute how far they lag behind the current observed frontier.
75+
*/
76+
private final long[] latenessSamples = new long[SAMPLE_SIZE];
77+
78+
/** Highest Kafka record timestamp observed by this generator. */
79+
private long maxTimestamp = Long.MIN_VALUE;
80+
81+
/** Last emitted watermark, used to preserve Flink's monotonic watermark contract. */
82+
private long lastEmittedWatermark = Long.MIN_VALUE;
83+
84+
/**
85+
* Number of records observed, including the first record that cannot produce a lateness sample.
86+
*/
87+
private long recordCount;
88+
89+
/** Next slot in the circular lateness sample buffer. */
90+
private int nextSampleIndex;
91+
92+
@Override
93+
public void onEvent(RowData event, long eventTimestamp, WatermarkOutput output) {
94+
if (maxTimestamp != Long.MIN_VALUE) {
95+
// Sample lateness before updating maxTimestamp, so the sample reflects disorder relative to
96+
// the already observed stream frontier.
97+
latenessSamples[nextSampleIndex] = Math.max(0L, maxTimestamp - eventTimestamp);
98+
nextSampleIndex = (nextSampleIndex + 1) % SAMPLE_SIZE;
99+
}
100+
101+
maxTimestamp = Math.max(maxTimestamp, eventTimestamp);
102+
recordCount++;
103+
}
104+
105+
@Override
106+
public void onPeriodicEmit(WatermarkOutput output) {
107+
if (recordCount < MIN_RECORDS || maxTimestamp == Long.MIN_VALUE) {
108+
// Avoid emitting during warm-up, so limited early samples won't distort the estimate.
109+
return;
110+
}
111+
112+
long outOfOrdernessMillis = calculateOutOfOrdernessMillis();
113+
long watermark = maxTimestamp - outOfOrdernessMillis - 1;
114+
115+
if (watermark > lastEmittedWatermark) {
116+
output.emitWatermark(new Watermark(watermark));
117+
lastEmittedWatermark = watermark;
118+
}
119+
}
120+
121+
private long calculateOutOfOrdernessMillis() {
122+
int sampleCount = (int) Math.min(recordCount - 1, SAMPLE_SIZE);
123+
if (sampleCount <= 0) {
124+
return MIN_OUT_OF_ORDERNESS_MILLIS;
125+
}
126+
127+
long[] samples = Arrays.copyOf(latenessSamples, sampleCount);
128+
Arrays.sort(samples);
129+
130+
// Use a high quantile rather than the maximum, so a single pathological record does not stall
131+
// all event-time progress until it rotates out of the sample buffer.
132+
int quantileIndex =
133+
Math.min(sampleCount - 1, (int) Math.ceil(sampleCount * OUT_OF_ORDERNESS_QUANTILE) - 1);
134+
135+
return Math.min(
136+
MAX_OUT_OF_ORDERNESS_MILLIS,
137+
Math.max(MIN_OUT_OF_ORDERNESS_MILLIS, samples[quantileIndex]));
138+
}
139+
}
140+
}

‎connectors/kafka-safe-connector/src/main/java/org/apache/flink/streaming/connectors/kafka/table/SafeKafkaDynamicSource.java‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package org.apache.flink.streaming.connectors.kafka.table;
1717

18+
import com.datasqrl.flinkrunner.connector.kafka.KafkaRecordTimestampWatermarkStrategy;
1819
import org.apache.flink.annotation.Internal;
1920
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
2021
import org.apache.flink.api.common.serialization.DeserializationSchema;
@@ -40,6 +41,7 @@
4041
import org.apache.flink.table.connector.source.DynamicTableSource;
4142
import org.apache.flink.table.connector.source.ScanTableSource;
4243
import org.apache.flink.table.connector.source.abilities.SupportsReadingMetadata;
44+
import org.apache.flink.table.connector.source.abilities.SupportsSourceWatermark;
4345
import org.apache.flink.table.connector.source.abilities.SupportsWatermarkPushDown;
4446
import org.apache.flink.table.data.GenericMapData;
4547
import org.apache.flink.table.data.RowData;
@@ -82,7 +84,10 @@
8284
/** A version-agnostic Kafka {@link ScanTableSource}. */
8385
@Internal
8486
public class SafeKafkaDynamicSource
85-
implements ScanTableSource, SupportsReadingMetadata, SupportsWatermarkPushDown {
87+
implements ScanTableSource,
88+
SupportsReadingMetadata,
89+
SupportsSourceWatermark,
90+
SupportsWatermarkPushDown {
8691

8792
private static final String KAFKA_TRANSFORMATION = "kafka";
8893

@@ -99,6 +104,9 @@ public class SafeKafkaDynamicSource
99104
/** Watermark strategy that is used to generate per-partition watermark. */
100105
protected @Nullable WatermarkStrategy<RowData> watermarkStrategy;
101106

107+
/** Whether to use Kafka record timestamps to generate source watermarks. */
108+
protected boolean sourceWatermarkEnabled;
109+
102110
// --------------------------------------------------------------------------------------------
103111
// Format attributes
104112
// --------------------------------------------------------------------------------------------
@@ -215,6 +223,7 @@ public SafeKafkaDynamicSource(
215223
this.producedDataType = physicalDataType;
216224
this.metadataKeys = Collections.emptyList();
217225
this.watermarkStrategy = null;
226+
this.sourceWatermarkEnabled = false;
218227
// Kafka-specific attributes
219228
Preconditions.checkArgument(
220229
(topics != null && topicPattern == null)
@@ -264,9 +273,7 @@ public ScanRuntimeProvider getScanRuntimeProvider(ScanContext context) {
264273
@Override
265274
public DataStream<RowData> produceDataStream(
266275
ProviderContext providerContext, StreamExecutionEnvironment execEnv) {
267-
if (watermarkStrategy == null) {
268-
watermarkStrategy = WatermarkStrategy.noWatermarks();
269-
}
276+
final WatermarkStrategy<RowData> watermarkStrategy = getWatermarkStrategy();
270277
DataStreamSource<RowData> sourceStream =
271278
execEnv.fromSource(
272279
kafkaSource, watermarkStrategy, "KafkaSource-" + tableIdentifier);
@@ -340,6 +347,11 @@ public void applyWatermark(WatermarkStrategy<RowData> watermarkStrategy) {
340347
this.watermarkStrategy = watermarkStrategy;
341348
}
342349

350+
@Override
351+
public void applySourceWatermark() {
352+
this.sourceWatermarkEnabled = true;
353+
}
354+
343355
@Override
344356
public DynamicTableSource copy() {
345357
final SafeKafkaDynamicSource copy =
@@ -366,6 +378,7 @@ public DynamicTableSource copy() {
366378
copy.producedDataType = producedDataType;
367379
copy.metadataKeys = metadataKeys;
368380
copy.watermarkStrategy = watermarkStrategy;
381+
copy.sourceWatermarkEnabled = sourceWatermarkEnabled;
369382
return copy;
370383
}
371384

@@ -403,6 +416,7 @@ public boolean equals(Object o) {
403416
&& Objects.equals(upsertMode, that.upsertMode)
404417
&& Objects.equals(tableIdentifier, that.tableIdentifier)
405418
&& Objects.equals(watermarkStrategy, that.watermarkStrategy)
419+
&& sourceWatermarkEnabled == that.sourceWatermarkEnabled
406420
&& Objects.equals(parallelism, that.parallelism);
407421
}
408422

@@ -429,6 +443,7 @@ public int hashCode() {
429443
upsertMode,
430444
tableIdentifier,
431445
watermarkStrategy,
446+
sourceWatermarkEnabled,
432447
parallelism);
433448
}
434449

@@ -509,6 +524,18 @@ protected KafkaSource<RowData> createKafkaSource(
509524
return kafkaSourceBuilder.build();
510525
}
511526

527+
private WatermarkStrategy<RowData> getWatermarkStrategy() {
528+
if (watermarkStrategy != null) {
529+
return watermarkStrategy;
530+
}
531+
532+
if (sourceWatermarkEnabled) {
533+
return KafkaRecordTimestampWatermarkStrategy.INSTANCE;
534+
}
535+
536+
return WatermarkStrategy.noWatermarks();
537+
}
538+
512539
private OffsetResetStrategy getResetStrategy(String offsetResetConfig) {
513540
return Arrays.stream(OffsetResetStrategy.values())
514541
.filter(ors -> ors.name().equals(offsetResetConfig.toUpperCase(Locale.ROOT)))

‎connectors/kafka-safe-connector/src/test/java/com/datasqrl/flinkrunner/connector/kafka/KafkaRecordTimestampWatermarkStrategyTest.java‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
/*
2+
* Copyright © 2026 DataSQRL (contact@datasqrl.com)
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.datasqrl.flinkrunner.connector.kafka;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
20+
import java.util.ArrayList;
21+
import java.util.List;
22+
import org.apache.flink.api.common.eventtime.Watermark;
23+
import org.apache.flink.api.common.eventtime.WatermarkGenerator;
24+
import org.apache.flink.api.common.eventtime.WatermarkOutput;
25+
import org.apache.flink.table.data.RowData;
26+
import org.junit.jupiter.api.Test;
27+
28+
/** Tests for {@link KafkaRecordTimestampWatermarkStrategy}. */
29+
class KafkaRecordTimestampWatermarkStrategyTest {
30+
31+
@Test
32+
void testDoesNotEmitBeforeWarmup() {
33+
final WatermarkGenerator<RowData> generator = createGenerator();
34+
final CollectingWatermarkOutput output = new CollectingWatermarkOutput();
35+
36+
for (int i = 0; i < KafkaRecordTimestampWatermarkStrategy.MIN_RECORDS - 1; i++) {
37+
generator.onEvent(null, i, output);
38+
}
39+
generator.onPeriodicEmit(output);
40+
41+
assertThat(output.watermarks).isEmpty();
42+
}
43+
44+
@Test
45+
void testEmitsWatermarkFromKafkaRecordTimestampAfterWarmup() {
46+
final WatermarkGenerator<RowData> generator = createGenerator();
47+
final CollectingWatermarkOutput output = new CollectingWatermarkOutput();
48+
49+
for (int i = 0; i < KafkaRecordTimestampWatermarkStrategy.MIN_RECORDS; i++) {
50+
generator.onEvent(null, i * 100L, output);
51+
}
52+
generator.onPeriodicEmit(output);
53+
54+
assertThat(output.watermarks)
55+
.containsExactly(
56+
(KafkaRecordTimestampWatermarkStrategy.MIN_RECORDS - 1) * 100L
57+
- KafkaRecordTimestampWatermarkStrategy.MIN_OUT_OF_ORDERNESS_MILLIS
58+
- 1);
59+
}
60+
61+
@Test
62+
void testWatermarksAreMonotonic() {
63+
final WatermarkGenerator<RowData> generator = createGenerator();
64+
final CollectingWatermarkOutput output = new CollectingWatermarkOutput();
65+
66+
for (int i = 0; i < KafkaRecordTimestampWatermarkStrategy.MIN_RECORDS; i++) {
67+
generator.onEvent(null, i * 100L, output);
68+
}
69+
generator.onPeriodicEmit(output);
70+
71+
for (int i = 0; i < KafkaRecordTimestampWatermarkStrategy.MIN_RECORDS; i++) {
72+
generator.onEvent(null, 0L, output);
73+
}
74+
generator.onPeriodicEmit(output);
75+
76+
assertThat(output.watermarks).hasSize(1);
77+
}
78+
79+
private static WatermarkGenerator<RowData> createGenerator() {
80+
return KafkaRecordTimestampWatermarkStrategy.INSTANCE.createWatermarkGenerator(null);
81+
}
82+
83+
private static final class CollectingWatermarkOutput implements WatermarkOutput {
84+
85+
private final List<Long> watermarks = new ArrayList<>();
86+
87+
@Override
88+
public void emitWatermark(Watermark watermark) {
89+
watermarks.add(watermark.getTimestamp());
90+
}
91+
92+
@Override
93+
public void markIdle() {}
94+
95+
@Override
96+
public void markActive() {}
97+
}
98+
}

0 commit comments

Comments
 (0)