Another approach can be a Custom Function , catered to the Logic you want to apply.
However this approach can get a bit complex , as it has few subtle things to take care of while implementing it.
Both the above answers are better and faster as they use Spark Native Functions
This is merely an additional way of doing it using applyInPandas
ApplyInPandas Approach
Custom Function - Contains the Custom Logic
def generate_ctas_list(inp,base_column=None,target_column=None,column_lst=None):
### Sort the values for your Target Columns
tgt_col_values = np.sort(inp[target_column].values)
### Getting the Max Value , to use to generate the List
max_tgt = tgt_col_values[-1]
zeros_lst = np.zeros(max_tgt)
### Substract 1 to ensure indexing starts from 0
tgt_col_values = tgt_col_values - 1
### Update the Zeros List with tgt_col_values to 1
zeros_lst[tgt_col_values] = 1
### Get the singular value for Base Column
base_value = inp.loc[:,base_column].values[0]
### Generate DataFrame from the above custom logic to return
data = [base_value,zeros_lst]
res = pd.DataFrame(data,column_lst).T
return res
Implementation
from functools import partial
import pandas as pd
import numpy as np
input_list = [
(11,1)
,(11,2)
,(11,7)
,(45,7)
]
sparkDF = sql.createDataFrame(input_list,['ID','CTA'])
schema = StructType([
StructField('ID', DoubleType(), True),
StructField('CTA', ArrayType(DoubleType()), True)
])
partial_func = partial(generate_ctas_list,base_column='ID',target_column='CTA',column_lst=['ID','CTA'])
sparkDF_agg = sparkDF.groupby('ID').applyInPandas(partial_func,schema)
sparkDF_agg.show(truncate=False)
+----+-----------------------------------+
|ID |CTA |
+----+-----------------------------------+
|11.0|[1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0]|
|45.0|[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]|
+----+-----------------------------------+