创建多个标记输出时Google数据流超出堆

Question

我有许多未分区的大型BigQuery表和文件，我希望以各种方式对它们进行分区。因此，我决定尝试编写一个数据流作业来实现这一点。我认为这工作很简单。我尝试使用泛型编写，以便轻松地应用TextIO和BigQueryIO源代码。它在小型表上工作得很好，但在大型表上运行时，我总是得到Java.lang.outofMemoryError：Java堆空间。

在我的主类中，我要么读取一个带有目标键的文件（由另一个DF作业生成），要么对一个BigQuery表运行一个查询，以获得要根据其进行分解的键列表。我的主类如下所示：

Pipeline sharder = Pipeline.create(opts);

// a functional interface that shows the tag map how to get a tuple tag
KeySelector<String, TableRow> bqSelector = (TableRow row) -> (String) row.get("COLUMN") != null ? (String) row.get("COLUMN") : "null";

// a utility class to store a tuple tag list and hash map of String TupleTag
TupleTagMap<String, TableRow> bqTags = new TupleTagMap<>(new ArrayList<>(inputKeys),bqSelector);

// custom transorm
ShardedTransform<String, TableRow> bqShard = new ShardedTransform<String, TableRow>(bqTags, TableRowJsonCoder.of());

String source = "PROJECTID:ADATASET.A_BIG_TABLE";
String destBase = "projectid:dataset.a_big_table_sharded_";

TableSchema schema = bq.tables().get("PROJECTID","ADATASET","A_BIG_TABLE").execute().getSchema();


PCollectionList<TableRow> shards = sharder.apply(BigQueryIO.Read.from(source)).apply(bqShard);
for (PCollection<TableRow> shard : shards.getAll()) {
    String shardName = StringUtils.isNotEmpty(shard.getName()) ? shard.getName() : "NULL";
    shard.apply(BigQueryIO.Write.to(destBase + shardName)
            .withWriteDisposition(WriteDisposition.WRITE_TRUNCATE)
            .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
            .withSchema(schema));
    System.out.println(destBase+shardName);
} 
sharder.run();

我生成一组tupleTags用于自定义转换。我创建了一个实用工具类，它存储TupleTagList和HashMap，以便通过键引用元组标记：

public class TupleTagMap<Key, Type> implements Serializable {

private static final long serialVersionUID = -8762959703864266959L;
final private TupleTagList tList;
final private Map<Key, TupleTag<Type>> map;
final private KeySelector<Key, Type> selector;

public TupleTagMap(List<Key> t, KeySelector<Key, Type> selector) {
    map = new HashMap<>();
    for (Key key : t)
        map.put(key, new TupleTag<Type>());
    this.tList = TupleTagList.of(new ArrayList<>(map.values()));
    this.selector = selector;

}

public Map<Key, TupleTag<Type>> getMap() {
    return map;
}

public TupleTagList getTagList() {
    return tList;
}

public TupleTag<Type> getTag(Type t){
    return map.get(selector.getKey(t));
}

然后我有一个自定义转换，它基本上有一个函数，使用元组映射输出pCollectionTuple，然后将其移动到pCollectionList以返回到主类：

public class ShardedTransform<Key, Type> extends
    PTransform<PCollection<Type>, PCollectionList<Type>> {


private static final long serialVersionUID = 3320626732803297323L;
private final TupleTagMap<Key, Type> tags;
private final Coder<Type> coder;


public ShardedTransform(TupleTagMap<Key, Type> tags, Coder<Type> coder) {
    this.tags = tags;
    this.coder = coder;
}

@Override
public PCollectionList<Type> apply(PCollection<Type> in) {

    PCollectionTuple shards = in.apply(ParDo.of(
            new ShardFn<Key, Type>(tags)).withOutputTags(
            new TupleTag<Type>(), tags.getTagList()));

    List<PCollection<Type>> shardList = new ArrayList<>(tags.getMap().size());

    for (Entry<Key, TupleTag<Type>> e : tags.getMap().entrySet()){
        PCollection<Type> shard = shards.get(e.getValue()).setName(e.getKey().toString()).setCoder(coder);
        shardList.add(shard);
    }
        return PCollectionList.of(shardList);
    } 
}

实际的DoFn非常简单，它只使用主类中提供的lambda，在hash映射中找到匹配的元组标记，用于侧输出：

public class ShardFn<Key, Type> extends DoFn<Type, Type> {

private static final long serialVersionUID = 961325260858465105L;

private final TupleTagMap<Key, Type> tags;

ShardFn(TupleTagMap<Key, Type> tags) {

    this.tags = tags;
}

@Override
public void processElement(DoFn<Type, Type>.ProcessContext c)
        throws Exception {
    Type element = c.element();
    TupleTag<Type> tag = tags.getTag(element);

    if (tag != null)
        c.sideOutput(tags.getTag(element), element);
    } 
}

Answer 1

Beam模型目前对动态分区/大量分区没有很好的支持。您的方法在图构建时选择碎片的数量，然后产生的pardo可能会全部融合在一起，因此每个工作人员都试图同时写入80个不同的BQ表。每次写入都需要一些本地缓冲，所以可能太多了。

还有一种替代方法可以跨表（但不跨元素）进行并行化。如果您有大量相对较小的输出表，这将很好地工作。使用ParDo标记每个元素的表，然后执行GroupByKey。这为您提供了pcollection >> 。然后通过将元素写入表来处理每个kv > 。

不幸的是，现在您必须用BQ手写才能使用此选项。我们正在考虑用内置的支持来扩展接收器API。由于Dataflow SDK作为Apache Beam的一部分正在进一步开发，我们在这里跟踪该请求:https://issues.Apache.org/jira/browse/Beam-92