在Pyspark中选择一个范围

Question

我在python中有一个spark数据帧。并且，它是根据列进行排序的。如何选择特定范围的数据（例如中间数据的50％）？例如，如果我有1M数据，我想要从250K到750K索引的take数据。如果不在pyspark中使用collect，我怎么能这样做？

更确切地说，我想要像take函数一样在一个范围之间得到结果。例如，像take(250000, 750000)。

Answer 1

以下是在pyspark DF中选择范围的一种方法：

创建DF

df = spark.createDataFrame(
    data = [(10, "2018-01-01"), (22, "2017-01-01"), (13, "2014-01-01"), (4, "2015-01-01")\
           ,(35, "2013-01-01"),(26, "2016-01-01"),(7, "2012-01-01"),(18, "2011-01-01")]
    , schema =  ["amount", "date"]
)

df.show()

+------+----------+
|amount|      date|
+------+----------+
|    10|2018-01-01|
|    22|2017-01-01|
|    13|2014-01-01|
|     4|2015-01-01|
|    35|2013-01-01|
|    26|2016-01-01|
|     7|2012-01-01|
|    18|2011-01-01|
+------+----------+

排序（日期）和插入索引（基于行号）

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

w = Window.orderBy("date")
df = df.withColumn("index", F.row_number().over(w))

df.show()

+------+----------+-----+
|amount|      date|index|
+------+----------+-----+
|    18|2011-01-01|    1|
|     7|2012-01-01|    2|
|    35|2013-01-01|    3|
|    13|2014-01-01|    4|
|     4|2015-01-01|    5|
|    26|2016-01-01|    6|
|    22|2017-01-01|    7|
|    10|2018-01-01|    8|
+------+----------+-----+

获取所需范围（假设想要第3行和第6行之间的所有内容）

df1=df.filter(df.index.between(3, 6))

df1.show()
+------+----------+-----+
|amount|      date|index|
+------+----------+-----+
|    35|2013-01-01|    3|
|    13|2014-01-01|    4|
|     4|2015-01-01|    5|
|    26|2016-01-01|    6|
+------+----------+-----+

Answer 2

这很简单，使用between，例如假设你的排序列名是index -

df_sample = df.select(df.somecolumn, df.index.between(250000, 750000)) 

创建新的数据帧df_sample后，您可以根据需要执行任何操作（包括获取或收集）。