想要对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)]
扫描下方二维码关注领取程序员必备千套ppt模板,300本精选好书,丰富面经: