使用PySpark中的列值组合查找最小和最大范围

时间:2019-09-15 06:55:19

标签: python dataframe pyspark

我有一个像这样的pyspark数据框,

+----------+--------+----------+----------+
|id_       | p      |d1        |  d2      |
+----------+--------+----------+----------+
|  1       | A      |2018-09-26|2018-10-26|
|  2       | B      |2018-06-21|2018-07-19|
|  2       | C      |2018-07-13|2018-10-07|
|  2       | B      |2018-12-31|2019-02-27|
|  2       | A      |2019-01-28|2019-06-25|
-------------------------------------------

从这个数据框中,我必须制作一个这样的数据框,

+----------+--------+----------+----------+
|id_       | q      |d1        |  d2      |
+----------+--------+----------+----------+
|  1       | A      |2018-09-26|2018-10-26|
|  2       | B      |2018-06-21|2018-07-12|
|  2       | B C    |2018-07-13|2018-07-19|
|  2       | C      |2018-07-20|2019-10-07|
|  2       | B      |2018-12-31|2019-01-27|
|  2       | B A    |2019-01-28|2019-02-27|
|  2       | A      |2019-02-28|2019-06-25|
-------------------------------------------

这就像是从某个时间到何时查找特定p的数据中存在id_的值。如果同一天有多个p,则两者都应存在于数据中,并以空格分隔。

我试图通过创建min(d1)max(d2)范围内的每个日期并相应地填充它们来做到这一点。从该数据框中,经过一些融合和分组,我可以获得所需的结果。

但是该过程需要很长时间并且效率很低。

我正在寻找一种有效的方法来执行此任务。

我也可以遇到更复杂的重叠情况,即两个以上p值之间的重叠。

在下面查看示例数据,

+----------+--------+----------+----------+
|id_       | p      |d1        |  d2      |
+----------+--------+----------+----------+
|  1       | A      |2018-09-26|2018-10-26|
|  2       | B      |2018-06-21|2018-07-19|
|  2       | C      |2018-06-27|2018-07-07|
|  2       | A      |2018-07-02|2019-02-27|
|  2       | A      |2019-03-28|2019-06-25|
-------------------------------------------

此必须转换为

+----------+--------+----------+----------+
|id_       | q      |d1        |  d2      |
+----------+--------+----------+----------+
|  1       | A      |2018-09-26|2018-10-26|
|  2       | B      |2018-06-21|2018-06-26|
|  2       | B C    |2018-06-27|2018-07-01|
|  2       | B C A  |2018-07-02|2018-07-07|
|  2       | A B    |2018-07-08|2018-07-19|
|  2       | A      |2018-07-20|2019-02-27|
|  2       | A      |2019-03-28|2019-06-25|
-------------------------------------------

q中各个项目的顺序无关紧要。即,如果A,B和C重叠。它既可以显示为A B C,也可以显示为B C A或A C B,等等。

我还添加了一个很难获得的保护套,即d2 == lead(d1).over(window)。在这种情况下,可以安全地假设p的值是不同的。即p != lead(p).over(window)

+---+---+----------+----------+
|id_| p |    d1    | d2       |
+---+---+----------+----------+
|100| 12|2013-10-16|2014-01-17|
|100| 12|2014-01-20|2014-04-15|
|100| 12|2014-04-22|2014-05-19|
|100| 12|2014-05-22|2014-06-19|
|100| 12|2014-07-23|2014-09-18|
|100| 12|2014-09-23|2014-12-18|
|100| 12|2014-12-20|2015-01-16|
|100| 12|2015-01-23|2015-02-19|
|100| 12|2015-02-21|2015-04-20|
|100| 7 |2015-04-20|2015-05-17|
|100| 7 |2015-05-19|2015-06-15|
|100| 7 |2015-06-18|2015-09-01|
|100| 7 |2015-09-09|2015-11-26|
+---+---+----------+----------+

在以上数据中,从底部开始的第4和第5行显示了这种情况。在这种情况下,预期结果是

+---+-----+----------+----------+
|id_| p   | d1       | d2       |
+---+-----+----------+----------+
|100| 12  |2013-10-16|2014-01-17|
|100| 12  |2014-01-20|2014-04-15|
|100| 12  |2014-04-22|2014-05-19|
|100| 12  |2014-05-22|2014-06-19|
|100| 12  |2014-07-23|2014-09-18|
|100| 12  |2014-09-23|2014-12-18|
|100| 12  |2014-12-20|2015-01-16|
|100| 12  |2015-01-23|2015-02-19|
|100| 12  |2015-02-21|2015-04-19|
|100| 12 7|2015-04-20|2015-04-20|
|100| 7   |2015-04-21|2015-05-17|
|100| 7   |2015-05-19|2015-06-15|
|100| 7   |2015-06-18|2015-09-01|
|100| 7   |2015-09-09|2015-11-26|
+---+-----+----------+----------+

下面是相同情况的另一个示例

+---+---+----------+----------+
|id_| p | d1       | d2       |
+---+---+----------+----------+
|101| 12|2015-02-24|2015-03-23|
|101| 12|2015-04-01|2015-05-19|
|101| 12|2015-05-29|2015-06-25|
|101| 12|2015-07-03|2015-07-30|
|101| 12|2015-09-02|2015-09-29|
|101| 12|2015-10-02|2015-10-29|
|101| 9 |2015-10-29|2015-11-11|
|101| 9 |2015-11-25|2015-12-22|
+---+---+----------+----------+

同样的预期结果是,

+---+-----+----------+----------+
|id_| q   | d1       | d2       |
+---+-----+----------+----------+
|101| 12  |2015-02-24|2015-03-23|
|101| 12  |2015-04-01|2015-05-19|
|101| 12  |2015-05-29|2015-06-25|
|101| 12  |2015-07-03|2015-07-30|
|101| 12  |2015-09-02|2015-09-29|
|101| 12  |2015-10-02|2015-10-28|
|101| 12 9|2015-10-29|2015-10-29|
|101| 9   |2015-10-30|2015-11-11|
|101| 9   |2015-11-25|2015-12-22|
+---+---+------------+----------+

1 个答案:

答案 0 :(得分:2)

更新:基于OP的评论和更新,由于可能发生许多重叠,因此我认为dataframe-JOIN可能是最直接的方法。以下是我在Spark 2.4.0上测试过的全新解决方案(array_join,transform,sequence等都需要Spark 2.4+):

更新2:在注释/聊天中的每次讨论中,我添加了代码逻辑,以设置每个drange(d1, d2)的边界/如何调整d1 / d2,flag中需要一个新的df_drange字段才能完成此逻辑。详细信息请参见下面的 Set up boundaries 部分

更新3::调整了代码以处理(d1 == d2)在df_drange中。最初删除了这种情况。

设置数据:

注意:我添加了df2,并将d1和d2转换为DateType(),而原始df则将两个字段保留为StringType(),因为我们需要进行一些串联操作。

from pyspark.sql import Window
from pyspark.sql.functions import lead, expr, to_date, collect_set, array_sort, array_join, broadcast

df = spark.createDataFrame([
      (1, 'A', '2018-09-26', '2018-10-26')
    , (2, 'B', '2018-06-21', '2018-07-19')
    , (2, 'C', '2018-06-27', '2018-07-07')
    , (2, 'A', '2018-07-02', '2019-02-27')
    , (2, 'A', '2019-03-28', '2019-06-25')
  ], ['id_', 'p', 'd1', 'd2'])

# convert d1, d2 to DateType() if they are StringType()
df2 = df.withColumn('d1', to_date('d1')).withColumn('d2', to_date('d2'))

df2.printSchema()
root
 |-- id_: long (nullable = true)
 |-- p: string (nullable = true)
 |-- d1: date (nullable = true)
 |-- d2: date (nullable = true)

创建参考数据帧:df_drange

df_drange包含d1和d2中所有不同的日期, 加上一个标志,当1来自df_drange.d1(在原始df中)时设置为df.d2,否则设置为0。排序日期并将其细分为间隔日期范围。检索字段d1d2flag(仅d1)并将其转换为适当的DataType()

df_drange = df.select('id_', 'd1', lit(0).alias('flag')).union(df.select('id_', 'd2', lit(1))) \
    .groupby('id_') \
    .agg(array_sort(collect_set(concat('d1', lit('-'), 'flag'))).alias('dates')) \
    .withColumn('dates', expr("""
         explode(transform(sequence(0, size(dates)-2), i -> named_struct('d1', dates[i], 'd2', dates[i+1])))
       """)) \
    .selectExpr(
         'id_'
       , "to_date(substring_index(dates.d1, '-', 3)) as d1"
       , "to_date(substring_index(dates.d2, '-', 3)) as d2"
       , "boolean(substring_index(dates.d1, '-', -1)) as flag"
     )

df_drange.orderBy('id_','d1').show()
+---+----------+----------+-----+
|id_|        d1|        d2| flag|
+---+----------+----------+-----+
|  1|2018-09-26|2018-10-26|false|
|  2|2018-06-21|2018-06-27|false|
|  2|2018-06-27|2018-07-02|false|
|  2|2018-07-02|2018-07-07|false|
|  2|2018-07-07|2018-07-19| true|
|  2|2018-07-19|2019-02-27| true|
|  2|2019-02-27|2019-03-28| true|
|  2|2019-03-28|2019-06-25|false|
+---+----------+----------+-----+

df_drange.printSchema()
root
 |-- id_: long (nullable = true)
 |-- d1: date (nullable = true)
 |-- d2: date (nullable = true)
 |-- flag: boolean (nullable = true)

使用Join设置df1

与原始df以及每个id_重叠的左连接 在 df_dranges 的( d1 d2 )和(em)的( d1 d2 )之间原始df 。后 df_drange 中的groupby( id _ d1 d2 flag ),得到array_join(collect_set(p),''):

df1 = broadcast(df_drange).join(
      df2
    , (df2.id_ == df_drange.id_) & (
            ((df2.d1 < df_drange.d2) & (df2.d2 > df_drange.d1)) 
          | ((df_drange.d1 == df_drange.d2) & df_drange.d1.between(df2.d1, df2.d2)) 
      )
    , how = 'left'
).groupby(df_drange.id_, df_drange.d1, df_drange.d2, df_drange.flag) \
 .agg(array_join(collect_set('p'), ' ').alias('q'))

df1.show()
+---+----------+----------+-----+-----+
|id_|        d1|        d2| flag|    q|
+---+----------+----------+-----+-----+
|  1|2018-09-26|2018-10-26|false|    A|
|  2|2018-06-21|2018-06-27|false|    B|
|  2|2018-06-27|2018-07-02|false|  C B|
|  2|2018-07-02|2018-07-07|false|C B A|
|  2|2018-07-07|2018-07-19| true|  B A|
|  2|2018-07-19|2019-02-27| true|    A|
|  2|2019-02-27|2019-03-28| true|     |
|  2|2019-03-28|2019-06-25|false|    A|
+---+----------+----------+-----+-----+

设置边界

对于df1,如果q ==”,则存在间隙,应删除此类行。 根据标记,next_flag,next_d1定义每个Drange的边界 如评论/聊天中所述。以下是伪代码,用于显示当前逻辑如何/何时调整d1 / d2:

flag = (if d1 is from original_d2) ? true : false
both next_d1 and next_flag defined on WindowSpec-w1

# for df1.d1: if flag is true, add 1 day, otherwise keep as-is
d1 = IF(flag, date_add(d1,1), d1)

# for df1.d2: keep as-is when there is gap with the next row or 
# the next_flag is true, else minus 1 day
d2 = IF((next_d1 != d2) or next_flag, d2, date_sub(d2,1))

实际代码:

# WindowSpec to calculate next_d1
w1 = Window.partitionBy('id_').orderBy('d1')

# filter out gaps and calculate next_d1 and the adjusted d1 and d2
df_new = df1.where('q!= ""') \
            .withColumn('next_d1', lead('d1').over(w1)) \
            .withColumn('next_flag', coalesce(lead('flag').over(w1), lit(True))) \
            .selectExpr(
                    'id_'
                  , 'q'
                  , 'IF(flag, date_add(d1,1), d1) AS d1'
                  , 'IF((next_d1 != d2) or next_flag, d2, date_sub(d2,1)) AS d2'
             )

df_new.show()
+---+-----+----------+----------+
|id_|    q|        d1|        d2|
+---+-----+----------+----------+
|  1|    A|2018-09-26|2018-10-26|
|  2|    B|2018-06-21|2018-06-26|
|  2|  C B|2018-06-27|2018-07-01|
|  2|C B A|2018-07-02|2018-07-07|
|  2|  B A|2018-07-08|2018-07-19|
|  2|    A|2018-07-20|2019-02-27|
|  2|    A|2019-03-28|2019-06-25|
+---+-----+----------+----------+