python之从整数列表中选择值的 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]|

YearBLT 是特性 build 的年份，ADP_Range 代表建筑规范更新的年份，MinYear/MaxYear 代表 ADP 范围的最小和最大年份。

我正在尝试添加具有最适用建筑规范的列 (ADP_Year)，其逻辑如下:

如果 YearBLT 小于 MinYear，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 < MinYear 所以它设置为 NA

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_max 和filter 高阶函数都需要 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 应该更快，但仍然比本地表达式慢)。