从整数列表中选择值的 PySpark 新列



我在PySpark中有以下数据帧:

|ID    |YearBLT|MinYear|MaxYear|ADP_Range               |
---------------------------------------------------------
|164876|2010   |2004   |2009   |[2004,2009]             |
|164877|2008   |2000   |2011   |[2000, 2002, 2011]      |
|164878|2000   |2003   |2011   |[2003, 2011]            |
|164879|2013   |1999   |2015   |[2003, 2007, 2015, 1999]|

其中,年份BLT是房产建成的年份,ADP_Range表示建筑规范更新的年份,最小年份/最大年份表示ADP范围的最小和最大年份。

我正在尝试添加具有最适用的构建代码的列 (ADP_Year),其逻辑如下:

  • 如果年份BLT小于最小年份,则ADP_Year=="NA">
  • 如果 YearBLT 大于 MaxYear,则ADP_Year == Max(ADP_Range)
  • 如果 YearBLT 介于两者之间,它将选择ADP_Range中低于 YearBLT的最接近的日期

预期输出如下:

|ID    |YearBLT|MinYear|MaxYear|ADP_Range               |ADP_Year|
------------------------------------------------------------------
|164876|2010   |2004   |2009   |[2004,2009]             |2009    |
|164877|2008   |2000   |2011   |[2000, 2002, 2011]      |2002    |
|164878|2000   |2003   |2011   |[2003, 2011]            |NA      |
|164879|2013   |1999   |2015   |[2003, 2007, 2015, 1999]|2007    |

2010> MaxYear,因此它从 MaxYear 中选择值,

2008 年介于 2000 年和 2011 年之间;因为有第三个值 2002,所以选择它,因为它比 2000 年更近

。2000

2013 年介于 1999 年和 2015 年之间;由于有 2007 年和 2015 年的第三和第四个值,因此选择 2007 年

前两种情况很简单,我有它们的工作代码:

dfADP = dfADP.withColumn("ADP_Year",when(dfADP['YearBLT'] < dfADP['MinYear'], lit("NA")
.when(dfADP['YearBLT'] > dfADP['MaxYear'],dfADP['MaxYear'])))

我正在为此旋转轮子,并希望就这是否可能提出一些建议。

首先,让我们找到范围的最大值

from pyspark.sql.functions import array_max, col, expr, when
max_adp_range = array_max("ADP_Range")

和最接近的值:

closest_adp_range = array_max(expr("""
filter(ADP_Range, y -> y < YearBLT)
"""))

并将这两者组合成一个表达式:

adp_year = when(
# If the YearBLT is greater than the MaxYear, ADP_Year == Max(ADP_Range)
col("YearBLT") > col("MaxYear"), max_adp_range
).when(
# If the YearBLT is in between, it chooses 
# the closest date below the YearBLT in the ADP_Range
col("YearBLT").between(col("MinYear"), col("MaxYear")), closest_adp_range
).otherwise(
# If the YearBLT is less than the MinYear, ADP_Year == "NA"
# Note: not required. Included just for clarity.
None
)

最后选择:

df = spark.createDataFrame([                                    
(164876, 2010, 2004, 2009, [2004,2009]),
(164877, 2008, 2000, 2011, [2000, 2002, 2011]),   
(164878, 2000, 2003, 2011, [2003, 2011]),         
(164879, 2013, 1999, 2015, [2003, 2007, 2015, 1999])
], ("id", "YearBLT", "MinYear", "MaxYear", "ADP_Range"))
df.withColumn("ADP_YEAR", adp_year).show()

这应该给出预期的结果:

+------+-------+-------+-------+--------------------+--------+
|    id|YearBLT|MinYear|MaxYear|           ADP_Range|ADP_YEAR|
+------+-------+-------+-------+--------------------+--------+
|164876|   2010|   2004|   2009|        [2004, 2009]|    2009|
|164877|   2008|   2000|   2011|  [2000, 2002, 2011]|    2002|
|164878|   2000|   2003|   2011|        [2003, 2011]|    null|
|164879|   2013|   1999|   2015|[2003, 2007, 2015...|    2007|
+------+-------+-------+-------+--------------------+--------+

array_maxfilter高阶函数都需要 Spark 2.4 或更高版本。在 2.3 或之前,您可以将上述表达式重新定义为

from pyspark.sql.functions import udf
max_adp_range = udf(max, "bigint")("ADP_Range")
closest_adp_range = udf(
lambda xs, y: max(x for x in xs if x < y), "bigint"
)("ADP_Range", "YearBLT")

但是您应该预期会显著的性能损失(单个udf应该更快,但仍然比本机表达式慢)。

最新更新