pyspark中dataframe切片

想要对pyspark中dataframe实现pandas.dataframe中iloc的切片功能，发现spark中没有相关函数可以直接实现该功能，因此自己琢磨了一个方法。
首先创建一个dataframe。

dfs = spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["letter", "name"])

长这样子
±-----±—+
|letter|name|
±-----±—+
| a| 1|
| b| 2|
| c| 3|
±-----±—+

# 定义切片函数
def getrows(df, rownums=None):
    return df.rdd.zipWithIndex().filter(lambda x: x[1] in rownums).map(lambda x: x[0])

#取0和1行
getrows(dfs, rownums=[0,1]).collect()

结果：
[Row(letter=u’a’, name=1), Row(letter=u’b’, name=2)]

#取0和2行
getrows(dfs, rownums=[0,2]).collect()

[Row(letter=u’a’, name=1), Row(letter=u’c’, name=3)]

切片后是rdd，再转换成dataframe即可。

getrows(dfs, rownums=[0,2]).toDF().show()

±-----±—+
|letter|name|
±-----±—+
| a| 1|
| c| 3|
±-----±—+
搞定。
再提一下实现该功能主要用到的函数zipWithIndex()，该函数将RDD中的元素和这个元素在RDD中的ID（索引号）组合成键/值对,如下。

dfs.rdd.zipWithIndex().collect()

[(Row(letter=u’a’, name=1), 0),
(Row(letter=u’b’, name=2), 1),
(Row(letter=u’c’, name=3), 2)]