按日期将pyspark数据集分为两部分

Question

我有pyspark数据集，我想按datetime列将其分为训练和测试数据集，其中训练数据集的日期时间应小于datetime列的中位数，而测试数据集应具有其余时间

我试图按日期时间列对数据集进行排序，然后选择前半部分。但这只能解决火车零件问题，我不知道如何从PySpark的初始数据集中“减去”火车数据

train = data.orderBy('datetime').limit(data.count() // 2)
# test = ?

如果PySpark具有与Pandas tail（）函数的类比，那会很好，但没有。

Answer 1

您可以添加一个对日期时间进行排名的列，然后使用该排名对数据框进行分区。 percent_rank函数给出百分位数iirc。

from pyspark.sql import functions as F
from pyspark.window import Window

data_win = Window.partitionBy().orderBy('datetime')
dt_rank = data.withColumn('percent_rank', F.percent_rank().over(data_win))
train = dt_rank.filter(F.col('percent_rank') <= 0.5)
test = dt_rank.filter(F.col('percent_rank') > 0.5)