BatchCollector 自定义微批处理工具

需求

在处理大量数据的，希望每5000条保存一次Elasticsearch或者MongoDB。

处理方法

1. BatchCollector实现

package com.oppo.ait.business.utils;

import com.google.common.base.Preconditions;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.function.*;
import java.util.stream.Collector;

/**
 * 批量提交工具集
 *
 * @param <T>
 * @author 
 */
public class BatchCollector<T> implements Collector<T, List<T>, List<T>> {

	private final int batchSize;
	private final Consumer<List<T>> batchProcessor;

	public BatchCollector(int batchSize, Consumer<List<T>> batchProcessor) {
		Preconditions.checkNotNull(batchProcessor);
		Preconditions.checkArgument(batchSize > 0, "batchSize必须大于0");
		this.batchSize = batchSize;
		this.batchProcessor = batchProcessor;
	}

	@Override
	public Supplier<List<T>> supplier() {
		return ArrayList::new;
	}

	@Override
	public BiConsumer<List<T>, T> accumulator() {
		return (ts, t) -> {
			ts.add(t);
			if (ts.size() >= batchSize) {
				batchProcessor.accept(ts);
				ts.clear();
			}
		};
	}

	@Override
	public BinaryOperator<List<T>> combiner() {
		return (ts, ots) -> {
			if (ts.size() > 0) {
				batchProcessor.accept(ts);
			}
			if (ots.size() > 0) {
				batchProcessor.accept(ots);
			}
			return Collections.emptyList();
		};
	}

	@Override
	public Function<List<T>, List<T>> finisher() {
		return ts -> {
			if (ts.size() > 0) {
				batchProcessor.accept(ts);
			}
			return Collections.emptyList();
		};
	}

	@Override
	public Set<Characteristics> characteristics() {
		return Collections.emptySet();
	}
}

2. 代码使用

 /**
   * 需要写入的stream数据到HDFS
   *
   * @param dataType 该数据的类型，决定了如何组织最终的字符串为HDFS的每一行，详见 {@link HdfsWriter#getFunctionChain}
   */
  public void writeToHdfs(DataType dataType, Stream<BotAnnotationData> stream) {
    //过滤掉不存在domain的数据，domain是文件路径的一部分
    stream
        .filter(botAnnotationData -> StringUtils.isNotEmpty(botAnnotationData.getDomain()))
        .collect(Collectors.groupingBy(BotAnnotationData::getDomain,
            new BatchCollector<>(5000, list -> {
              String domain = list.stream().findFirst().map(BotAnnotationData::getDomain).get();
              //如果不存在该domain的HDFS文件句柄，则创建一个句柄
              FSDataOutputStream fsDataOutputStream = writerMap
                  .computeIfAbsent(domain, (k) -> createTextFileWriter(k, dataType));
              //刷新数据到HDFS
              flushData(fsDataOutputStream, list, dataType, this::parseBotAnnotationData);
            })));
  }