# Copyright 2020 The Feast Authors
#
# Licensed 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
#
# https://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.
import enum
import warnings
from abc import ABC, abstractmethod
from datetime import timedelta
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
from google.protobuf.duration_pb2 import Duration
from google.protobuf.json_format import MessageToJson
from typeguard import typechecked
from feast import type_map
from feast.data_format import StreamFormat
from feast.field import Field
from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
from feast.repo_config import RepoConfig, get_data_source_class_from_type
from feast.types import from_value_type
from feast.value_type import ValueType
class KafkaOptions:
"""
DataSource Kafka options used to source features from Kafka messages
"""
def __init__(
self,
kafka_bootstrap_servers: str,
message_format: StreamFormat,
topic: str,
watermark_delay_threshold: Optional[timedelta] = None,
):
self.kafka_bootstrap_servers = kafka_bootstrap_servers
self.message_format = message_format
self.topic = topic
self.watermark_delay_threshold = watermark_delay_threshold or None
@classmethod
def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):
"""
Creates a KafkaOptions from a protobuf representation of a kafka option
Args:
kafka_options_proto: A protobuf representation of a DataSource
Returns:
Returns a KafkaOptions object based on the kafka_options protobuf
"""
watermark_delay_threshold = None
if kafka_options_proto.HasField("watermark_delay_threshold"):
watermark_delay_threshold = (
timedelta(days=0)
if kafka_options_proto.watermark_delay_threshold.ToNanoseconds() == 0
else kafka_options_proto.watermark_delay_threshold.ToTimedelta()
)
kafka_options = cls(
kafka_bootstrap_servers=kafka_options_proto.kafka_bootstrap_servers,
message_format=StreamFormat.from_proto(kafka_options_proto.message_format),
topic=kafka_options_proto.topic,
watermark_delay_threshold=watermark_delay_threshold,
)
return kafka_options
def to_proto(self) -> DataSourceProto.KafkaOptions:
"""
Converts an KafkaOptionsProto object to its protobuf representation.
Returns:
KafkaOptionsProto protobuf
"""
watermark_delay_threshold = None
if self.watermark_delay_threshold is not None:
watermark_delay_threshold = Duration()
watermark_delay_threshold.FromTimedelta(self.watermark_delay_threshold)
kafka_options_proto = DataSourceProto.KafkaOptions(
kafka_bootstrap_servers=self.kafka_bootstrap_servers,
message_format=self.message_format.to_proto(),
topic=self.topic,
watermark_delay_threshold=watermark_delay_threshold,
)
return kafka_options_proto
class KinesisOptions:
"""
DataSource Kinesis options used to source features from Kinesis records
"""
def __init__(
self,
record_format: StreamFormat,
region: str,
stream_name: str,
):
self.record_format = record_format
self.region = region
self.stream_name = stream_name
@classmethod
def from_proto(cls, kinesis_options_proto: DataSourceProto.KinesisOptions):
"""
Creates a KinesisOptions from a protobuf representation of a kinesis option
Args:
kinesis_options_proto: A protobuf representation of a DataSource
Returns:
Returns a KinesisOptions object based on the kinesis_options protobuf
"""
kinesis_options = cls(
record_format=StreamFormat.from_proto(kinesis_options_proto.record_format),
region=kinesis_options_proto.region,
stream_name=kinesis_options_proto.stream_name,
)
return kinesis_options
def to_proto(self) -> DataSourceProto.KinesisOptions:
"""
Converts an KinesisOptionsProto object to its protobuf representation.
Returns:
KinesisOptionsProto protobuf
"""
kinesis_options_proto = DataSourceProto.KinesisOptions(
record_format=self.record_format.to_proto(),
region=self.region,
stream_name=self.stream_name,
)
return kinesis_options_proto
_DATA_SOURCE_OPTIONS = {
DataSourceProto.SourceType.BATCH_FILE: "feast.infra.offline_stores.file_source.FileSource",
DataSourceProto.SourceType.BATCH_BIGQUERY: "feast.infra.offline_stores.bigquery_source.BigQuerySource",
DataSourceProto.SourceType.BATCH_REDSHIFT: "feast.infra.offline_stores.redshift_source.RedshiftSource",
DataSourceProto.SourceType.BATCH_SNOWFLAKE: "feast.infra.offline_stores.snowflake_source.SnowflakeSource",
DataSourceProto.SourceType.BATCH_TRINO: "feast.infra.offline_stores.contrib.trino_offline_store.trino_source.TrinoSource",
DataSourceProto.SourceType.BATCH_SPARK: "feast.infra.offline_stores.contrib.spark_offline_store.spark_source.SparkSource",
DataSourceProto.SourceType.BATCH_ATHENA: "feast.infra.offline_stores.contrib.athena_offline_store.athena_source.AthenaSource",
DataSourceProto.SourceType.STREAM_KAFKA: "feast.data_source.KafkaSource",
DataSourceProto.SourceType.STREAM_KINESIS: "feast.data_source.KinesisSource",
DataSourceProto.SourceType.REQUEST_SOURCE: "feast.data_source.RequestSource",
DataSourceProto.SourceType.PUSH_SOURCE: "feast.data_source.PushSource",
}
@typechecked
class DataSource(ABC):
"""
DataSource that can be used to source features.
Args:
name: Name of data source, which should be unique within a project
timestamp_field (optional): Event timestamp field used for point-in-time joins of
feature values.
created_timestamp_column (optional): Timestamp column indicating when the row
was created, used for deduplicating rows.
field_mapping (optional): A dictionary mapping of column names in this data
source to feature names in a feature table or view. Only used for feature
columns, not entity or timestamp columns.
description (optional) A human-readable description.
tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
owner (optional): The owner of the data source, typically the email of the primary
maintainer.
date_partition_column (optional): Timestamp column used for partitioning. Not supported by all offline stores.
"""
name: str
timestamp_field: str
created_timestamp_column: str
field_mapping: Dict[str, str]
description: str
tags: Dict[str, str]
owner: str
date_partition_column: str
def __init__(
self,
*,
name: str,
timestamp_field: Optional[str] = None,
created_timestamp_column: Optional[str] = None,
field_mapping: Optional[Dict[str, str]] = None,
description: Optional[str] = "",
tags: Optional[Dict[str, str]] = None,
owner: Optional[str] = "",
date_partition_column: Optional[str] = None,
):
"""
Creates a DataSource object.
Args:
name: Name of data source, which should be unique within a project.
timestamp_field (optional): Event timestamp field used for point-in-time joins of
feature values.
created_timestamp_column (optional): Timestamp column indicating when the row
was created, used for deduplicating rows.
field_mapping (optional): A dictionary mapping of column names in this data
source to feature names in a feature table or view. Only used for feature
columns, not entity or timestamp columns.
description (optional): A human-readable description.
tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
owner (optional): The owner of the data source, typically the email of the primary
maintainer.
date_partition_column (optional): Timestamp column used for partitioning. Not supported by all stores
"""
self.name = name
self.timestamp_field = timestamp_field or ""
self.created_timestamp_column = (
created_timestamp_column if created_timestamp_column else ""
)
self.field_mapping = field_mapping if field_mapping else {}
if (
self.timestamp_field
and self.timestamp_field == self.created_timestamp_column
):
raise ValueError(
"Please do not use the same column for 'timestamp_field' and 'created_timestamp_column'."
)
self.description = description or ""
self.tags = tags or {}
self.owner = owner or ""
self.date_partition_column = (
date_partition_column if date_partition_column else ""
)
def __hash__(self):
return hash((self.name, self.timestamp_field))
def __str__(self):
return str(MessageToJson(self.to_proto()))
def __eq__(self, other):
if other is None:
return False
if not isinstance(other, DataSource):
raise TypeError("Comparisons should only involve DataSource class objects.")
if (
self.name != other.name
or self.timestamp_field != other.timestamp_field
or self.created_timestamp_column != other.created_timestamp_column
or self.field_mapping != other.field_mapping
or self.date_partition_column != other.date_partition_column
or self.description != other.description
or self.tags != other.tags
or self.owner != other.owner
):
return False
return True
@staticmethod
@abstractmethod
def from_proto(data_source: DataSourceProto) -> Any:
"""
Converts data source config in protobuf spec to a DataSource class object.
Args:
data_source: A protobuf representation of a DataSource.
Returns:
A DataSource class object.
Raises:
ValueError: The type of DataSource could not be identified.
"""
data_source_type = data_source.type
if not data_source_type or (
data_source_type
not in list(_DATA_SOURCE_OPTIONS.keys())
+ [DataSourceProto.SourceType.CUSTOM_SOURCE]
):
raise ValueError("Could not identify the source type being added.")
if data_source_type == DataSourceProto.SourceType.CUSTOM_SOURCE:
cls = get_data_source_class_from_type(data_source.data_source_class_type)
return cls.from_proto(data_source)
cls = get_data_source_class_from_type(_DATA_SOURCE_OPTIONS[data_source_type])
return cls.from_proto(data_source)
@abstractmethod
def to_proto(self) -> DataSourceProto:
"""
Converts a DataSourceProto object to its protobuf representation.
"""
raise NotImplementedError
def validate(self, config: RepoConfig):
"""
Validates the underlying data source.
Args:
config: Configuration object used to configure a feature store.
"""
raise NotImplementedError
@staticmethod
@abstractmethod
def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
"""
Returns the callable method that returns Feast type given the raw column type.
"""
raise NotImplementedError
def get_table_column_names_and_types(
self, config: RepoConfig
) -> Iterable[Tuple[str, str]]:
"""
Returns the list of column names and raw column types.
Args:
config: Configuration object used to configure a feature store.
"""
raise NotImplementedError
def get_table_query_string(self) -> str:
"""
Returns a string that can directly be used to reference this table in SQL.
"""
raise NotImplementedError
@typechecked
class KafkaSource(DataSource):
def __init__(
self,
*,
name: str,
timestamp_field: str,
message_format: StreamFormat,
bootstrap_servers: Optional[str] = None,
kafka_bootstrap_servers: Optional[str] = None,
topic: Optional[str] = None,
created_timestamp_column: Optional[str] = "",
field_mapping: Optional[Dict[str, str]] = None,
description: Optional[str] = "",
tags: Optional[Dict[str, str]] = None,
owner: Optional[str] = "",
batch_source: Optional[DataSource] = None,
watermark_delay_threshold: Optional[timedelta] = None,
):
"""
Creates a KafkaSource object.
Args:
name: Name of data source, which should be unique within a project
timestamp_field: Event timestamp field used for point-in-time joins of feature values.
message_format: StreamFormat of serialized messages.
bootstrap_servers: (Deprecated) The servers of the kafka broker in the form "localhost:9092".
kafka_bootstrap_servers (optional): The servers of the kafka broker in the form "localhost:9092".
topic (optional): The name of the topic to read from in the kafka source.
created_timestamp_column (optional): Timestamp column indicating when the row
was created, used for deduplicating rows.
field_mapping (optional): A dictionary mapping of column names in this data
source to feature names in a feature table or view. Only used for feature
columns, not entity or timestamp columns.
description (optional): A human-readable description.
tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
owner (optional): The owner of the data source, typically the email of the primary
maintainer.
batch_source (optional): The datasource that acts as a batch source.
watermark_delay_threshold (optional): The watermark delay threshold for stream data.
Specifically how late stream data can arrive without being discarded.
"""
if bootstrap_servers:
warnings.warn(
(
"The 'bootstrap_servers' parameter has been deprecated in favor of 'kafka_bootstrap_servers'. "
"Feast 0.25 and onwards will not support the 'bootstrap_servers' parameter."
),
DeprecationWarning,
)
super().__init__(
name=name,
timestamp_field=timestamp_field,
created_timestamp_column=created_timestamp_column,
field_mapping=field_mapping,
description=description,
tags=tags,
owner=owner,
)
self.batch_source = batch_source
kafka_bootstrap_servers = kafka_bootstrap_servers or bootstrap_servers or ""
topic = topic or ""
self.kafka_options = KafkaOptions(
kafka_bootstrap_servers=kafka_bootstrap_servers,
message_format=message_format,
topic=topic,
watermark_delay_threshold=watermark_delay_threshold,
)
def __eq__(self, other):
if not isinstance(other, KafkaSource):
raise TypeError(
"Comparisons should only involve KafkaSource class objects."
)
if not super().__eq__(other):
return False
if (
self.kafka_options.kafka_bootstrap_servers
!= other.kafka_options.kafka_bootstrap_servers
or self.kafka_options.message_format != other.kafka_options.message_format
or self.kafka_options.topic != other.kafka_options.topic
or self.kafka_options.watermark_delay_threshold
!= other.kafka_options.watermark_delay_threshold
):
return False
return True
def __hash__(self):
return super().__hash__()
@staticmethod
def from_proto(data_source: DataSourceProto):
watermark_delay_threshold = None
if data_source.kafka_options.watermark_delay_threshold:
watermark_delay_threshold = (
timedelta(days=0)
if data_source.kafka_options.watermark_delay_threshold.ToNanoseconds()
== 0
else data_source.kafka_options.watermark_delay_threshold.ToTimedelta()
)
return KafkaSource(
name=data_source.name,
field_mapping=dict(data_source.field_mapping),
kafka_bootstrap_servers=data_source.kafka_options.kafka_bootstrap_servers,
message_format=StreamFormat.from_proto(
data_source.kafka_options.message_format
),
watermark_delay_threshold=watermark_delay_threshold,
topic=data_source.kafka_options.topic,
created_timestamp_column=data_source.created_timestamp_column,
timestamp_field=data_source.timestamp_field,
description=data_source.description,
tags=dict(data_source.tags),
owner=data_source.owner,
batch_source=DataSource.from_proto(data_source.batch_source)
if data_source.batch_source
else None,
)
def to_proto(self) -> DataSourceProto:
data_source_proto = DataSourceProto(
name=self.name,
type=DataSourceProto.STREAM_KAFKA,
field_mapping=self.field_mapping,
kafka_options=self.kafka_options.to_proto(),
description=self.description,
tags=self.tags,
owner=self.owner,
)
data_source_proto.timestamp_field = self.timestamp_field
data_source_proto.created_timestamp_column = self.created_timestamp_column
if self.batch_source:
data_source_proto.batch_source.MergeFrom(self.batch_source.to_proto())
return data_source_proto
def validate(self, config: RepoConfig):
raise NotImplementedError
def get_table_column_names_and_types(
self, config: RepoConfig
) -> Iterable[Tuple[str, str]]:
raise NotImplementedError
@staticmethod
def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
return type_map.redshift_to_feast_value_type
def get_table_query_string(self) -> str:
raise NotImplementedError
@typechecked
class RequestSource(DataSource):
"""
RequestSource that can be used to provide input features for on demand transforms
Attributes:
name: Name of the request data source
schema: Schema mapping from the input feature name to a ValueType
description: A human-readable description.
tags: A dictionary of key-value pairs to store arbitrary metadata.
owner: The owner of the request data source, typically the email of the primary
maintainer.
"""
name: str
schema: List[Field]
description: str
tags: Dict[str, str]
owner: str
def __init__(
self,
*,
name: str,
schema: List[Field],
description: Optional[str] = "",
tags: Optional[Dict[str, str]] = None,
owner: Optional[str] = "",
):
"""Creates a RequestSource object."""
super().__init__(name=name, description=description, tags=tags, owner=owner)
self.schema = schema
def validate(self, config: RepoConfig):
raise NotImplementedError
def get_table_column_names_and_types(
self, config: RepoConfig
) -> Iterable[Tuple[str, str]]:
raise NotImplementedError
def __eq__(self, other):
if not isinstance(other, RequestSource):
raise TypeError(
"Comparisons should only involve RequestSource class objects."
)
if not super().__eq__(other):
return False
if isinstance(self.schema, List) and isinstance(other.schema, List):
for field1, field2 in zip(self.schema, other.schema):
if field1 != field2:
return False
return True
else:
return False
def __hash__(self):
return super().__hash__()
@staticmethod
def from_proto(data_source: DataSourceProto):
schema_pb = data_source.request_data_options.schema
list_schema = []
for field_proto in schema_pb:
list_schema.append(Field.from_proto(field_proto))
return RequestSource(
name=data_source.name,
schema=list_schema,
description=data_source.description,
tags=dict(data_source.tags),
owner=data_source.owner,
)
def to_proto(self) -> DataSourceProto:
schema_pb = []
if isinstance(self.schema, Dict):
for key, value in self.schema.items():
schema_pb.append(
Field(name=key, dtype=from_value_type(value.value)).to_proto()
)
else:
for field in self.schema:
schema_pb.append(field.to_proto())
data_source_proto = DataSourceProto(
name=self.name,
type=DataSourceProto.REQUEST_SOURCE,
description=self.description,
tags=self.tags,
owner=self.owner,
)
data_source_proto.request_data_options.schema.extend(schema_pb)
return data_source_proto
def get_table_query_string(self) -> str:
raise NotImplementedError
@staticmethod
def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
raise NotImplementedError
@typechecked
class KinesisSource(DataSource):
def validate(self, config: RepoConfig):
raise NotImplementedError
def get_table_column_names_and_types(
self, config: RepoConfig
) -> Iterable[Tuple[str, str]]:
raise NotImplementedError
@staticmethod
def from_proto(data_source: DataSourceProto):
return KinesisSource(
name=data_source.name,
timestamp_field=data_source.timestamp_field,
field_mapping=dict(data_source.field_mapping),
record_format=StreamFormat.from_proto(
data_source.kinesis_options.record_format
),
region=data_source.kinesis_options.region,
stream_name=data_source.kinesis_options.stream_name,
created_timestamp_column=data_source.created_timestamp_column,
description=data_source.description,
tags=dict(data_source.tags),
owner=data_source.owner,
batch_source=DataSource.from_proto(data_source.batch_source)
if data_source.batch_source
else None,
)
@staticmethod
def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
raise NotImplementedError
def get_table_query_string(self) -> str:
raise NotImplementedError
def __init__(
self,
*,
name: str,
record_format: StreamFormat,
region: str,
stream_name: str,
timestamp_field: Optional[str] = "",
created_timestamp_column: Optional[str] = "",
field_mapping: Optional[Dict[str, str]] = None,
description: Optional[str] = "",
tags: Optional[Dict[str, str]] = None,
owner: Optional[str] = "",
batch_source: Optional[DataSource] = None,
):
if record_format is None:
raise ValueError("Record format must be specified for kinesis source")
super().__init__(
name=name,
timestamp_field=timestamp_field,
created_timestamp_column=created_timestamp_column,
field_mapping=field_mapping,
description=description,
tags=tags,
owner=owner,
)
self.batch_source = batch_source
self.kinesis_options = KinesisOptions(
record_format=record_format, region=region, stream_name=stream_name
)
def __eq__(self, other):
if not isinstance(other, KinesisSource):
raise TypeError(
"Comparisons should only involve KinesisSource class objects."
)
if not super().__eq__(other):
return False
if (
self.kinesis_options.record_format != other.kinesis_options.record_format
or self.kinesis_options.region != other.kinesis_options.region
or self.kinesis_options.stream_name != other.kinesis_options.stream_name
):
return False
return True
def __hash__(self):
return super().__hash__()
def to_proto(self) -> DataSourceProto:
data_source_proto = DataSourceProto(
name=self.name,
type=DataSourceProto.STREAM_KINESIS,
field_mapping=self.field_mapping,
kinesis_options=self.kinesis_options.to_proto(),
description=self.description,
tags=self.tags,
owner=self.owner,
)
data_source_proto.timestamp_field = self.timestamp_field
data_source_proto.created_timestamp_column = self.created_timestamp_column
if self.batch_source:
data_source_proto.batch_source.MergeFrom(self.batch_source.to_proto())
return data_source_proto
class PushMode(enum.Enum):
ONLINE = 1
OFFLINE = 2
ONLINE_AND_OFFLINE = 3
@typechecked
class PushSource(DataSource):
"""
A source that can be used to ingest features on request
"""
# TODO(adchia): consider adding schema here in case where Feast manages pushing events to the offline store
# TODO(adchia): consider a "mode" to support pushing raw vs transformed events
batch_source: DataSource
def __init__(
self,
*,
name: str,
batch_source: DataSource,
description: Optional[str] = "",
tags: Optional[Dict[str, str]] = None,
owner: Optional[str] = "",
):
"""
Creates a PushSource object.
Args:
name: Name of the push source
batch_source: The batch source that backs this push source. It's used when materializing from the offline
store to the online store, and when retrieving historical features.
description (optional): A human-readable description.
tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
owner (optional): The owner of the data source, typically the email of the primary
maintainer.
"""
super().__init__(name=name, description=description, tags=tags, owner=owner)
self.batch_source = batch_source
def __eq__(self, other):
if not isinstance(other, PushSource):
return False
if not super().__eq__(other):
return False
if self.batch_source != other.batch_source:
return False
return True
def __hash__(self):
return super().__hash__()
def validate(self, config: RepoConfig):
raise NotImplementedError
def get_table_column_names_and_types(
self, config: RepoConfig
) -> Iterable[Tuple[str, str]]:
raise NotImplementedError
@staticmethod
def from_proto(data_source: DataSourceProto):
assert data_source.HasField("batch_source")
batch_source = DataSource.from_proto(data_source.batch_source)
return PushSource(
name=data_source.name,
batch_source=batch_source,
description=data_source.description,
tags=dict(data_source.tags),
owner=data_source.owner,
)
def to_proto(self) -> DataSourceProto:
batch_source_proto = None
if self.batch_source:
batch_source_proto = self.batch_source.to_proto()
data_source_proto = DataSourceProto(
name=self.name,
type=DataSourceProto.PUSH_SOURCE,
description=self.description,
tags=self.tags,
owner=self.owner,
batch_source=batch_source_proto,
)
return data_source_proto
def get_table_query_string(self) -> str:
raise NotImplementedError
@staticmethod
def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
raise NotImplementedError