如何在PySpark中从数据框中选择一系列行

Question

我有一个包含10609行的数据框，我想将每次100行转换成JSON格式并发送回Web服务。

我尝试使用SQL的LIMIT子句，例如：

temptable = spark.sql("select item_code_1 from join_table limit 100")

这会返回前100行，但如果我想要接下来的100行，我尝试了这个方法但没有成功。

temptable = spark.sql("select item_code_1 from join_table limit 100, 200")

错误：Py4JJavaError：调用o22.sql时发生错误。： org.apache.spark.sql.catalyst.parser.ParseException：不匹配的输入 “，”预期（第1行，第44个位置）

== SQL ==

select item_code_1 from join_table limit 100, 200

- Sridhar Lanka

---

可能是Is there a way to slice dataframe based on index in pyspark?的重复问题。 - pault

Answer 1

您需要创建一个行号列，它将为列分配顺序编号，并使用该列通过筛选器在范围内获取数据。

df = spark.createDataFrame([('a',),
                            ('b',),
                            ('c',),
                            ('d',),
                            ('e',)
                            ],'item : string')
df.show()

#+----+
#|item|
#+----+
#|   a|
#|   b|
#|   c|
#|   d|
#|   e|
#+----+

我使用一个虚拟静态列 lit('a') 来生成行号。请根据你的实际数据更新以下逻辑（此逻辑生成行号）。

partitionBy(lit('a')).orderBy(lit('a')

数据框示例-

from pyspark.sql.functions import lit,row_number,col
from pyspark.sql.window import Window

w = Window().partitionBy(lit('a')).orderBy(lit('a'))

df1 = df.withColumn("row_num", row_number().over(w))

df1.filter(col("row_num").between(1,2)).show()     

#+----+-------+
#|item|row_num|
#+----+-------+
#|   a|      1|
#|   b|      2|
#+----+-------+

df1.filter(col("row_num").between(3,4)).show()

#+----+-------+
#|item|row_num|
#+----+-------+
#|   c|      3|
#|   d|      4|
#+----+-------+

Spark SQL 示例 -

df1.createOrReplaceTempView("dfTable")

spark.sql("SELECT * FROM dfTable WHERE row_num between 1 and 2").show()

#+----+-------+
#|item|row_num|
#+----+-------+
#|   a|      1|
#|   b|      2|
#+----+-------+

- Shantanu Sharma

---

谢谢。这将给我前n行，但不是从中间位置开始的n行。 - Sridhar Lanka

我已经编辑了我的答案，现在你可以传递范围到特定数量的记录。请确保根据你的实际数据修改row_num的逻辑。 - Shantanu Sharma