java - PCollection 与其自身的笛卡尔积

Question

假设我有一个类型为 KV<String, Integer> 的有界 PCollection p 。假设 p 无法放入内存，因此不能成为 DoFn 的侧面输入.

示例p:

("foo", 0)
("bar", 1)
("baz", 2)

如何求 p 与其自身的笛卡尔积？

例如，p x p可能如下所示:

("foo+foo", [("foo", 0), ("foo", 0)])
("foo+bar", [("foo", 0), ("bar", 1)])
("foo+baz", [("foo", 0), ("baz", 2)])
("bar+foo", [("bar", 1), ("foo", 0)])
("bar+bar", [("bar", 1), ("bar", 1)])
("bar+baz", [("bar", 1), ("baz", 2)])
("baz+foo", [("baz", 2), ("foo", 0)])
("baz+bar", [("baz", 2), ("bar", 1)])
("baz+baz", [("baz", 2), ("baz", 2)])

Answer 1

正如您所猜测的，最简单的方法是使用一个 DoFn 将您的 PCollection 作为主输入和辅助输入进行处理。

如果这不起作用，因为 PCollection 太大而无法放入内存，您可以将其划分为 N 个不相交的 PCollection，将其传递给每个 PCollection，然后展平结果。例如，您可以编写类似的内容

class CrossProduct(beam.PTransform):
  def expand(self, pcoll):
    N = 10
    parts = pcoll | beam.Partition(lambda element, n: hash(element) % n, N)
    cross_parts = [
        pcoll | str(ix) >> beam.FlatMap(
            lambda x, side: [(x, s) for s in side],
            beam.pvalue.AsIter(part))
        for ix, part in enumerate(parts)]
    return cross_parts | beam.Flatten()

output = input | CrossProduct()

但请注意，除非 PCollection 的元素特别大，否则如果 PCollection 无法装入内存，则其叉积可能会过大而无法生成(和处理)。