Выберите диапазон в Pyspark

Question

У меня есть искровой датафрейм в python.И это было отсортировано на основе столбца.Как я могу выбрать определенный диапазон данных (например, 50% данных в середине)?Например, если у меня есть данные 1М, я хочу take данных с индексом от 250К до 750К.Как я могу сделать это без использования collect в pyspark?

Если быть более точным, я хочу что-то вроде take функции для получения результатов между диапазонами.Например, что-то вроде take(250000, 750000).

Answer 1

Вот один из способов выбрать диапазон в Dyspark 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, вы можете выполнить любую операцию (включая take илисобирать) согласно вашей потребности.