¿Cómo pivotar en múltiples columnas en Spark SQL?

Necesito pivotar más de una columna en un dataframe de pyspark. Marco de datos de muestra,

>>> d = [(100,1,23,10),(100,2,45,11),(100,3,67,12),(100,4,78,13),(101,1,23,10),(101,2,45,13),(101,3,67,14),(101,4,78,15),(102,1,23,10),(102,2,45,11),(102,3,67,16),(102,4,78,18)] >>> mydf = spark.createDataFrame(d,['id','day','price','units']) >>> mydf.show() +---+---+-----+-----+ | id|day|price|units| +---+---+-----+-----+ |100| 1| 23| 10| |100| 2| 45| 11| |100| 3| 67| 12| |100| 4| 78| 13| |101| 1| 23| 10| |101| 2| 45| 13| |101| 3| 67| 14| |101| 4| 78| 15| |102| 1| 23| 10| |102| 2| 45| 11| |102| 3| 67| 16| |102| 4| 78| 18| +---+---+-----+-----+ 

Ahora, si necesito poner la columna de precios en una fila para cada ID según el día, entonces puedo usar el método de pivote como,

 >>> pvtdf = mydf.withColumn('combcol',F.concat(F.lit('price_'),mydf['day'])).groupby('id').pivot('combcol').agg(F.first('price')) >>> pvtdf.show() +---+-------+-------+-------+-------+ | id|price_1|price_2|price_3|price_4| +---+-------+-------+-------+-------+ |100| 23| 45| 67| 78| |101| 23| 45| 67| 78| |102| 23| 45| 67| 78| +---+-------+-------+-------+-------+ 

así que cuando necesito la columna de unidades para ser transpuesta como precio, tengo que crear un dataframe más como anteriormente para las unidades y luego unirlas usando id. Pero cuando tengo más columnas como tal, probé una función para hacerlo. ,

 >>> def pivot_udf(df,*cols): ... mydf = df.select('id').drop_duplicates() ... for c in cols: ... mydf = mydf.join(df.withColumn('combcol',F.concat(F.lit('{}_'.format(c)),df['day'])).groupby('id').pivot('combcol').agg(F.first(c)),'id') ... return mydf ... >>> pivot_udf(mydf,'price','units').show() +---+-------+-------+-------+-------+-------+-------+-------+-------+ | id|price_1|price_2|price_3|price_4|units_1|units_2|units_3|units_4| +---+-------+-------+-------+-------+-------+-------+-------+-------+ |100| 23| 45| 67| 78| 10| 11| 12| 13| |101| 23| 45| 67| 78| 10| 13| 14| 15| |102| 23| 45| 67| 78| 10| 11| 16| 18| +---+-------+-------+-------+-------+-------+-------+-------+-------+ 

Necesita sugerencias sobre si es una buena práctica hacerlo y si hay alguna otra forma mejor de hacerlo. ¡Gracias por adelantado!

Al igual que en la versión de la chispa 1.6, creo que esa es la única manera porque el pivote toma solo una columna y hay un segundo valor de atributo en el que puede pasar los distintos valores de esa columna que harán que su código se ejecute más rápido porque de lo contrario, chispa tiene que ejecutarlo. , así que sí, esa es la manera correcta de hacerlo.

La solución en la pregunta es lo mejor que puedo conseguir. La única mejora sería cache en cache el conjunto de datos de entrada para evitar la doble exploración, es decir,

 mydf.cache pivot_udf(mydf,'price','units').show()