com.flaptor.indextank.util.Union.java Source code

Introduction

Here is the source code for com.flaptor.indextank.util.Union.java

Source

/*
 * Copyright (c) 2011 LinkedIn, Inc
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package com.flaptor.indextank.util;

import java.util.BitSet;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;

import com.flaptor.util.Pair;
import com.flaptor.util.CollectionsUtil.PeekingIterator;
import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

public abstract class Union<K, V extends Comparable<V>> extends AbstractSkippableIterable<V> {

    private Iterable<SkippableIterable<K>> cursors;

    public Union(Iterable<SkippableIterable<K>> cursors) {
        this.cursors = cursors;
    }

    protected abstract V transform(K k);

    protected boolean shouldUse(V v, List<K> ks) {
        return true;
    }

    protected int comp(K a, K b) {
        return transform(a).compareTo(transform(b));
    }

    private static class Head<K> {
        final int position;
        K value;

        public Head(int position) {
            this.position = position;
        }
    }

    @Override
    public SkippableIterator<V> iterator() {
        return new AbstractSkippableIterator<V>() {
            private List<PeekingSkippableIterator<K>> iterators = Lists
                    .newArrayList(Iterables.transform(cursors, Union.<K>peekingIteratorFunction()));
            /*
             * SortedSet with all the heads (first elements) for all the iterators and the position in the iterators list of the
             * SkippableIterator the element is in.
             */
            private TreeSet<Head<K>> heads = new TreeSet<Head<K>>(new Comparator<Head<K>>() {
                @Override
                public int compare(Head<K> o1, Head<K> o2) {
                    int compare = comp(o1.value, o2.value);
                    if (compare != 0) {
                        return compare;
                    } else {
                        return o1.position - o2.position;
                    }
                }
            });

            private List<Head<K>> prebuiltHeads = Lists.newArrayListWithExpectedSize(iterators.size());
            {
                for (int i = 0; i < iterators.size(); i++) {
                    prebuiltHeads.add(new Head<K>(i));
                }
            }

            private Head<K> getHead(int position, K value) {
                Head<K> head = prebuiltHeads.get(position);
                head.value = value;
                return head;
            }

            {
                /*
                 *  Keep all the heads of all the iterators sorted.  
                 */
                for (int i = 0; i < iterators.size(); i++) {
                    PeekingSkippableIterator<K> it = iterators.get(i);
                    if (it.hasNext()) {
                        heads.add(getHead(i, it.next()));
                    }
                }

            }

            private List<K> ks = Lists.newArrayList();
            private BitSet nextsToDo = new BitSet();

            @Override
            protected V computeNext() {
                while (true) {
                    int pos = 0;
                    while (true) {
                        pos = nextsToDo.nextSetBit(pos);
                        if (pos < 0) {
                            break;
                        }
                        if (iterators.get(pos).hasNext()) {
                            heads.add(getHead(pos, iterators.get(pos).next()));
                        }
                        nextsToDo.clear(pos);
                        pos++;
                    }

                    Head<K> currentElement = heads.pollFirst();
                    if (currentElement == null) {
                        return endOfData();
                    }
                    int position = currentElement.position;
                    nextsToDo.set(position);

                    // advance the iterator from which we took the last element
                    PeekingSkippableIterator<K> peekingSkippableIterator = iterators.get(position);

                    ks.clear();

                    K currentK = currentElement.value;
                    ks.add(currentK);

                    // add all other heads while the key is the same
                    while (heads.size() > 0 && comp(heads.first().value, currentK) == 0) {
                        Head<K> newElement = heads.pollFirst();
                        nextsToDo.set(newElement.position);

                        PeekingSkippableIterator<K> newElementIterator = iterators.get(newElement.position);
                        ks.add(newElement.value);
                    }

                    V v = transform(ks.get(0));
                    if (shouldUse(v, ks)) {
                        return v;
                    }
                }
            }

            @Override
            public void skipTo(int i) {
                for (SkippableIterator<K> it : iterators) {
                    it.skipTo(i);
                }
            }
        };
    }

    public static <T> Function<PeekingIterator<T>, T> peekFunction() {
        return new Function<PeekingIterator<T>, T>() {
            @Override
            public T apply(PeekingIterator<T> it) {
                return it.peek();
            }
        };
    }

    public static <T> Function<SkippableIterable<T>, PeekingSkippableIterator<T>> peekingIteratorFunction() {
        return new Function<SkippableIterable<T>, PeekingSkippableIterator<T>>() {
            @Override
            public PeekingSkippableIterator<T> apply(SkippableIterable<T> ts) {
                return new PeekingSkippableIterator<T>(ts.iterator());
            }
        };
    }

    @SuppressWarnings("unchecked")
    public static void main(String[] args) {
        List<Integer> l1 = Lists.newArrayList(1, 2, 3, 3, 4, 5, 6, 7, 8, 9, 10);
        List<Integer> l2 = Lists.newArrayList(1, 2, 3, 4, 5, 6, 8, 9, 10);
        IdentityUnion<Integer> union = new IdentityUnion<Integer>(Skippables.fromIterable(l1),
                Skippables.fromIterable(l2));
        System.out.println(union);
    }

}