AWS Glue pySpark Filter & Manual Mapping of Several Columns - pyspark

Question

I'm using AWS Glue Studio with DynamicFrameCollections. I created a custom transformation where I am looking to filter by multiple columns and modify 2 column in the row based off a static mapping list. I'm struggling to figure out what the most efficient way to do this - pandas, udfs, or something completely different?
Consider the sample dataframe:
data = [{"Category": 'A', "Subcategory": 2, "Value": 121.44, "Properties": {}},
{"Category": 'B', "Subcategory": 2, "Value": 300.01, "Properties": None},
{"Category": 'C', "Subcategory": 3, "Value": 10.99, "Properties": { "Active":True } },
{"Category": 'E', "Subcategory": 4, "Value": 33.87, "Properties": { "Active":True, "ReadOnly": False }},
{"Category": 'E', "Subcategory": 1, "Value": 11.37, "Properties": { "Active":True }}
]
df = spark.createDataFrame(data)
I need to filter and transform by Category and Subcategory. Below is the sample mapping with the key as the category and subcategory merged while the first value in the array must be created as a new column ActivityName and the second values must be merged with the Properties column:
mapping= {"A2": ["EatingFood", { "Visible": True }],
"A3": ["DrinkingWater", { "Visible": False }],
"B2": ["Sleep", { "Visible": False }],
"C3": ["Exercise", { "Visible": False }],
"E4": ["Running", { "Visible": False }],
}
The output data I am expecting is:
resultingData = [{"Category": 'A', "Subcategory": 2, "ActivityName":"EatingFood", "Value": 121.44, "Properties": { "Visible": True }},
{"Category": 'B', "Subcategory": 2, "ActivityName":"Sleep", "Value": 300.01, "Properties": {"Visible": False}},
{"Category": 'C', "Subcategory": 3, "ActivityName":"Exercise", "Value": 10.99, "Properties": { "Active":True, "Visible": False } },
{"Category": 'E', "Subcategory": 4, "ActivityName":"Running", "Value": 33.87, "Properties": { "Active":True, "ReadOnly": False, "Visible": False }}
]
Note that the last data entry, E1 is missing because it was not in my mapping filter.
Is there any way to achieve this? I have a large list of items that I need to manually filter/map/transform like this. Thank you.

Answer 1

I got this working by transforming the dynamicframe into a dataframe and processing it using glue functions. Here's what I did:
def FilterAndMap (glueContext, dfc) -> DynamicFrameCollection:
from pyspark.sql.types import StringType, ArrayType
from awsglue.dynamicframe import DynamicFrame
import pyspark.sql.functions as f
import json
mapping= {"A2": ["EatingFood", json.dumps({ "Visible": True })],
"A3": ["DrinkingWater", json.dumps({ "Visible": False })],
"B2": ["Sleep", json.dumps({ "Visible": False })],
"C3": ["Exercise", json.dumps({ "Visible": False })],
"E4": ["Running", json.dumps({ "Visible": False })],
}
df = dfc.select(list(dfc.keys())[0]).toDF()
def func_filter_udf(concat_str):
return mapping[concat_str]
def func_map_udf(map_str):
if map_str[1]:
map_string = json.loads(map_str[0])
ret_val = json.dumps({**map_string, **json.loads(map_str[1])})
else:
ret_val = map_str[0]
return ret_val
filter_udf = f.udf(func_filter_udf, ArrayType(StringType()))
map_udf = f.udf(func_map_udf, StringType())
df = df.filter(f.concat("Category", "Subcategory").isin([*mapping]))
df = df.withColumn("concat_col", filter_udf(f.concat("Category", "Subcategory")))
df = (df.withColumn("ActivityName", df.concat_col[0]).
withColumn("Properties", map_udf(f.struct(df.concat_col[1], df.Properties))))
df = df.drop("concat_col")
dyf_processed = DynamicFrame.fromDF(df, glueContext, "filtered")
return(DynamicFrameCollection({"filtered": dyf_processed }, glueContext))