forked from mage-ai/mage-ai
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathfill_in_missing_values
executable file
·45 lines (33 loc) · 1.29 KB
/
fill_in_missing_values
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
from pandas import DataFrame
import math
if 'transformer' not in globals():
from mage_ai.data_preparation.decorators import transformer
if 'test' not in globals():
from mage_ai.data_preparation.decorators import test
def select_number_columns(df: DataFrame) -> DataFrame:
return df[['Age', 'Fare', 'Parch', 'Pclass', 'SibSp', 'Survived']]
def fill_missing_values_with_median(df: DataFrame) -> DataFrame:
for col in df.columns:
values = sorted(df[col].dropna().tolist())
median_age = values[math.floor(len(values) / 2)]
df[[col]] = df[[col]].fillna(median_age)
return df
@transformer
def transform_df(df: DataFrame, *args, **kwargs) -> DataFrame:
"""
Template code for a transformer block.
Add more parameters to this function if this block has multiple parent blocks.
There should be one parameter for each output variable from each parent block.
Args:
df (DataFrame): Data frame from parent block.
Returns:
DataFrame: Transformed data frame
"""
# Specify your transformation logic here
return fill_missing_values_with_median(select_number_columns(df))
@test
def test_output(df) -> None:
"""
Template code for testing the output of the block.
"""
assert df is not None, 'The output is undefined'