How do we add the random samples taken from a column to another column?

Arshad 0
Viewed 33

Let's say I have a dataframe like -

target_id = [3733345, 3725312, 3717114, 3408996, 3354970]
test_df = spark.createDataFrame(target_id, IntegerType()).withColumnRenamed("value", "target_id")

enter image description here

I want to add random samples of values from this column to another column other_target_ids such that the output comes something like below:

target_id   other_ids
3733345     [3731634, 3729995, 3728014, 3708332, 3720...
3725312     [3711541, 3726052, 3733763, 900056057, 371...
3717114     [3701718, 3713481, 3715433, 3714825, 3731...
3408996     [3405896, 3250400, 3237054, 3242492, 3256...
3354970     [3354969, 3347893, 3348168, 3353273, 3356...
pyspark
1 Answers

I guest you could do that with a few steps to collect and filter ids. Check my sample below

from pyspark.sql import functions as F
from pyspark.sql import types as T
from pyspark.sql.window import Window as W

(test_df

    # collect all ids into a list
    .withColumn('ids', F.collect_list('target_id').over(W.Window.partitionBy(F.lit(1))))
 
    # remove `target_id` from the above list
    .withColumn('ids', F.array_except('ids', F.array(F.col('target_id'))))

    # shuffle the ids
    .withColumn('ids', F.shuffle('ids'))

    # "sampling" the ids by getting the first few items
    .withColumn('ids', F.slice('ids', 1, 2))

    # display
    .show(10, False)
)

+---------+------------------+
|target_id|ids               |
+---------+------------------+
|3733345  |[3408996, 3354970]|
|3725312  |[3408996, 3733345]|
|3717114  |[3408996, 3354970]|
|3408996  |[3733345, 3354970]|
|3354970  |[3733345, 3725312]|
+---------+------------------+
pltc0
Related