Pyspark-获取Spark数据框中值的第一列出现

时间:2019-05-02 22:21:49

标签: pyspark

我有一个数据帧,如下所示,我需要第一个,最后一次出现的值为0和非零值

Id  Col1  Col2  Col3  Col4
1    1     0     0      2
2    0     0     0      0
3    4     2     2      4
4    2     5     9      0
5    0     4     0      0

Expected Result:    
Id  Col1  Col2  Col3  Col4 First_0 Last_0 First_non_zero Last_non_zero
1    1     0     0      2   2        3         1               4
2    0     0     0      0   1        4         0               0
3    4     2     2      4   0        0         1               4
4    2     5     9      0   4        4         1               3
5    0     4     0      0   1        4         2               2

1 个答案:

答案 0 :(得分:0)

这是使用pyspark的F.array()F.greatest()F.least()的一种方法:

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [(1,1,0,0,2), (2,0,0,0,0), (3,4,2,2,4), (4,2,5,9,0), (5,0,4,0,0)]
  , ['Id','Col1','Col2','Col3','Col4']
)

df.show()
#+---+----+----+----+----+
#| Id|Col1|Col2|Col3|Col4|
#+---+----+----+----+----+
#|  1|   1|   0|   0|   2|
#|  2|   0|   0|   0|   0|
#|  3|   4|   2|   2|   4|
#|  4|   2|   5|   9|   0|
#|  5|   0|   4|   0|   0|
#+---+----+----+----+----+

# column names involved in the calculation
cols = df.columns[1:]

# create an array column `arr_0` with index of elements(having F.col(cols[index])==0) in array cols
# then select the greatest and least value to identify the first_0 and last_0
# fillna with '0' when none of the items is '0'
df.withColumn('arr_0', F.array([ F.when(F.col(cols[i])==0, i+1) for i in range(len(cols))])) \
  .withColumn('first_0', F.least(*[F.col('arr_0')[i] for i in range(len(cols))])) \
  .withColumn('last_0', F.greatest(*[F.col('arr_0')[i] for i in range(len(cols))])) \
  .fillna(0, subset=['first_0', 'last_0']) \
  .show()
#+---+----+----+----+----+------------+-------+------+
#| Id|Col1|Col2|Col3|Col4|       arr_0|first_0|last_0|
#+---+----+----+----+----+------------+-------+------+
#|  1|   1|   0|   0|   2|   [, 2, 3,]|      2|     3|
#|  2|   0|   0|   0|   0|[1, 2, 3, 4]|      1|     4|
#|  3|   4|   2|   2|   4|       [,,,]|      0|     0|
#|  4|   2|   5|   9|   0|     [,,, 4]|      4|     4|
#|  5|   0|   4|   0|   0|  [1,, 3, 4]|      1|     4|
#+---+----+----+----+----+------------+-------+------+

如果您使用的是 pyspark 2.4 ,则还可以尝试F.array_min()F.array_max()

df.withColumn('arr_0', F.array([ F.when(F.col(cols[i])==0, i+1) for i in range(len(cols)) ])) \
  .select('*', F.array_min('arr_0').alias('first_0'), F.array_max('arr_0').alias('last_0')) \
  .fillna(0, subset=['first_0', 'last_0']) \
  .show()
#+---+----+----+----+----+------------+-------+------+
#| Id|Col1|Col2|Col3|Col4|       arr_0|first_0|last_0|
#+---+----+----+----+----+------------+-------+------+
#|  1|   1|   0|   0|   2|   [, 2, 3,]|      2|     3|
#|  2|   0|   0|   0|   0|[1, 2, 3, 4]|      1|     4|
#|  3|   4|   2|   2|   4|       [,,,]|      0|     0|
#|  4|   2|   5|   9|   0|     [,,, 4]|      4|     4|
#|  5|   0|   4|   0|   0|  [1,, 3, 4]|      1|     4|
#+---+----+----+----+----+------------+-------+------+

您可以对last_non_zerofirst_non_zero进行同样的操作。