Skip to content

Commit e2e28c7

Browse files
Rodrigo-Palmantkathole
authored andcommitted
fix: Keep a zero watermark delay threshold on a Kafka source
KafkaOptions.__init__ stored 'watermark_delay_threshold or None', and timedelta(0) is falsy, so a delay of zero was read as 'not given': KafkaOptions(..., watermark_delay_threshold=timedelta(0)) .watermark_delay_threshold -> None to_proto().HasField(...) -> False The value never reached the proto, so it was lost on apply and on every registry round-trip, while timedelta(seconds=30) survived. Both from_proto sites carried a branch mapping an explicit zero Duration back to timedelta(days=0). That branch could not return anything: the constructor it hands the value to dropped it. Feeding a proto that does carry a zero Duration returned None. KafkaSource.from_proto also tested the Duration message for truth instead of presence. A protobuf message is always truthy, so an absent delay entered the branch and came out as timedelta(days=0); only the 'or None' above hid it. KafkaOptions.from_proto, its sibling in this file, already uses HasField, which is what both now use. An absent delay stays None; a zero delay stays zero. Signed-off-by: Rodrigo-Palma <[email protected]>
1 parent 4771156 commit e2e28c7

2 files changed

Lines changed: 50 additions & 9 deletions

File tree

‎sdk/python/feast/data_source.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def __init__(
4747
self.kafka_bootstrap_servers = kafka_bootstrap_servers
4848
self.message_format = message_format
4949
self.topic = topic
50-
self.watermark_delay_threshold = watermark_delay_threshold or None
50+
self.watermark_delay_threshold = watermark_delay_threshold
5151

5252
@classmethod
5353
def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):
@@ -63,9 +63,7 @@ def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):
6363
watermark_delay_threshold = None
6464
if kafka_options_proto.HasField("watermark_delay_threshold"):
6565
watermark_delay_threshold = (
66-
timedelta(days=0)
67-
if kafka_options_proto.watermark_delay_threshold.ToNanoseconds() == 0
68-
else kafka_options_proto.watermark_delay_threshold.ToTimedelta()
66+
kafka_options_proto.watermark_delay_threshold.ToTimedelta()
6967
)
7068
kafka_options = cls(
7169
kafka_bootstrap_servers=kafka_options_proto.kafka_bootstrap_servers,
@@ -552,12 +550,9 @@ def __hash__(self):
552550
@staticmethod
553551
def from_proto(data_source: DataSourceProto):
554552
watermark_delay_threshold = None
555-
if data_source.kafka_options.watermark_delay_threshold:
553+
if data_source.kafka_options.HasField("watermark_delay_threshold"):
556554
watermark_delay_threshold = (
557-
timedelta(days=0)
558-
if data_source.kafka_options.watermark_delay_threshold.ToNanoseconds()
559-
== 0
560-
else data_source.kafka_options.watermark_delay_threshold.ToTimedelta()
555+
data_source.kafka_options.watermark_delay_threshold.ToTimedelta()
561556
)
562557
return KafkaSource(
563558
name=data_source.name,

‎sdk/python/tests/unit/test_data_sources.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from datetime import timedelta
2+
13
import pytest
24

35
from feast.data_format import ProtoFormat
@@ -333,3 +335,47 @@ def test_data_source_eq_cross_type_returns_false():
333335
assert file_source == FileSource(
334336
name="src", path="/tmp/x.parquet", timestamp_field="ts"
335337
)
338+
339+
340+
def test_kafka_source_keeps_a_zero_watermark_delay_threshold():
341+
"""A watermark delay of zero is a value, not an unset option.
342+
343+
``KafkaOptions`` stored ``watermark_delay_threshold or None``, so a
344+
``timedelta(0)`` was read as "not given" and never reached the proto. The
345+
``from_proto`` branch that turned an explicit zero ``Duration`` back into
346+
``timedelta(0)`` could therefore never return it.
347+
"""
348+
file_source = FileSource(
349+
name="batch", path="/tmp/x.parquet", timestamp_field="event_timestamp"
350+
)
351+
352+
def kafka_source(watermark):
353+
return KafkaSource(
354+
name="test_source",
355+
kafka_bootstrap_servers="test_servers",
356+
message_format=ProtoFormat("class_path"),
357+
topic="test_topic",
358+
timestamp_field="event_timestamp",
359+
watermark_delay_threshold=watermark,
360+
batch_source=file_source,
361+
)
362+
363+
zero = kafka_source(timedelta(0))
364+
assert zero.kafka_options.watermark_delay_threshold == timedelta(0)
365+
366+
zero_proto = zero.to_proto()
367+
assert zero_proto.kafka_options.HasField("watermark_delay_threshold")
368+
assert KafkaSource.from_proto(
369+
zero_proto
370+
).kafka_options.watermark_delay_threshold == timedelta(0)
371+
372+
# An absent delay stays absent, and is not read as a zero
373+
unset = kafka_source(None)
374+
assert unset.kafka_options.watermark_delay_threshold is None
375+
376+
unset_proto = unset.to_proto()
377+
assert not unset_proto.kafka_options.HasField("watermark_delay_threshold")
378+
assert (
379+
KafkaSource.from_proto(unset_proto).kafka_options.watermark_delay_threshold
380+
is None
381+
)

0 commit comments

Comments
 (0)