Вот один из способов использования 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_zero и first_non_zero.