在 GroupByKey 之后减少 PCollection

Question

我正在尝试根据交易数据生成简单的客户摘要。例如，给定目标交易类型，发生了多少交易，总金额是多少？

原始输入示例：

custid  desc        amount
111     coffee      3.50
111     grocery     23.00
333     coffee      4.00
222     gas station 32.00
222     gas station 55.50
333     coffee      3.00

所需输出示例：

custid nbr_coffee amt_coffee nbr_gas_station amt_gas_station
111    1          3.50       0               0.00
222    0          0          2               87.50
333    2          7.00       0               0

我的目标跑步者将是 Dataflow（但目前使用 DirectRunner 进行测试）。

这是我所拥有的代码 sn-p：

def categorize_coffee(transaction):

    if transaction['trx_desc'] == 'coffee':
        transaction['coffee'] = True
    else:
        transaction['coffee'] = False

    return transaction

def categorize_gas_station(transaction):

    if transaction['trx_desc'] == 'gas station':
        transaction['gas_station'] = True
    else:
        transaction['gas_station'] = False

    return transaction

def summarize_coffee(grouping):

    key, values = grouping
    values = list(values)

    nbr = 0
    amt = 0

    for d in values:
        if d['coffee'] == True:
            nbr+=1
            amt+=d['amount']

    ret_val = {}
    ret_val['cust'] = d['cust']
    ret_val['nbr_coffee'] = nbr
    ret_val['amt_coffee'] = amt

    return ret_val

def summarize_gas_station(grouping):

    key, values = grouping
    values = list(values)

    nbr = 0
    amt = 0

    for d in values:
        if d['gas_station'] == True:
            nbr += 1
            amt += d['amount']

    ret_val = {}
    ret_val['cust'] = d['cust']
    ret_val['nbr_gas_station'] = nbr
    ret_val['amt_gas_station'] = amt

    return ret_val

def create_dict(row):

    vars = row.split(',')
    return {'cust': vars[0], 'trx_desc': str(vars[1]), 'amount': float(vars[2])}

with beam.Pipeline(options=pipeline_options) as p:

    categorized_trx = (
        p | 'get data' >> beam.io.ReadFromText('./test.csv')
        | beam.Map(create_dict)
        | beam.Map(categorize_coffee)
        | beam.Map(categorize_gas_station)
        | beam.Map(lambda trx: (trx['cust'], trx))
        | beam.GroupByKey()
    )

    coffee_trx = (categorized_trx | beam.Map(summarize_coffee))

    gas_station_trx = (categorized_trx | beam.Map(summarize_gas_station))

    result = (coffee_trx, gas_station_trx) | beam.Flatten()

现在的实际结果是：

{'amt_coffee': 7.0, 'cust': u'333', 'nbr_coffee': 2}
{'amt_coffee': 0, 'cust': u'222', 'nbr_coffee': 0}
{'amt_coffee': 3.5, 'cust': u'111', 'nbr_coffee': 1}
{'nbr_gas_station': 0, 'cust': u'333', 'amt_gas_station': 0}
{'nbr_gas_station': 2, 'cust': u'222', 'amt_gas_station': 87.5}
{'nbr_gas_station': 0, 'cust': u'111', 'amt_gas_station': 0}

没有像我预期的那样变平或加入。我是 Beam 新手 - 不确定我是否理解如何正确解决这个问题，所以希望能得到一些见解。

Answer 1

这应该可行：

...

def summarize_coffee(grouping):

    ...

    return (d['cust'], ret_val)


def summarize_gas_station(grouping):

    ...

    return (d['cust'], ret_val)

...

def processJoin(row):
    (customer, data) = row
    coffee_trx=data['coffee_trx']
    gas_station_trx=data['gas_station_trx']
    return (customer, coffee_trx, gas_station_trx)

result = ({coffee_trx: coffee_trx, gas_station_trx: gas_station_trx}
         | 'Group' >> beam.CoGroupByKey()    
         | 'Reshape' >> beam.Map(processJoin)
         | 'Unwind' >> beam.FlatMap(lambda x: x)
         )

Answer 2

Beam 提供Combine 转换，应该允许您组合 PCollection 的元素。看起来对于您的用例，您可以使用 Combine.perKey() 根据键组合 PCollection 的键控元素。作为 combine 函数，你既可以提供一个函数，也可以实现CombineFn。