FirstN.java

package storm.trident.operation.builtin;

import backtype.storm.tuple.Fields;
import java.util.Comparator;
import java.util.PriorityQueue;
import storm.trident.Stream;
import storm.trident.operation.Aggregator;
import storm.trident.operation.Assembly;
import storm.trident.operation.BaseAggregator;
import storm.trident.operation.TridentCollector;
import storm.trident.tuple.TridentTuple;


public class FirstN implements Assembly {

    Aggregator _agg;
    
    public FirstN(int n, String sortField) {
        this(n, sortField, false);
    }
    
    public FirstN(int n, String sortField, boolean reverse) {
        if(sortField!=null) {
            _agg = new FirstNSortedAgg(n, sortField, reverse);
        } else {
            _agg = new FirstNAgg(n);
        }
    }
    
    @Override
    public Stream apply(Stream input) {
        Fields outputFields = input.getOutputFields();
        return input.partitionAggregate(outputFields, _agg, outputFields)
                    .global()
                    .partitionAggregate(outputFields, _agg, outputFields);             
    }
    
    public static class FirstNAgg extends BaseAggregator<FirstNAgg.State> {
        int _n;
        
        public FirstNAgg(int n) {
            _n = n;
        }
        
        static class State {
            int emitted = 0;
        }
        
        @Override
        public State init(Object batchId, TridentCollector collector) {
            return new State();
        }

        @Override
        public void aggregate(State val, TridentTuple tuple, TridentCollector collector) {
            if(val.emitted < _n) {
                collector.emit(tuple);
                val.emitted++;
            }
        }

        @Override
        public void complete(State val, TridentCollector collector) {
        }
        
    }
    
    public static class FirstNSortedAgg extends BaseAggregator<PriorityQueue> {

        int _n;
        String _sortField;
        boolean _reverse;
        
        public FirstNSortedAgg(int n, String sortField, boolean reverse) {
            _n = n;
            _sortField = sortField;
            _reverse = reverse;
        }

        @Override
        public PriorityQueue init(Object batchId, TridentCollector collector) {
            return new PriorityQueue(_n, new Comparator<TridentTuple>() {
                @Override
                public int compare(TridentTuple t1, TridentTuple t2) {
                    Comparable c1 = (Comparable) t1.getValueByField(_sortField);
                    Comparable c2 = (Comparable) t2.getValueByField(_sortField);
                    int ret = c1.compareTo(c2);
                    if(_reverse) ret *= -1;
                    return ret;
                }                
            });
        }

        @Override
        public void aggregate(PriorityQueue state, TridentTuple tuple, TridentCollector collector) {
            state.add(tuple);
        }

        @Override
        public void complete(PriorityQueue val, TridentCollector collector) {
            int total = val.size();
            for(int i=0; i<_n && i < total; i++) {
                TridentTuple t = (TridentTuple) val.remove();
                collector.emit(t);
            }
        }
    }    
}