当第 n 列值为“x”时，访问 Pyspark 数据帧的第 (n+1) 列

Question

I have a dataframe with 1 million rows and 200 columns. I have to put only a few columns in my final dataframe. If for every row one of the column value is 3300, i need to put the next column's value in my dataframe.

For eg:

Here i have col3 value as 3300 so i need to put col4 in my final dataframe. Using column name won't be a good solution because i have 200 columns.

Answer 1

One way to do this is

df = spark.createDataFrame([(1100,1200,3300,4400,5500),(3300,1200,3200,4400,5500),(1100,1200,3200,4400,3300),(1100,3300,3300,4400,5500)],['col1','col2','col3','col4','col5'])

+----+----+----+----+----+
|col1|col2|col3|col4|col5|
+----+----+----+----+----+
|1100|1200|3300|4400|5500|
|3300|1200|3200|4400|5500|
|1100|1200|3200|4400|3300|
|1100|3300|3300|4400|5500|
+----+----+----+----+----+

from itertools import chain
column_map = create_map([lit(i) for i in chain(*enumerate(df.columns))])

df.withColumn('data',array(df.columns)).\
   withColumn('index',array_position(array(df.columns),3300)).\
   withColumn('value',col('data').getItem(col('index'))).\
   withColumn('columnName',column_map[col('index')]).\
   select('columnName','value').show()

+----------+-----+
|columnName|value|
+----------+-----+
|      col4| 4400|
|      col2| 1200|
|      null| null|
|      col3| 3300|
+----------+-----+

Update To fetch fixed cols and the value of column right after 3300 use this

df.withColumn('data',array(df.columns)).\
   withColumn('index',array_position(array(df.columns),3300)).\
   withColumn('value',col('data').getItem(col('index'))).\
   withColumn('columnName',column_map[col('index')]).\
   select('col1','col2','value')