Python Pyspark：扩展数据集以包括邻居

我是Spark的新手，正在尝试将现有的python应用程序迁移到pyspark

第一个函数之一（在本例中为

f（x）

）应针对数据集中的每个元素运行，但也应考虑数据集中的其他元素

我能得到的最佳简化是以下伪代码：

    def idx_gen_a(x):
        return x-5

    def idx_gen_b(x):
        return x*3

    def f(i, x, dataset):
        elem1 = dataset.get(idx_gen_a(i))
        elem2 = dataset.get(idx_gen_b(i))
        ...
        return some_calculation(x, elem1, elem2, ...)

    def main(dataset):
        result = []
        for i, x in enumerate(dataset):
            result.append(f(i, x,dataset))

---

有没有一种像火花一样的方法

foreachPartition

和

aggregate

似乎不太合适..

我认为您所称的

数据集。获取
大致映射到spark中的
连接。我已经使用pyspark和RDD编写了上述代码的粗略翻译
f1
和
f2
是您的两个功能。您可以使用数据帧执行类似的操作

data = spark.range(10).rdd.map(lambda row: (row[0], row[0] * 10))

def unNest(nested):
  key, ((v1, v2), v3) = nested
  return key, (v1, v2, v3)

def f1(a): return a + 1
def f2(a): return a - 1

one = data.map(lambda pair: (f1(pair[0]), pair[1]))
two = data.map(lambda pair: (f2(pair[0]), pair[1]))
data.join(one).join(two).map(unNest).take(10)

# [(1, (10, 0, 20)),
#  (2, (20, 10, 30)),
#  (3, (30, 20, 40)),
#  (4, (40, 30, 50)),
#  (5, (50, 40, 60)),
#  (6, (60, 50, 70)),
#  (7, (70, 60, 80)),
#  (8, (80, 70, 90))]

有不同类型的联接，例如内部联接和外部联接，但我希望这足以为您指明正确的方向。
请看一下@user10465355，谢谢，我已经看过sparksql的窗口功能，但我不太确定如何将其应用于此情况。。