Skip to main content
 首页 » 编程设计

python之按行索引拆分 Spark 数据帧

2025年05月04日29shihaiming

我想按行顺序拆分数据框。如果有 100 行,则需要拆分成 4 个相等的数据帧应该具有索引 0-24、25-49、分别为 50-74 和 75-99。

唯一可用的预定义函数是 randomSplit。但是 randomSplit 在拆分之前随机化数据。 我想到的另一种方法是使用 count 缩减操作找到数据的计数,然后使用 take 继续提取数据,但它非常昂贵。 有没有其他方法可以在保持相同顺序的情况下实现上述目标?

请您参考如下方法:

您可以使用 monotonically_increasing_id获取行号(如果您还没有)然后 ntile在行号窗口上拆分成你想要的任意多个分区:

from pyspark.sql.window import Window 
from pyspark.sql.functions import monotonically_increasing_id, ntile 
 
values = [(str(i),) for i in range(100)] 
df = spark.createDataFrame(values, ('value',)) 
 
def split_by_row_index(df, num_partitions=4): 
    # Let's assume you don't have a row_id column that has the row order 
    t = df.withColumn('_row_id', monotonically_increasing_id()) 
    # Using ntile() because monotonically_increasing_id is discontinuous across partitions 
    t = t.withColumn('_partition', ntile(num_partitions).over(Window.orderBy(t._row_id)))  
    return [t.filter(t._partition == i+1).drop('_row_id', '_partition') for i in range(partitions)] 
 
[i.collect() for i in split_by_row_index(df)]