原創(chuàng)文章，轉(zhuǎn)載請(qǐng)注明原作地址：www.itdecent.cn/p/2db35d7bb92f

在Storm開(kāi)發(fā)過(guò)程中，經(jīng)常性需要將符合某個(gè)條件的消息分發(fā)到同一個(gè)partition。官方提供了一個(gè)Stream.partitionBy("fieldName")的API，可以根據(jù)每條tuple某個(gè)字段進(jìn)行流劃分，劃分方式是：

field.hashCode() mod num-task

雖然Clojure實(shí)際上是運(yùn)行與java平臺(tái)上的一種Lisp方言，但是在java里直接計(jì)算field.hashCode() % num-task，發(fā)現(xiàn)得到的結(jié)果和實(shí)際partitionIndex并不一致。另外，這種原生的根據(jù)字段的哈希值取模進(jìn)行流劃分的方式也過(guò)于單一，因此直接去寫(xiě)了一個(gè)Storm自定義流劃分函數(shù)的方法。實(shí)現(xiàn)代碼如下：

import java.util.Arrays;
import java.util.List;

import org.apache.storm.generated.GlobalStreamId;
import org.apache.storm.grouping.CustomStreamGrouping;
import org.apache.storm.task.WorkerTopologyContext;
import org.apache.storm.tuple.Fields;

public class HashModStreamGrouping implements CustomStreamGrouping {
    
    private List<Integer> targetTasks;
    private String partitionKeyName;
    private int partitionKeyIndex;

    public HashModStreamGrouping(String partitionKeyName) {
        this.partitionKeyName = partitionKeyName;
    }
    
    @Override
    public List<Integer> chooseTasks(int taskId, List<Object> values) {
        String partitionStr = String.valueOf(values.get(this.partitionKeyIndex));
        int partitionVal = getTaskIndex(partitionStr, this.targetTasks.size());
        return Arrays.asList(this.targetTasks.get(partitionVal));
    }

    @Override
    public void prepare(WorkerTopologyContext context, GlobalStreamId streamId,
            List<Integer> targetTasks) {
        this.targetTasks = targetTasks;
        
        Fields outputFields = context.getComponentOutputFields(streamId);
        this.partitionKeyIndex = outputFields.fieldIndex(this.partitionKeyName);
    }
    
    public static int getTaskIndex(String partitionStr, int targetSize) {
        return  (Math.abs(partitionStr.hashCode()) % targetSize );
    }
} 

調(diào)用方式如下，其中partitionValue是用于流劃分的字段名

Stream.partition(new HashModStreamGrouping("partitionValue"))