如何遍历数据框并将数组添加到每一行

时间:2019-09-03 06:09:15

标签: pyspark pyspark-dataframes

我需要遍历pyspark数据框,并按活动月份数逐行显示。我主要关心的是内存管理,当我试图将数据放入配置单元中并消耗该过程所花费的时间时。

我曾经使用Empid date_active date_end 1234 2012-01-01 2012-10-27 2345 2012-01-01 2012-12-31 3456 2012-01-01 2012-08-15 的收集 但这杀死了我的代码的性能。

输入

EmpId effective_Month
1234  Jan-12
1234  Feb-12
1234  ....
1234  Oct-12
2345  Jan-12
2345  Feb-12
2345  ....
2345  Dec-12

输出

LocalizationRoute

2 个答案:

答案 0 :(得分:0)

通常,我建议尝试使用udf。我曾经遇到过类似的问题,但是我认为解决方案仍然太复杂了。

相反,您可以改变对问题的看法。当您知道所有(sequential? 1) => false的{​​{1}}和min(date_active)时,您可以遍历所有月份,例如max(date_end)并将其另存为数据框。 现在,您(广播)将结果数据框与表的每一行相连。最后,您只需要一个简单的过滤器,例如:EmpIds。 最后,您应该将'yyyy-mm-01'转换为您喜欢的字符串格式。

如果广播表很小,则广播联接非常快,因此运行时在这里应该不是问题。

答案 1 :(得分:0)

您可以在数据框API中解决以下问题

创建示例数据框

  df = spark.createDataFrame([["123","2012-01-01","2012-10-01"],['234', '2012-01-01', '2012-05-01'],["345","2012-01-01","2012-11-01"]], ("age","date_active", "date_end"))

    +---+-----------+----------+
    |age|date_active|  date_end|
    +---+-----------+----------+
    |123| 2012-01-01|2012-10-01|
    |234| 2012-01-01|2012-05-01|
    |345| 2012-01-01|2012-11-01|
    +---+-----------+----------+

将数据类型从字符串更改为时间戳记

df = df.withColumn('date_active', df['date_active'].cast('timestamp'))\
.withColumn('date_end', df['date_end'].cast('timestamp'))

在下面的代码中添加月份列

from pyspark.sql import functions as f

df.withColumn('month_diff', f.months_between('date_end', 'date_active')).withColumn("repeat", f.expr("split(repeat(',', month_diff), ',')"))\
.select("*", f.posexplode("repeat").alias("date", "val")).withColumn("date", f.expr("add_months(date_active, date)"))\
.withColumn('month', f.date_format('date','MMM')).select(['age', 'date_active', 'date_end', 'month']).show()


---+-----------+----------+-----+
|age|date_active|  date_end|month|
+---+-----------+----------+-----+
|123| 2012-01-01|2012-10-01|  Jan|
|123| 2012-01-01|2012-10-01|  Feb|
|123| 2012-01-01|2012-10-01|  Mar|
|123| 2012-01-01|2012-10-01|  Apr|
|123| 2012-01-01|2012-10-01|  May|
|123| 2012-01-01|2012-10-01|  Jun|
|123| 2012-01-01|2012-10-01|  Jul|
|123| 2012-01-01|2012-10-01|  Aug|
|123| 2012-01-01|2012-10-01|  Sep|
|123| 2012-01-01|2012-10-01|  Oct|
|234| 2012-01-01|2012-05-01|  Jan|
|234| 2012-01-01|2012-05-01|  Feb|
|234| 2012-01-01|2012-05-01|  Mar|
|234| 2012-01-01|2012-05-01|  Apr|
|234| 2012-01-01|2012-05-01|  May|
|345| 2012-01-01|2012-11-01|  Jan|
|345| 2012-01-01|2012-11-01|  Feb|
|345| 2012-01-01|2012-11-01|  Mar|
|345| 2012-01-01|2012-11-01|  Apr|
|345| 2012-01-01|2012-11-01|  May|
+---+-----------+----------+-----+