FilesExpand file tree

master

/

UDFDemo.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

48 lines (36 loc) · 1.32 KB

master

/

UDFDemo.py

Copy path

Top

File metadata and controls

• Code
• Blame

48 lines (36 loc) · 1.32 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

import re

from pyspark.sql import *

from pyspark.sql.functions import *

from pyspark.sql.types import *

from lib.logger import Log4j

def parse_gender(gender):

female_pattern = r"^f$|f.m|w.m"

male_pattern = r"^m$|ma|m.l"

if re.search(female_pattern, gender.lower()):

return "Female"

elif re.search(male_pattern, gender.lower()):

return "Male"

else:

return "Unknown"

if __name__ == "__main__":

spark = SparkSession \

.builder \

.appName("UDF Demo") \

.master("local[2]") \

.getOrCreate()

logger = Log4j(spark)

survey_df = spark.read \

.option("header", "true") \

.option("inferSchema", "true") \

.csv("data/survey.csv")

survey_df.show(10)

parse_gender_udf = udf(parse_gender, returnType=StringType())

logger.info("Catalog Entry:")

[logger.info(r) for r in spark.catalog.listFunctions() if "parse_gender" in r.name]

survey_df2 = survey_df.withColumn("Gender", parse_gender_udf("Gender"))

survey_df2.show(10)

spark.udf.register("parse_gender_udf", parse_gender, StringType())

logger.info("Catalog Entry:")

[logger.info(r) for r in spark.catalog.listFunctions() if "parse_gender" in r.name]

survey_df3 = survey_df.withColumn("Gender", expr("parse_gender_udf(Gender)"))

survey_df3.show(10)