spark sql到最近假期的距离

Question

在熊猫中我有一个类似于

的功能

indices = df.dateColumn.apply(holidays.index.searchsorted)
df['nextHolidays'] = holidays.index[indices]
df['previousHolidays'] = holidays.index[indices - 1]

计算到最近假​​期的距离并将其存储为新列。

searchsorted http://pandas.pydata.org/pandas-docs/version/0.18.1/generated/pandas.Series.searchsorted.html 是 pandas 的一个很好的解决方案，因为这给了我下一个假期的索引，而没有很高的算法复杂度 Parallelize pandas apply 例如这种方法比并行循环要快得多。

如何在 spark 或 hive 中实现这一点？

Answer 1

这可以使用聚合来完成，但这种方法比 pandas 方法具有更高的复杂性。但是您可以使用 UDF 实现类似的性能。它不会像 pandas 那样优雅，但是：

假设这个假期数据集：

holidays = ['2016-01-03', '2016-09-09', '2016-12-12', '2016-03-03']
index = spark.sparkContext.broadcast(sorted(holidays))

以及数据框中 2016 年日期的数据集：

from datetime import datetime, timedelta
dates_array = [(datetime(2016, 1, 1) + timedelta(i)).strftime('%Y-%m-%d') for i in range(366)]
from pyspark.sql import Row
df = spark.createDataFrame([Row(date=d) for d in dates_array])

UDF 可以使用 pandas searchsorted，但需要在执行程序上安装 pandas。 Insted 你可以像这样使用plan python：

def nearest_holiday(date):
    last_holiday = index.value[0]
    for next_holiday in index.value:
        if next_holiday >= date:
            break
        last_holiday = next_holiday
    if last_holiday > date:
        last_holiday = None
    if next_holiday < date:
        next_holiday = None
    return (last_holiday, next_holiday)


from pyspark.sql.types import *
return_type = StructType([StructField('last_holiday', StringType()), StructField('next_holiday', StringType())])

from pyspark.sql.functions import udf
nearest_holiday_udf = udf(nearest_holiday, return_type)

并且可以和withColumn一起使用：

df.withColumn('holiday', nearest_holiday_udf('date')).show(5, False)

+----------+-----------------------+
|date      |holiday                |
+----------+-----------------------+
|2016-01-01|[null,2016-01-03]      |
|2016-01-02|[null,2016-01-03]      |
|2016-01-03|[2016-01-03,2016-01-03]|
|2016-01-04|[2016-01-03,2016-03-03]|
|2016-01-05|[2016-01-03,2016-03-03]|
+----------+-----------------------+
only showing top 5 rows