Repository navigation
Expand file tree
/
Copy pathspark_transformation.py
More file actions
125 lines (114 loc) · 4.16 KB
/
Copy pathspark_transformation.py
File metadata and controls
125 lines (114 loc) · 4.16 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
from typing import Any, Dict, Optional, Union, cast
import pandas as pd
import pyspark.sql
from feast.infra.compute_engines.spark.utils import get_or_create_new_spark_session
from feast.transformation.base import Transformation
from feast.transformation.mode import TransformationMode
class SparkTransformation(Transformation):
r"""
SparkTransformation can be used to define a transformation using a Spark UDF or SQL query.
The current spark session will be used or a new one will be created if not available.
E.g.:
spark_transformation = SparkTransformation(
mode=TransformationMode.SPARK,
udf=remove_extra_spaces,
udf_string="remove extra spaces",
)
OR
spark_transformation = Transformation(
mode=TransformationMode.SPARK_SQL,
udf=remove_extra_spaces_sql,
udf_string="remove extra spaces sql",
)
OR
@transformation(mode=TransformationMode.SPARK)
def remove_extra_spaces_udf(df: pd.DataFrame) -> pd.DataFrame:
return df.assign(name=df['name'].str.replace('\s+', ' '))
"""
def __new__(
cls,
mode: Union[TransformationMode, str],
udf: Any,
udf_string: str,
spark_config: Dict[str, Any] = {},
name: Optional[str] = None,
tags: Optional[Dict[str, str]] = None,
description: str = "",
owner: str = "",
*args,
**kwargs,
) -> "SparkTransformation":
"""
Creates a SparkTransformation
Args:
mode: (required) The mode of the transformation. Choose one from TransformationMode.SPARK or TransformationMode.SPARK_SQL.
udf: (required) The user-defined transformation function.
udf_string: (required) The string representation of the udf. The dill get source doesn't
spark_config: (optional) The spark configuration to use for the transformation.
name: (optional) The name of the transformation.
tags: (optional) Metadata tags for the transformation.
description: (optional) A description of the transformation.
owner: (optional) The owner of the transformation.
"""
instance = super(SparkTransformation, cls).__new__(
cls,
mode=mode,
spark_config=spark_config,
udf=udf,
udf_string=udf_string,
name=name,
tags=tags,
description=description,
owner=owner,
)
return cast(SparkTransformation, instance)
def __init__(
self,
mode: Union[TransformationMode, str],
udf: Any,
udf_string: str,
spark_config: Dict[str, Any] = {},
name: Optional[str] = None,
tags: Optional[Dict[str, str]] = None,
description: str = "",
owner: str = "",
*args,
**kwargs,
):
super().__init__(
mode=mode,
udf=udf,
name=name,
udf_string=udf_string,
tags=tags,
description=description,
owner=owner,
)
self.spark_session = get_or_create_new_spark_session(spark_config)
def transform(
self,
*inputs: Union[str, pd.DataFrame],
) -> pd.DataFrame:
if self.mode == TransformationMode.SPARK_SQL:
return self._transform_spark_sql(*inputs)
else:
return self._transform_spark_udf(*inputs)
@staticmethod
def _create_temp_view_for_dataframe(df: pyspark.sql.DataFrame, name: str):
df_temp_view = f"feast_transformation_temp_view_{name}"
df.createOrReplaceTempView(df_temp_view)
return df_temp_view
def _transform_spark_sql(
self, *inputs: Union[pyspark.sql.DataFrame, str]
) -> pd.DataFrame:
inputs_str = [
self._create_temp_view_for_dataframe(v, f"index_{i}")
if isinstance(v, pyspark.sql.DataFrame)
else v
for i, v in enumerate(inputs)
]
return self.spark_session.sql(self.udf(*inputs_str))
def _transform_spark_udf(self, *inputs: Any) -> pd.DataFrame:
return self.udf(*inputs)
def infer_features(self, *args, **kwargs) -> Any:
pass