Pyspark: labeling unique arrays regardless of order in array column

user_15779694 0
Viewed 34

I have an array column in a pyspark df and would like to label the unique arrays that contains the same elements regardless of order, expected output:

| Array Column | Label  |
| ------------ | ------ |
| [1,2,3]      | Group1 |
| [3,2,1]      | Group1 |
| [2,1,5]      | Group2 |
| [1,2,5]      | Group2 |
| [2,3,1]      | Group1 |

Anyone has any ideas? array_except == empty would involve a row by row comparison and the data is pretty big. So wondering if there are any other solutions for pyspark.

pyspark label arrays
1 Answers

Data

df= spark.createDataFrame([(1,[1,2,3]),     
    (2, [3,2,1] ),    
    (3, [2,1,5]),      
    (4, [1,2,5]) ,     
    (5, [2,3,1] ) ],
      ('id', 'array_column'))
    
    df.show()

Solution

df= (df.withColumn('label', expr("array_join(array_sort(array_column),',')"))#Sort array elements, convert array by joining its elements to make string
  #.withColumn('label', dense_rank().over(Window.partitionBy().orderBy('label')))#dense rank to get order
  .withColumn('label', concat(lit('Group'),(dense_rank().over(Window.partitionBy().orderBy('label'))).cast('string')))#dense rank to get order

 
 .orderBy('id')#reorder
).show()

+---+------------+------+
| id|array_column| label|
+---+------------+------+
|  1|   [1, 2, 3]|Group1|
|  2|   [3, 2, 1]|Group1|
|  3|   [2, 1, 5]|Group2|
|  4|   [1, 2, 5]|Group2|
|  5|   [2, 3, 1]|Group1|
+---+------------+------+
user_89869750
Related