创建一个函数，将时间转换为二进制响应变量

Question

我目前有一个RDD，其中有两列

Row(pickup_time=datetime.datetime(2014, 2, 9, 14, 51)
   dropoff_time=datetime.datetime(2014, 2, 9, 14, 58)

我想将这些转换成二进制响应变量，其中1表示夜间时间，0表示白天时间。

我知道我们可以使用UserDefinedFunction来创建一个函数，它可以将这些更改为所需的格式。

例如，我有另一列，它是一个字符串，指定支付类型为'CSH'或'CRD'所以我可以解决这个问题

pay_map = {'CRD':1.0, 'CSH':0.0}
pay_bin = UserDefinedFunction(lambda z: pay_map[z], DoubleType())
df = df.withColumn('payment_type', pay_bin(df['payment_type']))

我如何将同样的逻辑应用到我正在问的问题中?如果有帮助的话，我试着变换这些变量因为我将运行一个决策树。

Answer 1

这里不需要UDF。您可以使用between和类型强制转换:

from pyspark.sql.functions import hour
def in_range(colname, lower_bound=6, upper_bound=17):
    """
    :param colname - Input column name (str)
    :lower_bound   - Lower bound for day hour (int, 0-23) 
    :upper_bound   - Upper bound for day hour (int, 0-23)
    """
    assert 0 <= lower_bound <= 23
    assert 0 <= upper_bound <= 23
    if lower_bound < upper_bound:
        return hour(colname).between(lower_bound, upper_bound).cast("integer")
    else:
        return (
            (hour(colname) >= lower_bound) |
            (hour(colname) <= upper_bound)
        ).cast("integer")

使用例子:

df = sc.parallelize([
    Row(
        pickup_time=datetime.datetime(2014, 2, 9, 14, 51), 
        dropoff_time=datetime.datetime(2014, 2, 9, 14, 58)
    ),
    Row(
        pickup_time=datetime.datetime(2014, 2, 9, 19, 51), 
        dropoff_time=datetime.datetime(2014, 2, 9, 1, 58)
    )
]).toDF()
(df
    .withColumn("dropoff_during_day", in_range("dropoff_time"))
    # between 6pm and 5am
    .withColumn("pickpup_during_night", in_range("pickup_time", 18, 5)))

+--------------------+--------------------+------------------+--------------------+
|        dropoff_time|         pickup_time|dropoff_during_day|pickpup_during_night|
+--------------------+--------------------+------------------+--------------------+
|2014-02-09 14:58:...|2014-02-09 14:51:...|                 1|                   0|
|2014-02-09 01:58:...|2014-02-09 19:51:...|                 0|                   1|
+--------------------+--------------------+------------------+--------------------+