Fix lint

kevjumba · kevjumba · commit 8d92603357c7 · 2022-08-17T11:56:57.000-07:00
Signed-off-by: Kevin Zhang &lt;kzhang@tecton.ai&gt;
diff --git a/sdk/python/feast/infra/contrib/azure_provider.py b/sdk/python/feast/infra/contrib/azure_provider.py
@@ -15,14 +15,13 @@
 from feast.infra.offline_stores.offline_store import RetrievalJob
 from feast.infra.offline_stores.offline_utils import get_offline_store_from_config
 from feast.infra.online_stores.helpers import get_online_store_from_config
+from feast.infra.passthrough_provider import PassthroughProvider
 from feast.infra.provider import Provider
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.registry import BaseRegistry
 from feast.repo_config import RepoConfig
 from feast.saved_dataset import SavedDataset
-from feast.infra.passthrough_provider import PassthroughProvider
-
 from feast.utils import (
     _convert_arrow_to_proto,
     _get_column_names,
diff --git a/sdk/python/feast/infra/offline_stores/contrib/mssql_offline_store/mssql.py b/sdk/python/feast/infra/offline_stores/contrib/mssql_offline_store/mssql.py
@@ -16,7 +16,7 @@
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import sessionmaker
 
-from feast import FileSource, errors
+from feast import FileSource, entity, errors
 from feast.data_source import DataSource
 from feast.errors import InvalidEntityType
 from feast.feature_logging import LoggingConfig, LoggingSource
@@ -176,19 +176,21 @@ def get_historical_features(
         expected_join_keys = _get_join_keys(project, feature_views, registry)
         assert isinstance(config.offline_store, MsSqlServerOfflineStoreConfig)
         engine = make_engine(config.offline_store)
-        entity_df["event_timestamp"] = pandas.to_datetime(
-            entity_df["event_timestamp"], utc=True
-        ).fillna(pandas.Timestamp.now())
+        if isinstance(entity_df, pandas.DataFrame):
+            entity_df_event_timestamp_col = (
+                offline_utils.infer_event_timestamp_from_entity_df(dict(zip(list(entity_df.columns), list(entity_df.dtypes))))
+            )
+            entity_df[entity_df_event_timestamp_col] = pandas.to_datetime(
+                entity_df[entity_df_event_timestamp_col], utc=True
+            ).fillna(pandas.Timestamp.now())
 
+        # TODO: figure out how to deal with entity dataframes that are strings
         (
             table_schema,
             table_name,
         ) = _upload_entity_df_into_sqlserver_and_get_entity_schema(
             engine, config, entity_df, full_feature_names=full_feature_names
         )
-        entity_df_event_timestamp_col = (
-            offline_utils.infer_event_timestamp_from_entity_df(table_schema)
-        )
 
         _assert_expected_columns_in_sqlserver(
             expected_join_keys,
diff --git a/sdk/python/feast/infra/offline_stores/contrib/mssql_offline_store/tests/data_source.py b/sdk/python/feast/infra/offline_stores/contrib/mssql_offline_store/tests/data_source.py
@@ -10,13 +10,12 @@
 
 from feast.data_source import DataSource
 from feast.infra.offline_stores.contrib.mssql_offline_store.mssql import (
-    MsSqlServerOfflineStoreConfig,
+    MsSqlServerOfflineStoreConfig, _df_to_create_table_sql
 )
 from feast.infra.offline_stores.contrib.mssql_offline_store.mssqlserver_source import (
     MsSqlServerSource,
 )
 from feast.saved_dataset import SavedDatasetStorage
-from feast.type_map import pa_to_mssql_type
 from tests.integration.feature_repos.universal.data_source_creator import (
     DataSourceCreator,
 )
@@ -43,17 +42,6 @@ def mssql_container():
     container.stop()
 
 
-def _df_to_create_table_sql(df: pd.DataFrame, table_name: str) -> str:
-    pa_table = pa.Table.from_pandas(df)
-
-    columns = [f""""{f.name}" {pa_to_mssql_type(f.type)}""" for f in pa_table.schema]
-    return f"""
-        CREATE TABLE "{table_name}" (
-            {", ".join(columns)}
-        );
-        """
-
-
 class MsSqlDataSourceCreator(DataSourceCreator):
     tables: List[str] = []
 
@@ -106,7 +94,6 @@ def create_data_source(
         # Create table
 
         destination_name = self.get_prefixed_table_name(destination_name)
-        # _df_to_create_table_sql(df, destination_name)
         engine.execute(_df_to_create_table_sql(df, destination_name))
         # Upload dataframe to azure table
         df.to_sql(destination_name, engine, index=False, if_exists="append")
