pyspark withcolumn вставка списка в каждой строке

Question

У меня есть df с одним столбцом type, и у меня есть два списка

women = ['0980981', '0987098']
men = ['1234567', '4567854']

, теперь я хочу добавить другой столбец на основе значения столбца type, например:

from pyspark.sql import functions as psf
df_ = df.withColumn('new_col', psf.when(psf.col('type') == 'men', men).when(psf.col('type') == 'women', women))

Но я думаю, что мы не можем вставить список напрямую, как мы можем вставить Array('1234567', '4567854') в Scala.Я тоже попробовал psf.lit(men), но не повезло.

Есть идеи, как это сделать?

Answer 1

Используйте pyspark.sql.functions.array, который принимает list выражений столбца и возвращает однозначное выражение типа Array в сочетании с list пониманием над men:

men = ['1234567', '4567854']

df = spark.createDataFrame([['women'], ['men']], 'type: string')
df.withColumn('new_col', F.when(F.col('type') == 'men', F.array([F.lit(string) for string in men]))).show()

Выход:

+-----+------------------+
| type|           new_col|
+-----+------------------+
|women|              null|
|  men|[1234567, 4567854]|
+-----+------------------+