PySpark - 如何根据列中的两个值从数据帧中筛选出连续的行块

Question

我有一个数据帧，我想使用 pyspark 基于某些列值创建另一个数据帧。例如：下面是我的主要数据帧 -

Part1   Part2   Part3   Part4
aaa      up      24     k-123
bbb     down     45     i-98
ccc     down     54     k-89
fff     int      23     l-34
xyz      up      22     o-89
www      up      89     u-56

现在，我想创建另一个数据帧，它将搜索第一次出现"向下"，并将直到第一次出现"向上"。因此，预期数据帧将是：

   Part1    Part2   Part3   Part4
    bbb     down     45     i-98
    ccc     down     54     k-89
    fff     int      23     l-34
    xyz      up      22     o-89

Answer 1

第 1 步：创建DataFrame 。

from pyspark.sql.functions import when, col, lit
df = spark.createDataFrame(
    [('aaa','up',24,'k-123'),('bbb','down',45,'i-98'),('ccc','down',54,'k-89'),
     ('fff','int', 23,'l-34'),('xyz','up',22,'o-89'),('www','up',89,'u-56')], 
    schema = ['Part1','Part2','Part3','Part4']
)
df.show()
+-----+-----+-----+-----+
|Part1|Part2|Part3|Part4|
+-----+-----+-----+-----+
|  aaa|   up|   24|k-123|
|  bbb| down|   45| i-98|
|  ccc| down|   54| k-89|
|  fff|  int|   23| l-34|
|  xyz|   up|   22| o-89|
|  www|   up|   89| u-56|
+-----+-----+-----+-----+

步骤2： 首先，我们需要找到down的第一个出现并删除其上方的所有行。为此，我们创建一个列cumulative，如果值为 1 == down，Part2否则为 0，最后取此列的累积总和。

df = df.withColumn('Dummy',lit('dummy'))
df = df.withColumn('cumulative',when(col('Part2')=='down',1).otherwise(0))
df = df.selectExpr(
    'Part1','Part2','Part3','Part4','Dummy',
    'sum(cumulative) over (order by row_number() over (order by Dummy)) as cumulative'
 )
df.show()
+-----+-----+-----+-----+-----+----------+
|Part1|Part2|Part3|Part4|Dummy|cumulative|
+-----+-----+-----+-----+-----+----------+
|  aaa|   up|   24|k-123|dummy|         0|
|  bbb| down|   45| i-98|dummy|         1|
|  ccc| down|   54| k-89|dummy|         2|
|  fff|  int|   23| l-34|dummy|         2|
|  xyz|   up|   22| o-89|dummy|         2|
|  www|   up|   89| u-56|dummy|         2|
+-----+-----+-----+-----+-----+----------+

现在，删除累积总和为 0 的所有行。这将删除所有行，直到down首次出现。

df = df.where(col('cumulative')>=1)

步骤

3：执行与上述步骤2相同的操作，除了对up执行此操作并删除第cumulative列中值小于或等于1的所有行。这样，我们将删除第一次出现up下方的所有行。

df = df.withColumn('cumulative',when(col('Part2')=='up',1).otherwise(0))
df = df.selectExpr(
    'Part1','Part2','Part3','Part4','Dummy',
    'sum(cumulative) over (order by row_number() over (order by Dummy)) as cumulative'
 )
df = df.where(col('cumulative')<=1).drop('Dummy','cumulative')
df.show()
+-----+-----+-----+-----+
|Part1|Part2|Part3|Part4|
+-----+-----+-----+-----+
|  bbb| down|   45| i-98|
|  ccc| down|   54| k-89|
|  fff|  int|   23| l-34|
|  xyz|   up|   22| o-89|
+-----+-----+-----+-----+