有没有一种方法可以基于pyspark中的索引对数据帧进行切片？

Question

在python或R中，有一些方法可以使用索引对DataFrame进行切片。

例如，在熊猫中：

df.iloc[5:10,:]

pyspark中是否有类似的方式根据行的位置对数据进行切片？

Answer 1

简短回答

如果您已经有一个索引列（假设它称为'id'），则可以使用pyspark.sql.Column.between进行过滤：

from pyspark.sql.functions import col
df.where(col("id").between(5, 10))

如果您还没有索引列，则可以add one yourself，然后使用上面的代码。您应该基于其他一些列（orderBy("someColumn")）对数据内置一些排序。

---

完整说明

否，除非索引已经以列形式存在，否则无法轻易地按索引对Spark DataFrame进行切片。

Spark DataFrame本质上是无序的，不支持随机访问。 （没有pandas中的内置索引概念）。每行被视为结构化数据的独立集合，因此可以进行分布式并行处理。因此，任何执行者都可以获取任何数据块并对其进行处理，而无需考虑行的顺序。

现在很明显，可以执行涉及排序的操作（lead，lag等），但是这些操作会比较慢因为需要火花才能在执行程序之间重新整理数据。 （数据改组通常是Spark作业中最慢的组件之一。）

相关/深入阅读

• PySpark DataFrames - way to enumerate without converting to Pandas?
• PySpark - get row number for each row in a group
• how to add Row id in pySpark dataframes

Answer 2

您可以将 spark 数据帧转换为 koalas 数据帧。 Koalas 是 Databricks 的一个数据框，它提供了一个几乎像熊猫一样的界面来触发数据框。请参阅此处https://pypi.org/project/koalas/

import databricks.koalas as ks
kdf = ks.DataFrame(your_spark_df)
kdf[0:500] # your indexes here