fixed bugs, cleaned code, added AthenaDataSourceCreator

Signed-off-by: Youngkyu OH <toping4445@gmail.com>
feast-dev · feast-ci-bot · Aug 10, 2022 · Jul 24, 2022 · Jul 31, 2022 · Aug 7, 2022
commit 693e9dfec5e9fcbc0b574e040805a2b093aa6661
@@ -49,6 +49,7 @@ message DataSource {
     PUSH_SOURCE = 9;
     BATCH_TRINO = 10;
     BATCH_SPARK = 11;
+    BATCH_ATHENA = 12;
   }
 
   // Unique name of data source within the project
@@ -171,6 +172,22 @@ message DataSource {
     string database = 4;
   }
 
+  // Defines options for DataSource that sources features from a Athena Query
+  message AthenaOptions {
+    // Athena table name
+    string table = 1;
+
+    // SQL query that returns a table containing feature data. Must contain an event_timestamp column, and respective
+    // entity columns
+    string query = 2;
+
+    // Athena database name
+    string database = 3;
+
+    // Athena schema name
+    string data_source = 4;
+  }
+
   // Defines options for DataSource that sources features from a Snowflake Query
   message SnowflakeOptions {
     // Snowflake table name
@@ -242,5 +259,6 @@ message DataSource {
     PushOptions push_options = 22;
     SparkOptions spark_options = 27;
     TrinoOptions trino_options = 30;
+    AthenaOptions athena_options = 35;
   }
 }
@@ -60,6 +60,7 @@ message LoggingConfig {
     RedshiftDestination redshift_destination = 5;
     SnowflakeDestination snowflake_destination = 6;
     CustomDestination custom_destination = 7;
+    AthenaDestination athena_destination = 8;
   }
 
   message FileDestination {
@@ -80,6 +81,11 @@ message LoggingConfig {
     string table_name = 1;
   }
 
+  message AthenaDestination {
+    // Destination table name. data_source and database will be taken from an offline store config
+    string table_name = 1;
+  }
+
   message SnowflakeDestination {
     // Destination table name. Schema and database will be taken from an offline store config
     string table_name = 1;

@@ -59,6 +59,7 @@ message SavedDatasetStorage {
     DataSource.TrinoOptions trino_storage = 8;
     DataSource.SparkOptions spark_storage = 9;
     DataSource.CustomSourceOptions custom_storage = 10;
+    DataSource.AthenaOptions athena_storage = 11;
   }
 }
 

@@ -8,6 +8,9 @@
 from feast.infra.offline_stores.file_source import FileSource
 from feast.infra.offline_stores.redshift_source import RedshiftSource
 from feast.infra.offline_stores.snowflake_source import SnowflakeSource
+from feast.infra.offline_stores.contrib.athena_offline_store.athena_source import (
+    AthenaSource,
+)
 
 from .batch_feature_view import BatchFeatureView
 from .data_source import KafkaSource, KinesisSource, PushSource, RequestSource
@@ -50,4 +53,5 @@
     "SnowflakeSource",
     "PushSource",
     "RequestSource",
+    "AthenaSource",
 ]
@@ -160,6 +160,7 @@ def to_proto(self) -> DataSourceProto.KinesisOptions:
     DataSourceProto.SourceType.STREAM_KINESIS: "feast.data_source.KinesisSource",
     DataSourceProto.SourceType.REQUEST_SOURCE: "feast.data_source.RequestSource",
     DataSourceProto.SourceType.PUSH_SOURCE: "feast.data_source.PushSource",
+    DataSourceProto.SourceType.BATCH_ATHENA: "feast.infra.offline_stores.contrib.athena_offline_store.athena_source.AthenaSource",
 }