Skip to content

Commit da67d7d

Browse files
HaoXuAIHaoXuAI
committed
fix lint
Signed-off-by: HaoXuAI <sduxuhao@gmail.com>
1 parent df993df commit da67d7d

7 files changed

Lines changed: 75 additions & 80 deletions

File tree

‎sdk/python/feast/batch_feature_view.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
import functools
22
import warnings
33
from datetime import datetime, timedelta
4-
from types import FunctionType
5-
from typing import Dict, List, Optional, Tuple, Union, Callable, Any
4+
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
65

76
import dill
87

‎sdk/python/feast/infra/compute_engines/spark/config.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
from typing import Optional, Dict
1+
from typing import Dict, Optional
22

3-
from feast.repo_config import FeastConfigBaseModel
43
from pydantic import StrictStr
54

5+
from feast.repo_config import FeastConfigBaseModel
6+
67

78
class SparkComputeConfig(FeastConfigBaseModel):
89
type: StrictStr = "spark"

‎sdk/python/feast/infra/compute_engines/spark/utils.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
from typing import Optional, Dict
1+
from typing import Dict, Optional
22

33
from pyspark import SparkConf
44
from pyspark.sql import SparkSession
55

66

77
def get_or_create_new_spark_session(
8-
spark_config: Optional[Dict[str, str]] = None
8+
spark_config: Optional[Dict[str, str]] = None,
99
) -> SparkSession:
1010
spark_session = SparkSession.getActiveSession()
1111
if not spark_session:

‎sdk/python/feast/stream_feature_view.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ def get_feature_transformation(self) -> Optional[Transformation]:
152152
TransformationMode.PANDAS,
153153
TransformationMode.PYTHON,
154154
TransformationMode.SPARK_SQL,
155-
TransformationMode.SPARK
155+
TransformationMode.SPARK,
156156
) or self.mode in ("pandas", "python", "spark_sql", "spark"):
157157
return Transformation(
158158
mode=self.mode, udf=self.udf, udf_string=self.udf_string or ""

‎sdk/python/feast/transformation/spark_transformation.py‎

Lines changed: 40 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,27 @@
1-
from typing import Any, Union, Dict, Optional, cast
1+
from typing import Any, Dict, Optional, Union, cast
22

33
import pandas as pd
44
import pyspark.sql
55

6+
from feast.infra.compute_engines.spark.utils import get_or_create_new_spark_session
67
from feast.transformation.base import Transformation
78
from feast.transformation.mode import TransformationMode
8-
from feast.infra.compute_engines.spark.utils import get_or_create_new_spark_session
99

1010

1111
class SparkTransformation(Transformation):
12-
13-
def __new__(cls,
14-
mode: Union[TransformationMode, str],
15-
udf: Any,
16-
udf_string: str,
17-
spark_config: Dict[str, Any] = {},
18-
name: Optional[str] = None,
19-
tags: Optional[Dict[str, str]] = None,
20-
description: str = "",
21-
owner: str = "",
22-
*args,
23-
**kwargs) -> "SparkTransformation":
12+
def __new__(
13+
cls,
14+
mode: Union[TransformationMode, str],
15+
udf: Any,
16+
udf_string: str,
17+
spark_config: Dict[str, Any] = {},
18+
name: Optional[str] = None,
19+
tags: Optional[Dict[str, str]] = None,
20+
description: str = "",
21+
owner: str = "",
22+
*args,
23+
**kwargs,
24+
) -> "SparkTransformation":
2425
instance = super(SparkTransformation, cls).__new__(
2526
cls,
2627
mode=mode,
@@ -34,17 +35,19 @@ def __new__(cls,
3435
)
3536
return cast(SparkTransformation, instance)
3637

37-
def __init__(self,
38-
mode: Union[TransformationMode, str],
39-
udf: Any,
40-
udf_string: str,
41-
spark_config: Dict[str, Any] = {},
42-
name: Optional[str] = None,
43-
tags: Optional[Dict[str, str]] = None,
44-
description: str = "",
45-
owner: str = "",
46-
*args,
47-
**kwargs):
38+
def __init__(
39+
self,
40+
mode: Union[TransformationMode, str],
41+
udf: Any,
42+
udf_string: str,
43+
spark_config: Dict[str, Any] = {},
44+
name: Optional[str] = None,
45+
tags: Optional[Dict[str, str]] = None,
46+
description: str = "",
47+
owner: str = "",
48+
*args,
49+
**kwargs,
50+
):
4851
super().__init__(
4952
mode=mode,
5053
udf=udf,
@@ -56,36 +59,34 @@ def __init__(self,
5659
)
5760
self.spark_session = get_or_create_new_spark_session(spark_config)
5861

59-
def transform(self,
60-
*inputs: Union[str, pd.DataFrame],
61-
) -> pd.DataFrame:
62+
def transform(
63+
self,
64+
*inputs: Union[str, pd.DataFrame],
65+
) -> pd.DataFrame:
6266
if self.mode == TransformationMode.SPARK_SQL:
6367
return self._transform_spark_sql(*inputs)
6468
else:
6569
return self._transform_spark_udf(*inputs)
6670

6771
@staticmethod
68-
def _create_temp_view_for_dataframe(df: pyspark.sql.DataFrame,
69-
name: str):
72+
def _create_temp_view_for_dataframe(df: pyspark.sql.DataFrame, name: str):
7073
df_temp_view = f"feast_transformation_temp_view_{name}"
7174
df.createOrReplaceTempView(df_temp_view)
7275
return df_temp_view
7376

74-
def _transform_spark_sql(self,
75-
*inputs: Union[pyspark.sql.DataFrame, str]
76-
) -> pd.DataFrame:
77+
def _transform_spark_sql(
78+
self, *inputs: Union[pyspark.sql.DataFrame, str]
79+
) -> pd.DataFrame:
7780
inputs_str = [
7881
self._create_temp_view_for_dataframe(v, f"index_{i}")
79-
if isinstance(v, pyspark.sql.DataFrame) else v
82+
if isinstance(v, pyspark.sql.DataFrame)
83+
else v
8084
for i, v in enumerate(inputs)
8185
]
8286
return self.spark_session.sql(self.udf(*inputs_str))
8387

84-
def _transform_spark_udf(self,
85-
*inputs: Any) -> pd.DataFrame:
88+
def _transform_spark_udf(self, *inputs: Any) -> pd.DataFrame:
8689
return self.udf(*inputs)
8790

88-
def infer_features(self,
89-
*args,
90-
**kwargs) -> Any:
91+
def infer_features(self, *args, **kwargs) -> Any:
9192
pass

‎sdk/python/tests/unit/transformation/test_pandas_transformation.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
1-
from feast.transformation.pandas_transformation import PandasTransformation
21
import pandas as pd
32

3+
from feast.transformation.pandas_transformation import PandasTransformation
4+
45

56
def pandas_udf(features_df: pd.DataFrame) -> pd.DataFrame:
67
df = pd.DataFrame()
@@ -10,14 +11,8 @@ def pandas_udf(features_df: pd.DataFrame) -> pd.DataFrame:
1011

1112

1213
def test_init_pandas_transformation():
13-
transformation = PandasTransformation(
14-
udf=pandas_udf,
15-
udf_string="udf1"
16-
)
17-
features_df = pd.DataFrame.from_dict({
18-
"feature1": [1, 2],
19-
"feature2": [2, 3]
20-
})
14+
transformation = PandasTransformation(udf=pandas_udf, udf_string="udf1")
15+
features_df = pd.DataFrame.from_dict({"feature1": [1, 2], "feature2": [2, 3]})
2116
transformed_df = transformation.transform(features_df)
2217
assert transformed_df["output1"].values[0] == 1
2318
assert transformed_df["output2"].values[1] == 3

‎sdk/python/tests/unit/transformation/test_spark_transformation.py‎

Lines changed: 24 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,46 @@
1+
from unittest.mock import patch
2+
13
import pytest
24
from pyspark.sql import SparkSession
35
from pyspark.sql.functions import col, regexp_replace
4-
from unittest.mock import patch
56
from pyspark.testing.utils import assertDataFrameEqual
67

7-
from feast.transformation.spark_transformation import SparkTransformation
8-
from feast.transformation.mode import TransformationMode
98
from feast.transformation.base import Transformation
9+
from feast.transformation.mode import TransformationMode
10+
from feast.transformation.spark_transformation import SparkTransformation
1011

1112

1213
def get_sample_df(spark):
13-
sample_data = [{"name": "John D.", "age": 30},
14-
{"name": "Alice G.", "age": 25},
15-
{"name": "Bob T.", "age": 35},
16-
{"name": "Eve A.", "age": 28}]
14+
sample_data = [
15+
{"name": "John D.", "age": 30},
16+
{"name": "Alice G.", "age": 25},
17+
{"name": "Bob T.", "age": 35},
18+
{"name": "Eve A.", "age": 28},
19+
]
1720
df = spark.createDataFrame(sample_data)
1821
return df
1922

2023

2124
def get_expected_df(spark):
22-
expected_data = [{"name": "John D.", "age": 30},
23-
{"name": "Alice G.", "age": 25},
24-
{"name": "Bob T.", "age": 35},
25-
{"name": "Eve A.", "age": 28}]
25+
expected_data = [
26+
{"name": "John D.", "age": 30},
27+
{"name": "Alice G.", "age": 25},
28+
{"name": "Bob T.", "age": 35},
29+
{"name": "Eve A.", "age": 28},
30+
]
2631

2732
expected_df = spark.createDataFrame(expected_data)
2833
return expected_df
2934

3035

31-
def remove_extra_spaces(df,
32-
column_name):
33-
df_transformed = df.withColumn(column_name, regexp_replace(col(column_name), "\\s+", " "))
36+
def remove_extra_spaces(df, column_name):
37+
df_transformed = df.withColumn(
38+
column_name, regexp_replace(col(column_name), "\\s+", " ")
39+
)
3440
return df_transformed
3541

3642

37-
def remove_extra_spaces_sql(df,
38-
column_name):
43+
def remove_extra_spaces_sql(df, column_name):
3944
sql = f"""
4045
SELECT
4146
age,
@@ -51,9 +56,7 @@ def spark_fixture():
5156
yield spark
5257

5358

54-
@patch(
55-
"feast.infra.compute_engines.spark.utils.get_or_create_new_spark_session"
56-
)
59+
@patch("feast.infra.compute_engines.spark.utils.get_or_create_new_spark_session")
5760
def test_spark_transformation(spark_fixture):
5861
spark = SparkSession.builder.appName("Testing PySpark Example").getOrCreate()
5962
df = get_sample_df(spark)
@@ -69,9 +72,7 @@ def test_spark_transformation(spark_fixture):
6972
assertDataFrameEqual(transformed_df, expected_df)
7073

7174

72-
@patch(
73-
"feast.infra.compute_engines.spark.utils.get_or_create_new_spark_session"
74-
)
75+
@patch("feast.infra.compute_engines.spark.utils.get_or_create_new_spark_session")
7576
def test_spark_transformation_init_transformation(spark_fixture):
7677
spark = SparkSession.builder.appName("Testing PySpark Example").getOrCreate()
7778
df = get_sample_df(spark)
@@ -87,9 +88,7 @@ def test_spark_transformation_init_transformation(spark_fixture):
8788
assertDataFrameEqual(transformed_df, expected_df)
8889

8990

90-
@patch(
91-
"feast.infra.compute_engines.spark.utils.get_or_create_new_spark_session"
92-
)
91+
@patch("feast.infra.compute_engines.spark.utils.get_or_create_new_spark_session")
9392
def test_spark_transformation_sql(spark_fixture):
9493
spark = SparkSession.builder.appName("Testing PySpark Example").getOrCreate()
9594
df = get_sample_df(spark)

0 commit comments

Comments
 (0)