Skip to content

fix: Simplify DataSource.from_proto logic #2424

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
feast-ci-bot merged 11 commits into from
Mar 22, 2022
Merged

fix: Simplify DataSource.from_proto logic #2424

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
07212b5
fix: Simplify DataSource.from_proto logic
achals Mar 20, 2022
e99fabe
simpler
achals Mar 21, 2022
a86fd8d
use enum instead
achals Mar 22, 2022
a6f9db2
rebase and add push
achals Mar 22, 2022
658f4e0
imports
achals Mar 22, 2022
886f97a
imports
achals Mar 22, 2022
fe3d69d
imports
achals Mar 22, 2022
772cc5c
CR
achals Mar 22, 2022
852c27e
Revert "CR"
achals Mar 22, 2022
03cf881
CR
achals Mar 22, 2022
30a3cf8
cr
achals Mar 22, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
fix: Simplify DataSource.from_proto logic 
Signed-off-by: Achal Shah <achals@gmail.com>
  • Loading branch information
@achals
achals committed Mar 22, 2022
commit 07212b574a90d8fceb5a49835f1b8584be532fea
64 changes: 31 additions & 33 deletions sdk/python/feast/data_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,22 @@ def to_proto(self) -> DataSourceProto.KinesisOptions:
return kinesis_options_proto


_BATCH_DATA_SOURCE_TYPES = {
"file_options": "feast.infra.offline_stores.file_source.FileSource",
"bigquery_options": "feast.infra.offline_stores.bigquery_source.BigQuerySource",
"redshift_options": "feast.infra.offline_stores.redshift_source.RedshiftSource",
"snowflake_options": "feast.infra.offline_stores.snowflake_source.SnowflakeSource",
}


_NON_BATCH_DATA_SOURCE_TYPES = {
"request_data_options",
"custom_options",
"kafka_options",
"kinesis_options",
}


class DataSource(ABC):
"""
DataSource that can be used to source features.
Expand Down Expand Up @@ -210,44 +226,26 @@ def from_proto(data_source: DataSourceProto) -> Any:
Raises:
ValueError: The type of DataSource could not be identified.
"""
if data_source.data_source_class_type:
cls = get_data_source_class_from_type(data_source.data_source_class_type)
return cls.from_proto(data_source)

if data_source.request_data_options and data_source.request_data_options.schema:
data_source_obj = RequestDataSource.from_proto(data_source)
elif data_source.file_options.file_format and data_source.file_options.file_url:
from feast.infra.offline_stores.file_source import FileSource

data_source_obj = FileSource.from_proto(data_source)
elif (
data_source.bigquery_options.table_ref or data_source.bigquery_options.query
option_type = data_source.WhichOneof("options")
if not option_type or (
option_type not in _BATCH_DATA_SOURCE_TYPES
and option_type not in _NON_BATCH_DATA_SOURCE_TYPES
):
from feast.infra.offline_stores.bigquery_source import BigQuerySource

data_source_obj = BigQuerySource.from_proto(data_source)
elif data_source.redshift_options.table or data_source.redshift_options.query:
from feast.infra.offline_stores.redshift_source import RedshiftSource

data_source_obj = RedshiftSource.from_proto(data_source)

elif data_source.snowflake_options.table or data_source.snowflake_options.query:
from feast.infra.offline_stores.snowflake_source import SnowflakeSource
raise ValueError("Could not identify the source type being added.")

data_source_obj = SnowflakeSource.from_proto(data_source)
if option_type in _BATCH_DATA_SOURCE_TYPES:
cls = get_data_source_class_from_type(_BATCH_DATA_SOURCE_TYPES[option_type])
return cls.from_proto(data_source)

elif (
data_source.kafka_options.bootstrap_servers
and data_source.kafka_options.topic
and data_source.kafka_options.message_format
):
if option_type == "custom_options":
cls = get_data_source_class_from_type(data_source.data_source_class_type)
data_source_obj = cls.from_proto(data_source)
elif option_type == "kafka_options":
data_source_obj = KafkaSource.from_proto(data_source)
elif (
data_source.kinesis_options.record_format
and data_source.kinesis_options.region
and data_source.kinesis_options.stream_name
):
elif option_type == "kinesis_options":
data_source_obj = KinesisSource.from_proto(data_source)
elif option_type == "request_data_options":
data_source_obj = RequestDataSource.from_proto(data_source)
else:
raise ValueError("Could not identify the source type being added.")

Expand Down