在Apache Beam的窗口中聚合数据

时间:2019-07-11 17:12:12

标签: python google-cloud-dataflow apache-beam

我收到一个复杂且嵌套的JSON对象流作为我对管道的输入。

我的目标是创建小批量,馈送给另一个pubsub主题以进行下游处理。我在beam.beam.GroupByKey()函数中苦苦挣扎-从我读到的内容来看,这是尝试进行汇总的正确方法。

一个简化的示例,输入事件:

{ data:['a', 'b', 'c'], url: 'websiteA.com' }
{ data:['a', 'b', 'c'], url: 'websiteA.com' }
{ data:['a'], url: 'websiteB.com' }

我正在尝试创建以下内容:

{
'websiteA.com': {a:2, b:2, c:2},
'websiteB.com': {a:1},
}

我的问题在于,尝试对最简单的元组引发ValueError: too many values to unpack的任何事情进行分组。

我可以分两个步骤来运行它,但是从我使用beam.GroupByKey()的阅读来看,它很昂贵,因此应尽量减少。

根据@Cubez的答案进行编辑。

这是我的Combine函数,似乎起作用了一半:(

class MyCustomCombiner(beam.CombineFn):
  def create_accumulator(self):
    logging.info('accum_created') #Logs OK!
    return {}

  def add_input(self, counts, input):
    counts = {}
    for i in input:
      counts[i] = 1
    logging.info(counts) #Logs OK!
    return counts

  def merge_accumulators(self, accumulators):
    logging.info('accumcalled') #never logs anything
    c = collections.Counter()
    for d in accumulators:
      c.update(d)
    logging.info('accum: %s', accumulators) #never logs anything
    return dict(c)

  def extract_output(self, counts):
    logging.info('Counts2: %s', counts) #never logs anything
    return counts

似乎add_input过去什么都没叫?

添加管道代码:

with beam.Pipeline(argv=pipeline_args) as p:
    raw_loads_dict = (p 
      | 'ReadPubsubLoads' >> ReadFromPubSub(topic=PUBSUB_TOPIC_NAME).with_output_types(bytes)
      | 'JSONParse' >> beam.Map(lambda x: json.loads(x))
    )
    fixed_window_events = (raw_loads_dict
      | 'KeyOnUrl' >> beam.Map(lambda x: (x['client_id'], x['events']))
      | '1MinWindow' >> beam.WindowInto(window.FixedWindows(60))
      | 'CustomCombine' >> beam.CombinePerKey(MyCustomCombiner())
    )
    fixed_window_events | 'LogResults2' >> beam.ParDo(LogResults())

1 个答案:

答案 0 :(得分:3)

这是需要使用combiners的完美示例。这些是用于汇总或组合多个工作程序中的集合的转换。正如文档所说,CombineFns的工作方式是读取您的元素(beam.CombineFn.add_input),合并多个元素(beam.CombineFn.merge_accumulators),然后最终输出最终的组合值(beam.CombineFn.extract_output)。有关父类here的信息,请参见Python文档。

例如,创建一个组合器以输出一组数字的平均值,如下所示:

class AverageFn(beam.CombineFn):
  def create_accumulator(self):
    return (0.0, 0)

  def add_input(self, sum_count, input):
    (sum, count) = sum_count
    return sum + input, count + 1

  def merge_accumulators(self, accumulators):
    sums, counts = zip(*accumulators)
    return sum(sums), sum(counts)

  def extract_output(self, sum_count):
    (sum, count) = sum_count
    return sum / count if count else float('NaN')

pc = ...
average = pc | beam.CombineGlobally(AverageFn())

对于您的用例,我建议如下:

values = [
          {'data':['a', 'b', 'c'], 'url': 'websiteA.com'},
          {'data':['a', 'b', 'c'], 'url': 'websiteA.com'},
          {'data':['a'], 'url': 'websiteB.com'}
]

# This counts the number of elements that are the same.
def combine(counts):
  # A counter is a dictionary from keys to the number of times it has
  # seen that particular key.
  c = collections.Counter()
  for d in counts:
    c.update(d)
  return dict(c)

with beam.Pipeline(options=pipeline_options) as p:
  pc = (p
        # You should replace this step with reading data from your
        # source and transforming it to the proper format for below.
        | 'create' >> beam.Create(values)

        # This step transforms the dictionary to a tuple. For this
        # example it returns:
        # [ ('url': 'websiteA.com', 'data':['a', 'b', 'c']),
        #   ('url': 'websiteA.com', 'data':['a', 'b', 'c']),
        #   ('url': 'websiteB.com', 'data':['a'])]
        | 'url as key' >> beam.Map(lambda x: (x['url'], x['data']))

        # This is the magic that combines all elements with the same
        # URL and outputs a count based on the keys in 'data'.
        # This returns the elements:
        # [ ('url': 'websiteA.com', {'a': 2, 'b': 2, 'c': 2}),
        #   ('url': 'websiteB.com', {'a': 1})]
        | 'combine' >> beam.CombinePerKey(combine))

  # Do something with pc
  new_pc = pc | ...