PySparkで横持ちを縦持ちに変換するより良い方法を知りたい

df.rdd.fratmapを使用する事で実現できたものの、functionsをうまく使う事でスマートにできるような気がしています。
できないという回答でも構いませんので、ご意見いただけますでしょうか。

コード

df = spark.createDataFrame([
  ['a','11','12'],['b','21','22'],['c','31','32']
], ['id','code1','code2'])
df.show()
result = spark.createDataFrame([
  ['a','11','code1'],['a','12','code2'],['b','21','code1'],['b','22','code2'],['c','31','code1'],['c','32','code2']
], ['id','code','type'])
result.show()

変換前

+---+-----+-----+
| id|code1|code2|
+---+-----+-----+
|  a|   11|   12|
|  b|   21|   22|
|  c|   31|   32|
+---+-----+-----+

変換後

+---+----+-----+
| id|code| type|
+---+----+-----+
|  a|  11|code1|
|  a|  12|code2|
|  b|  21|code1|
|  b|  22|code2|
|  c|  31|code1|
|  c|  32|code2|
+---+----+-----+

試してみたこと

rddからflatmapを使用する事で実現はできました。

from pyspark.sql import Row
def func(row):
  return [
    Row(id=row['id'], code=row['code1'], type='code1'),
    Row(id=row['id'], code=row['code2'], type='code2')
  ]
spark.createDataFrame(df.rdd.flatMap(func)).show()

typeがないのであればexplodeを使う事で実現できたため、もう少し良い方法があるのではないかと思った次第です。
（これはこれでもう少し良い方法がある気がしますが。。。）

from pyspark.sql import functions as f
df.select('id', f.explode(f.array('code1','code2')).alias('code')).show()