Fix lint

Signed-off-by: Kevin Zhang <kzhang@tecton.ai>
feast-dev · feast-ci-bot · Aug 19, 2022 · Aug 5, 2022 · Aug 10, 2022 · Aug 10, 2022
commit 32992e364fe96665e178d0c86a69fb51d0da9bbd
@@ -21,14 +21,18 @@
 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,
     _run_pyarrow_field_mapping,
     make_tzaware,
 )
 
-# TODO: Refactor the provider code.
+DEFAULT_BATCH_SIZE = 10_000
+
+
 class AzureProvider(PassthroughProvider):
     def materialize_single_feature_view(
         self,

@@ -7,8 +7,8 @@
 
 import numpy as np
 import pandas
-import pyarrow as pa
 import pyarrow
+import pyarrow as pa
 import sqlalchemy
 from pydantic.types import StrictStr
 from pydantic.typing import Literal
@@ -29,14 +29,14 @@
     build_point_in_time_query,
     get_feature_view_query_context,
 )
-from feast.type_map import pa_to_mssql_type
-
 from feast.infra.provider import RetrievalJob
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.registry import BaseRegistry
 from feast.repo_config import FeastBaseModel, RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
+from feast.type_map import pa_to_mssql_type
 from feast.usage import log_exceptions_and_usage
+
 # Make sure warning doesn't raise more than once.
 warnings.simplefilter("once", RuntimeWarning)
 
@@ -176,7 +176,9 @@ 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())
+        entity_df["event_timestamp"] = pandas.to_datetime(
+            entity_df["event_timestamp"], utc=True
+        ).fillna(pandas.Timestamp.now())
 
         (
             table_schema,
@@ -337,7 +339,7 @@ def full_feature_names(self) -> bool:
     @property
     def on_demand_feature_views(self) -> List[OnDemandFeatureView]:
         return self._on_demand_feature_views
-    
+
     def _to_df_internal(self) -> pandas.DataFrame:
         return pandas.read_sql(self.query, con=self.engine).fillna(value=np.nan)
 
@@ -405,17 +407,18 @@ def _upload_entity_df_into_sqlserver_and_get_entity_schema(
     elif isinstance(entity_df, pandas.DataFrame):
         # Drop the index so that we don't have unnecessary columns
         engine.execute(_df_to_create_table_sql(entity_df, table_id))
-        entity_df.to_sql(name=table_id, con=engine, index=False, if_exists='append')
+        entity_df.to_sql(name=table_id, con=engine, index=False, if_exists="append")
         entity_schema = dict(zip(entity_df.columns, entity_df.dtypes)), table_id
 
     else:
         raise ValueError(
             f"The entity dataframe you have provided must be a SQL Server SQL query,"
             f" or a Pandas dataframe. But we found: {type(entity_df)} "
-    )
+        )
 
     return entity_schema
 
+
 def _df_to_create_table_sql(df: pandas.DataFrame, table_name: str) -> str:
     pa_table = pa.Table.from_pandas(df)
 
@@ -427,6 +430,7 @@ def _df_to_create_table_sql(df: pandas.DataFrame, table_name: str) -> str:
         );
         """
 
+
 def _get_entity_df_event_timestamp_range(
     entity_df: Union[pandas.DataFrame, str],
     entity_df_event_timestamp_col: str,

@@ -5,8 +5,8 @@
 import pytest
 from sqlalchemy import create_engine
 from testcontainers.core.container import DockerContainer
-from testcontainers.mssql import SqlServerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
+from testcontainers.mssql import SqlServerContainer
 
 from feast.data_source import DataSource
 from feast.infra.offline_stores.contrib.mssql_offline_store.mssql import (
@@ -54,8 +54,6 @@ def _df_to_create_table_sql(df: pd.DataFrame, table_name: str) -> str:
         """
 
 
-
-
 class MsSqlDataSourceCreator(DataSourceCreator):
     tables: List[str] = []
 
@@ -65,9 +63,9 @@ def __init__(
         super().__init__(project_name)
         self.tables_created: List[str] = []
         self.container = SqlServerContainer(user=MSSQL_USER, password=MSSQL_PASSWORD)
-        #self.container = fixture_request.getfixturevalue("mssql_container")
+        # self.container = fixture_request.getfixturevalue("mssql_container")
         self.container.start()
-        #self.exposed_port = self.container.get_exposed_port("1433")
+        # self.exposed_port = self.container.get_exposed_port("1433")
         if not self.container:
             raise RuntimeError(
                 "In order to use this data source "
@@ -94,20 +92,24 @@ def create_data_source(
         **kwargs,
     ) -> DataSource:
         if timestamp_field in df:
-            df[timestamp_field] = pd.to_datetime(df[timestamp_field], utc=True).fillna(pd.Timestamp.now())
+            df[timestamp_field] = pd.to_datetime(df[timestamp_field], utc=True).fillna(
+                pd.Timestamp.now()
+            )
             # Make sure the field mapping is correct and convert the datetime datasources.
         if created_timestamp_column in df:
-            df[created_timestamp_column] = pd.to_datetime(df[created_timestamp_column], utc=True).fillna(pd.Timestamp.now())
-
+            df[created_timestamp_column] = pd.to_datetime(
+                df[created_timestamp_column], utc=True
+            ).fillna(pd.Timestamp.now())
+
         connection_string = self.create_offline_store_config().connection_string
         engine = create_engine(connection_string)
         # Create table
-        
+
         destination_name = self.get_prefixed_table_name(destination_name)
-        #_df_to_create_table_sql(df, 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')
+        df.to_sql(destination_name, engine, index=False, if_exists="append")
         self.tables.append(destination_name)
         return MsSqlServerSource(
             name="ci_mssql_source",

@@ -391,7 +391,7 @@ def validate_dataframes(expected_df, actual_df, keys, event_timestamp=None):
         .drop_duplicates()
         .reset_index(drop=True)
     )
-    
+
     new_df = expected_df.drop("event_timestamp", axis=1)
     new_actual_df = actual_df.drop("event_timestamp", axis=1)
     keys = keys.remove("event_timestamp")