Refactor Environment class and DataSourceCreator API, and use fixtures for datasets and data sources

sdk/python/feast/feature_view.py

@@ -325,6 +325,28 @@ def infer_features_from_batch_source(self, config: RepoConfig):
                 self.batch_source.created_timestamp_column,
             } | set(self.entities)
 
+            if (
+                self.batch_source.event_timestamp_column
+                in self.batch_source.field_mapping
+            ):
+                columns_to_exclude.add(
+                    self.batch_source.field_mapping[
+                        self.batch_source.event_timestamp_column
+                    ]
+                )
+            if (
+                self.batch_source.created_timestamp_column
+                in self.batch_source.field_mapping
+            ):
+                columns_to_exclude.add(
+                    self.batch_source.field_mapping[
+                        self.batch_source.created_timestamp_column
+                    ]
+                )
+            for e in self.entities:
+                if e in self.batch_source.field_mapping:
+                    columns_to_exclude.add(self.batch_source.field_mapping[e])
+
             for (
                 col_name,
                 col_datatype,
@@ -335,7 +357,7 @@ def infer_features_from_batch_source(self, config: RepoConfig):
                 ):
                     feature_name = (
                         self.batch_source.field_mapping[col_name]
-                        if col_name in self.batch_source.field_mapping.keys()
+                        if col_name in self.batch_source.field_mapping
                         else col_name
                     )
                     self.features.append(

sdk/python/feast/infra/offline_stores/bigquery.py

@@ -86,6 +86,7 @@ def pull_latest_from_table_or_query(
             )
             WHERE _feast_row = 1
             """
+
         return BigQueryRetrievalJob(query=query, client=client, config=config)
 
     @staticmethod

sdk/python/feast/infra/provider.py

@@ -238,9 +238,11 @@ def _get_column_names(
             reverse_field_mapping[col] if col in reverse_field_mapping.keys() else col
             for col in feature_names
         ]
+    _feature_names = set(feature_names) - set(join_keys)
+    _feature_names = _feature_names - {event_timestamp_column, created_timestamp_column}
     return (
         join_keys,
-        feature_names,
+        list(_feature_names),
         event_timestamp_column,
         created_timestamp_column,
     )

sdk/python/feast/infra/utils/aws_utils.py

@@ -82,7 +82,7 @@ class RedshiftStatementNotFinishedError(Exception):
 
 
 @retry(
-    wait=wait_exponential(multiplier=0.1, max=30),
+    wait=wait_exponential(multiplier=1, max=30),
     retry=retry_if_exception_type(RedshiftStatementNotFinishedError),
 )
 def wait_for_redshift_statement(redshift_data_client, statement: dict) -> None:

sdk/python/feast/type_map.py

@@ -13,11 +13,13 @@
 # limitations under the License.
 
 import re
+from datetime import datetime
 from typing import Any, Dict, Union
 
 import numpy as np
 import pandas as pd
 from google.protobuf.json_format import MessageToDict
+from google.protobuf.timestamp_pb2 import Timestamp
 
 from feast.protos.feast.types.Value_pb2 import (
     BoolList,
@@ -104,6 +106,8 @@ def python_type_to_feast_value_type(
         "int8": ValueType.INT32,
         "bool": ValueType.BOOL,
         "timedelta": ValueType.UNIX_TIMESTAMP,
+        "Timestamp": ValueType.UNIX_TIMESTAMP,
+        "datetime": ValueType.UNIX_TIMESTAMP,
         "datetime64[ns]": ValueType.UNIX_TIMESTAMP,
         "datetime64[ns, tz]": ValueType.UNIX_TIMESTAMP,
         "category": ValueType.STRING,
@@ -160,7 +164,8 @@ def _type_err(item, dtype):
     raise ValueError(f'Value "{item}" is of type {type(item)} not of type {dtype}')
 
 
-def _python_value_to_proto_value(feast_value_type, value) -> ProtoValue:
+# TODO(achals): Simplify this method and remove the noqa.
+def _python_value_to_proto_value(feast_value_type, value) -> ProtoValue:  # noqa: C901
     """
     Converts a Python (native, pandas) value to a Feast Proto Value based
     on a provided value type
@@ -281,6 +286,10 @@ def _python_value_to_proto_value(feast_value_type, value) -> ProtoValue:
         elif feast_value_type == ValueType.INT64:
             return ProtoValue(int64_val=int(value))
         elif feast_value_type == ValueType.UNIX_TIMESTAMP:
+            if isinstance(value, datetime):
+                return ProtoValue(int64_val=int(value.timestamp()))
+            elif isinstance(value, Timestamp):
+                return ProtoValue(int64_val=int(value.ToSeconds()))
             return ProtoValue(int64_val=int(value))
         elif feast_value_type == ValueType.FLOAT:
             return ProtoValue(float_val=float(value))

sdk/python/tests/conftest.py

@@ -18,6 +18,14 @@
 import pandas as pd
 import pytest
 
+from tests.integration.feature_repos.repo_configuration import (
+    FULL_REPO_CONFIGS,
+    construct_test_environment,
+    construct_universal_data_sources,
+    construct_universal_datasets,
+    construct_universal_entities,
+)
+
 
 def pytest_configure(config):
     if platform in ["darwin", "windows"]:
@@ -87,3 +95,24 @@ def simple_dataset_2() -> pd.DataFrame:
         ],
     }
     return pd.DataFrame.from_dict(data)
+
+
+@pytest.fixture(params=FULL_REPO_CONFIGS, scope="session")
+def environment(request):
+    with construct_test_environment(request.param) as e:
+        yield e
+
+
+@pytest.fixture(scope="session")
+def universal_data_sources(environment):
+    entities = construct_universal_entities()
+    datasets = construct_universal_datasets(
+        entities, environment.start_date, environment.end_date
+    )
+    datasources = construct_universal_data_sources(
+        datasets, environment.data_source_creator
+    )
+
+    yield entities, datasets, datasources
+
+    environment.data_source_creator.teardown()

sdk/python/tests/integration/e2e/test_universal_e2e.py

@@ -3,27 +3,45 @@
 from typing import Optional
 
 import pandas as pd
+import pytest
 from pytz import utc
 
 from feast import FeatureStore, FeatureView
-from tests.integration.feature_repos.test_repo_configuration import (
-    Environment,
-    parametrize_e2e_test,
+from tests.data.data_creator import create_dataset
+from tests.integration.feature_repos.repo_configuration import (
+    FULL_REPO_CONFIGS,
+    construct_test_environment,
+    vary_infer_feature,
 )
 from tests.integration.feature_repos.universal.entities import driver
 from tests.integration.feature_repos.universal.feature_views import driver_feature_view
 
 
-@parametrize_e2e_test
-def test_e2e_consistency(test_environment: Environment):
-    fs, fv = (
-        test_environment.feature_store,
-        driver_feature_view(test_environment.data_source),
-    )
-    entity = driver()
-    fs.apply([fv, entity])
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    "config", vary_infer_feature(FULL_REPO_CONFIGS), ids=lambda v: str(v)
+)
+def test_e2e_consistency(config):
+    df = create_dataset()
+    test_suite_name = test_e2e_consistency.__name__
+    with construct_test_environment(config, test_suite_name) as test_environment:
+        fs = test_environment.feature_store
+        infer_features = test_environment.test_repo_config.infer_features
+        _, _, data_source, _ = (
+            None,
+            df,
+            test_environment.data_source_creator.create_data_source(
+                df, fs.project, field_mapping={"ts_1": "ts", "id": "driver_id"}
+            ),
+            test_environment.data_source_creator,
+        )
+
+        fv = driver_feature_view(data_source=data_source, infer_features=infer_features)
+
+        entity = driver()
+        fs.apply([fv, entity])
 
-    run_offline_online_store_consistency_test(fs, fv)
+        run_offline_online_store_consistency_test(fs, fv)
 
 
 def check_offline_and_online_features(