From pyspark.sql.dataframe.DataFrame to arraytype

Question

Suppose I have the following DataFrame.

import pyspark.sql.functions as f
from pyspark.sql.window import Window

l =[( 9    , 1,  'A' ),
    ( 9    , 2, 'B'  ),
    ( 9    , 3, 'C'  ),
    ( 9    , 4, 'D'  ),
    ( 10   , 1, 'A'  ),
    ( 10   , 2, 'B' )]
df = spark.createDataFrame(l, ['prod','rank', 'value'])
df.show()

+----+----+-----+
|prod|rank|value|
+----+----+-----+
|   9|   1|    A|
|   9|   2|    B|
|   9|   3|    C|
|   9|   4|    D|
|  10|   1|    A|
|  10|   2|    B|
+----+----+-----+

How can I create a new frame with an array with the values of the value column sorted based on the rank?

Desired Output:

l =[( 9    , ['A','B','C','D'] ),
    ( 10   , ['A','B'])]

l = spark.createDataFrame(l, ['prod', 'conc'])

+----+------------+
|prod|        conc|
+----+------------+
|   9|[A, B, C, D]|
|  10|      [A, B]|
+----+------------+

Possible duplicate of [collect\_list by preserving order based on another variable](https://stackoverflow.com/questions/46580253/collect-list-by-preserving-order-based-on-another-variable). — pault, Dec 17 '18 at 21:08
Specifically look at [this answer](https://stackoverflow.com/a/50668635/5858851) for a non-udf solution. — pault, Dec 17 '18 at 21:15

cph_sto · Answer 1 · 2018-12-18T08:32:27.310

0

df = df.orderBy(["prod", "rank"], ascending=[1, 1])
df = df.rdd.map(lambda r: (r.prod, r.value)).reduceByKey(lambda x,y: list(x) + list(y)).toDF(['prod','conc'])
df.show()
+----+------------+
|prod|        conc|
+----+------------+
|   9|[A, B, C, D]|
|  10|      [A, B]|
+----+------------+

edited Dec 18 '18 at 08:32

answered Dec 17 '18 at 21:37

cph_sto

7,189
12
42
78

score -1 · Answer 2 · answered Dec 17 '18 at 22:11

Here's a quick solution based on what you've specified. Hope it helps

w = Window.partitionBy('prod').orderBy('rank')
desiredDF = df.withColumn('values_list', f.collect_list('value').over(w)).groupBy('prod').agg(f.max('values_list').alias('conc'))
desiredDF.show()

+----+------------+
|prod|        conc|
+----+------------+
|   9|[A, B, C, D]|
|  10|      [A, B]|
+----+------------+

From pyspark.sql.dataframe.DataFrame to arraytype

2 Answers2