Package: pyspark-nested-functions
Nested fields transformation for pysparkMotivationApplying transformations to nested structures is tricky in Spark.
Assume we have below nested JSON data:[{"data":{"city":{"addresses":[{"id":"my-id"},{"id":"my-id2"}]}}}]To hash the nestedidfield you need to write the following PySpark code:importpyspark.sql.functionsasFhashed=df.withColumn("data",(F.col("data").withField("city",F.col("data.city").withField("addresses",F.transform("data.city.addresses",lambdac:c.withField("id",F.sha2(c.getField("id"),256)))))))With the library the code above can be simplified to:fromnestedfunctions.functions.hashimporthash_fieldhashed=hash_field(df,"data.city.addresses.id",num_bits=256)InstallTo install the current release$ pip install pyspark-nested-functionsAvailable functionsAdd nested fieldAdding a nested field called new_column_name based on a lambda function working on the column_to_process nested field.
Fields column_to_process and new_column_name need to have the same parent or be at the root!fromnestedfunctions.functions.add_nested_fieldimportadd_nested_fieldfrompyspark.sql.functionsimportwhenprocessed=add_nested_field(df,column_to_process="payload.array.booleanField",new_column_name="payload.array.booleanFieldAsString",f=lambdacolumn:when(column,"Y").when(~column,"N").otherwise(""),)Date FormatFormat a nested date field from current_date_format to target_date_format.fromnestedfunctions.functions.date_formatimportformat_datedate_formatted_df=format_date(df,field="customDimensions.value",current_date_format="y-d-M",target_date_format="y-MM")DropRecursively drop fields on any nested level.fromnestedfunctions.functions.dropimportdropdropped_df=drop(df,field="root_level.children1.children2")DuplicateDuplicate the nested field column_to_duplicate as duplicated_column_name.
Fields column_to_duplicate and duplicated_column_name need to have the same parent or be at the root!fromnestedfunctions.functions.duplicateimportduplicateduplicated_df=duplicate(df,column_to_duplicate="payload.lineItems.comments",duplicated_column_name="payload.lineItems.commentsDuplicate")ExprAdd or overwrite a nested field based on an expression.fromnestedfunctions.functions.exprimportexprfield="emails.unverified"processed=expr(df,field=field,expr=f"transform({field}, x -> (upper(x)))")Field RenameRename all the fields based on any rename function.(If you only want to rename specific fields filter on them in your rename function)fromnestedfunctions.functions.field_renameimportrenamedefcapitalize_field_name(field_name:str)->str:returnfield_name.upper()renamed_df=rename(df,rename_func=capitalize_field_name())FillnaThis function mimics the vanilla pyspark fillna functionality with added support for filling nested fields.
The use of the input parameters value and subset is exactly the same as for the vanilla pyspark implementation as describedhere.fromnestedfunctions.functions.fillnaimportfillna# Fill all null boolean fields with Falsefilled_df=fillna(df,value=False)# Fill nested field with valuefilled_df=fillna(df,subset="payload.lineItems.availability.stores.availableQuantity",value=0)# To fill array which is null specify list of valuesfilled_df=fillna(df,value={"payload.comments":["Automatically triggered stock check"]})# To fill elements of array that are null specify single valuefilled_df=fillna(df,value={"payload.comments":"Empty comment"})FlattenerReturn the flattened representation of the dataframe's schema.fromnestedfunctions.spark_schema.utilityimportSparkSchemaUtilityflattened_schema=SparkSchemaUtility().flatten_schema(df.schema)# flattened_schema = ["root-element",#                   "root-element-array-primitive",#                   "root-element-array-of-structs.d1.d2",#                   "nested-structure.n1",#                   "nested-structure.d1.d2"]HashReplace a nested field by its SHA-2 hash value.
By default the number of bits in the output hash value will be 256 but a different value can be set.fromnestedfunctions.functions.hashimporthash_fieldhashed_df=hash_field(df,"data.city.addresses.id",num_bits=256)NullifyMaking a field null on any nested level.fromnestedfunctions.functions.nullifyimportnullifynullified_df=nullify(df,field="creditCard.id")Overwrite nested fieldOverwrites a nested field based on a lambda function working on this nested field.fromnestedfunctions.functions.terminal_operationsimportapply_terminal_operationfrompyspark.sql.functionsimportwhenprocessed=apply_terminal_operation(df,field="payload.array.someBooleanField",f=lambdacolumn,type:when(column,"Y").when(~column,"N").otherwise(""),)RedactReplace a field by the default value of its data type.
The default value of a data type is typically its min or max value and can be foundhere.fromnestedfunctions.functions.redactimportredactredacted_df=redact(df,field="customDimensions.metabolicsConditions")WhitelistPreserving all fields listed in parameters. All other fields will be droppedfromnestedfunctions.functions.whitelistimportwhitelistwhitelisted_df=whitelist(df,["addresses.postalCode","creditCard"])Predicate variations of above functionsSome of the above functions likehash,nullifyanddate_formathave predicate variations.
For these variations you can specify a singlepredicate_key/predicate_valuepair for which the function will be run. This is mainly handy when you only want to adapt a nested value when one of the root columns has a specific value.LicenseApache License 2.0