Replace elements in an array by Position in Dataframe - Pyspark

Viewed 324

I have a dataframe:

|ID|CTA|
|------|
|11|1  |
|11|2  |
|11|7  |
|45|7  |

I need to group by ID and that each ID has its ARRAY of length 7 but indicating a 1 when in CTA it has the position

So my output dataframe should look like:

|ID|CTAS             |
|------------------- |
|11|[1,1,0,0,0,0,1]  |
|45|[0,0,0,0,0,0,1]  |

can you help me?


UPDATE:

Code

How can I leave the zeros in the array?

3 Answers

You can apply a TRANSFORM expression and iterate over sequence(1, 7) to check if the value contains in CTAS column:

import pyspark.sql.functions as f

group_df = (df
            .groupBy('ID')
            .agg(f.collect_list('CTA').alias('CTAS')))
# +---+---------+
# |ID |CTAS     |
# +---+---------+
# |11 |[1, 2, 7]|
# |45 |[7]      |
# +---+---------+

pos_df = (group_df
          .withColumn('CTAS', 
                      f.expr('transform(sequence(1, 7), value -> cast(array_contains(CTAS, value) as int))')))
pos_df.sort('ID').show(truncate=False)
# +---+---------------------+
# |ID |CTAS                 |
# +---+---------------------+
# |11 |[1, 1, 0, 0, 0, 0, 1]|
# |45 |[0, 0, 0, 0, 0, 0, 1]|
# +---+---------------------+

each ID has its ARRAY of length 7

Just adding another way (I personally like this answer posted above). Using a distinct crossjoin to spark.range and a left join to populate missing rows, then groupby and collect:

out = (df.withColumn("CTAS",F.lit(1)).join(
  df.select("ID").dropDuplicates().crossJoin(spark.range(1,8)
                                        .withColumnRenamed("id","CTA"))
       ,on=["ID","CTA"],how='outer').fillna({"CTAS":0})
       .groupBy("ID").agg(F.collect_list("CTAS").alias("CTAS")))

out.show(truncate=False)

+---+---------------------+
|ID |CTAS                 |
+---+---------------------+
|11 |[1, 1, 0, 0, 0, 0, 1]|
|45 |[0, 0, 0, 0, 0, 0, 1]|
+---+---------------------+

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]|
+----+-----------------------------------+
Related