Как я могу извлечь конкретную подвыборку из фрейма данных и сохранить в другой фрейм данных в pyspark?

Question

У меня есть фрейм данных с именем 'df1', в котором есть X строк, предположим 1000. Я хочу получить конкретный подвыбор этого фрейма данных и сохранить его как другой. Например, я хочу извлечь строки 400–700 из "df1" и сохранить его как "df2".

Я знаю, что одним из возможных способов является получение содержимого 'df1' в виде вектора с помощью:

list = df1.collect()
subsample = list[400:700]
df2 = sc.createDataFrame(subsample, attributes)

Но мой вопрос: есть ли другой способ получить тот же результат, не загружая данные в списке? Я спрашиваю об этом, потому что, когда у вас огромный набор данных, возможно, не будет эффективной загрузки данных с помощью сбора и генерации другого кадра данных.

Благодарю.

Answer 1

Фреймы данных Pyspark не имеют индексов. Вы можете создать один, но обратите внимание, что любая операция перемешивания (group, join...) это произошло до того, как создание индекса могло изменить порядок ваших строк.

import pyspark.sql.functions as psf
start = 400
end = 700
df2 = df1.rdd.zipWithIndex()\
    .map(lambda l: [l[1]] + list(l[0]))\
    .toDF(["index"] + df1.columns)\
    .filter(psf.col("index").between(start, end))

Другой способ - собрать в список только первые строки вашего фрейма данных:

df2 = spark.createDataFrame(df1.head(end)[start:], df1.columns)

или используя Pandas:

df2 = spark.createDataFrame(df1.limit(end).toPandas().iloc[start:, :])