Skip to content

Commit 2580835

Browse files
committed
Rename watermark to watermark_delay_threshold for KafkaSource
Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent 6589a2e commit 2580835

4 files changed

Lines changed: 31 additions & 27 deletions

File tree

protos/feast/core/DataSource.proto

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,8 @@ message DataSource {
136136
// Defines the stream data format encoding feature/entity data in Kafka messages.
137137
StreamFormat message_format = 3;
138138

139-
google.protobuf.Duration watermark = 4;
139+
// Watermark delay threshold for stream data
140+
google.protobuf.Duration watermark_delay_threshold = 4;
140141
}
141142

142143
// Defines options for DataSource that sources features from Kinesis records.

sdk/python/feast/data_source.py

Lines changed: 26 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,12 @@ def __init__(
5353
kafka_bootstrap_servers: str,
5454
message_format: StreamFormat,
5555
topic: str,
56-
watermark: Optional[timedelta] = None,
56+
watermark_delay_threshold: Optional[timedelta] = None,
5757
):
5858
self.kafka_bootstrap_servers = kafka_bootstrap_servers
5959
self.message_format = message_format
6060
self.topic = topic
61-
self.watermark = watermark or None
61+
self.watermark_delay_threshold = watermark_delay_threshold or None
6262

6363
@classmethod
6464
def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):
@@ -71,18 +71,18 @@ def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):
7171
Returns:
7272
Returns a BigQueryOptions object based on the kafka_options protobuf
7373
"""
74-
watermark = None
75-
if kafka_options_proto.HasField("watermark"):
76-
watermark = (
74+
watermark_delay_threshold = None
75+
if kafka_options_proto.HasField("watermark_delay_threshold"):
76+
watermark_delay_threshold = (
7777
timedelta(days=0)
78-
if kafka_options_proto.watermark.ToNanoseconds() == 0
79-
else kafka_options_proto.watermark.ToTimedelta()
78+
if kafka_options_proto.watermark_delay_threshold.ToNanoseconds() == 0
79+
else kafka_options_proto.watermark_delay_threshold.ToTimedelta()
8080
)
8181
kafka_options = cls(
8282
kafka_bootstrap_servers=kafka_options_proto.kafka_bootstrap_servers,
8383
message_format=StreamFormat.from_proto(kafka_options_proto.message_format),
8484
topic=kafka_options_proto.topic,
85-
watermark=watermark,
85+
watermark_delay_threshold=watermark_delay_threshold,
8686
)
8787

8888
return kafka_options
@@ -94,16 +94,16 @@ def to_proto(self) -> DataSourceProto.KafkaOptions:
9494
Returns:
9595
KafkaOptionsProto protobuf
9696
"""
97-
watermark_duration = None
98-
if self.watermark is not None:
99-
watermark_duration = Duration()
100-
watermark_duration.FromTimedelta(self.watermark)
97+
watermark_delay_threshold = None
98+
if self.watermark_delay_threshold is not None:
99+
watermark_delay_threshold = Duration()
100+
watermark_delay_threshold.FromTimedelta(self.watermark_delay_threshold)
101101

102102
kafka_options_proto = DataSourceProto.KafkaOptions(
103103
kafka_bootstrap_servers=self.kafka_bootstrap_servers,
104104
message_format=self.message_format.to_proto(),
105105
topic=self.topic,
106-
watermark=watermark_duration,
106+
watermark_delay_threshold=watermark_delay_threshold,
107107
)
108108

109109
return kafka_options_proto
@@ -381,7 +381,7 @@ def __init__(
381381
owner: Optional[str] = "",
382382
timestamp_field: Optional[str] = "",
383383
batch_source: Optional[DataSource] = None,
384-
watermark: Optional[timedelta] = None,
384+
watermark_delay_threshold: Optional[timedelta] = None,
385385
):
386386
"""
387387
Creates a KafkaSource object.
@@ -407,7 +407,8 @@ def __init__(
407407
timestamp_field (optional): Event timestamp field used for point
408408
in time joins of feature values.
409409
batch_source: The datasource that acts as a batch source.
410-
watermark: The watermark for stream data. Specifically how late stream data can arrive without being discarded.
410+
watermark_delay_threshold: The watermark delay threshold for stream data. Specifically how
411+
late stream data can arrive without being discarded.
411412
"""
412413
positional_attributes = [
413414
"name",
@@ -478,7 +479,7 @@ def __init__(
478479
kafka_bootstrap_servers=_kafka_bootstrap_servers,
479480
message_format=_message_format,
480481
topic=_topic,
481-
watermark=watermark,
482+
watermark_delay_threshold=watermark_delay_threshold,
482483
)
483484

484485
def __eq__(self, other):
@@ -495,7 +496,8 @@ def __eq__(self, other):
495496
!= other.kafka_options.kafka_bootstrap_servers
496497
or self.kafka_options.message_format != other.kafka_options.message_format
497498
or self.kafka_options.topic != other.kafka_options.topic
498-
or self.kafka_options.watermark != other.kafka_options.watermark
499+
or self.kafka_options.watermark_delay_threshold
500+
!= other.kafka_options.watermark_delay_threshold
499501
):
500502
return False
501503

@@ -506,12 +508,13 @@ def __hash__(self):
506508

507509
@staticmethod
508510
def from_proto(data_source: DataSourceProto):
509-
watermark = None
510-
if data_source.kafka_options.watermark:
511-
watermark = (
511+
watermark_delay_threshold = None
512+
if data_source.kafka_options.watermark_delay_threshold:
513+
watermark_delay_threshold = (
512514
timedelta(days=0)
513-
if data_source.kafka_options.watermark.ToNanoseconds() == 0
514-
else data_source.kafka_options.watermark.ToTimedelta()
515+
if data_source.kafka_options.watermark_delay_threshold.ToNanoseconds()
516+
== 0
517+
else data_source.kafka_options.watermark_delay_threshold.ToTimedelta()
515518
)
516519
return KafkaSource(
517520
name=data_source.name,
@@ -521,7 +524,7 @@ def from_proto(data_source: DataSourceProto):
521524
message_format=StreamFormat.from_proto(
522525
data_source.kafka_options.message_format
523526
),
524-
watermark=watermark,
527+
watermark_delay_threshold=watermark_delay_threshold,
525528
topic=data_source.kafka_options.topic,
526529
created_timestamp_column=data_source.created_timestamp_column,
527530
timestamp_field=data_source.timestamp_field,

sdk/python/tests/integration/registration/test_registry.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,7 @@ def simple_udf(x: int):
319319
message_format=AvroFormat(""),
320320
topic="topic",
321321
batch_source=FileSource(path="some path"),
322-
watermark=timedelta(days=1),
322+
watermark_delay_threshold=timedelta(days=1),
323323
)
324324

325325
sfv = StreamFeatureView(

sdk/python/tests/integration/registration/test_stream_feature_view_apply.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def test_apply_stream_feature_view(simple_dataset_1) -> None:
3333
message_format=AvroFormat(""),
3434
topic="topic",
3535
batch_source=file_source,
36-
watermark=timedelta(days=1),
36+
watermark_delay_threshold=timedelta(days=1),
3737
)
3838

3939
@stream_feature_view(
@@ -97,7 +97,7 @@ def test_stream_feature_view_udf(simple_dataset_1) -> None:
9797
message_format=AvroFormat(""),
9898
topic="topic",
9999
batch_source=file_source,
100-
watermark=timedelta(days=1),
100+
watermark_delay_threshold=timedelta(days=1),
101101
)
102102

103103
@stream_feature_view(

0 commit comments

Comments
 (0)