如何从pyspark中的数据框中选择一系列行

Posted 2023-04-18

【中文标题】如何从pyspark中的数据框中选择一系列行

【英文标题】：How to select a range of rows from a dataframe in pyspark

【发布时间】：2019-04-15 13:51:59

【问题描述】：

我有一个包含 10609 行的数据框，我想一次将 100 行转换为 JSON 并将它们发送回网络服务。

我尝试过使用类似 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

【问题讨论】：

Is there a way to slice dataframe based on index in pyspark?的可能重复

【参考方案1】：

您必须创建一个行号列，该列将为列分配序号，并使用该列通过过滤器获取范围内的数据。

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

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

我正在使用一个虚拟静态列lit('a') 来生成row_num。请根据您的实际数据更新以下逻辑（这会生成 row_num）。

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|
#+----+-------+

【讨论】：

谢谢。这会给我前 n 行，但不是中间位置的 n 行。

我已经编辑了我的答案，现在您可以将范围传递给一定数量的记录。确保根据实际数据修改 row_num 的逻辑。