org.apache.cassandra.utils.MergeIterator Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of cassandra-all Show documentation
Show all versions of cassandra-all Show documentation
A fork of the Apache Cassandra Project that uses Lucene indexes for providing near real time search such as ElasticSearch or Solr, including full text search capabilities, multi-dimensional queries, and relevance scoring.
The newest version!
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 org.apache.cassandra.utils;
import java.io.Closeable;
import java.io.IOException;
import java.util.*;
import com.google.common.collect.AbstractIterator;
/** Merges sorted input iterators which individually contain unique items. */
public abstract class MergeIterator extends AbstractIterator implements IMergeIterator
{
protected final Reducer reducer;
protected final List> iterators;
protected MergeIterator(List> iters, Reducer reducer)
{
this.iterators = iters;
this.reducer = reducer;
}
public static IMergeIterator get(List> sources,
Comparator comparator,
Reducer reducer)
{
if (sources.size() == 1)
{
return reducer.trivialReduceIsTrivial()
? new TrivialOneToOne<>(sources, reducer)
: new OneToOne<>(sources, reducer);
}
return new ManyToOne<>(sources, comparator, reducer);
}
public Iterable> iterators()
{
return iterators;
}
public void close()
{
for (Iterator iterator : this.iterators)
{
try
{
((Closeable)iterator).close();
}
catch (IOException e)
{
throw new RuntimeException(e);
}
}
reducer.close();
}
/** A MergeIterator that consumes multiple input values per output value. */
private static final class ManyToOne extends MergeIterator
{
// a queue for return: all candidates must be open and have at least one item
protected final PriorityQueue> queue;
// a stack of the last consumed candidates, so that we can lazily call 'advance()'
// TODO: if we had our own PriorityQueue implementation we could stash items
// at the end of its array, so we wouldn't need this storage
protected final ArrayDeque> candidates;
public ManyToOne(List> iters, Comparator comp, Reducer reducer)
{
super(iters, reducer);
this.queue = new PriorityQueue<>(Math.max(1, iters.size()));
for (Iterator iter : iters)
{
Candidate candidate = new Candidate<>(iter, comp);
if (!candidate.advance())
// was empty
continue;
this.queue.add(candidate);
}
this.candidates = new ArrayDeque<>(queue.size());
}
protected final Out computeNext()
{
advance();
return consume();
}
/** Consume values by sending them to the reducer while they are equal. */
protected final Out consume()
{
reducer.onKeyChange();
Candidate candidate = queue.peek();
if (candidate == null)
return endOfData();
do
{
candidate = queue.poll();
candidates.push(candidate);
reducer.reduce(candidate.item);
}
while (queue.peek() != null && queue.peek().compareTo(candidate) == 0);
return reducer.getReduced();
}
/** Advance and re-enqueue all items we consumed in the last iteration. */
protected final void advance()
{
Candidate candidate;
while ((candidate = candidates.pollFirst()) != null)
if (candidate.advance())
queue.add(candidate);
}
}
// Holds and is comparable by the head item of an iterator it owns
protected static final class Candidate implements Comparable>
{
private final Iterator iter;
private final Comparator comp;
private In item;
public Candidate(Iterator iter, Comparator comp)
{
this.iter = iter;
this.comp = comp;
}
/** @return True if our iterator had an item, and it is now available */
protected boolean advance()
{
if (!iter.hasNext())
return false;
item = iter.next();
return true;
}
public int compareTo(Candidate that)
{
return comp.compare(this.item, that.item);
}
}
/** Accumulator that collects values of type A, and outputs a value of type B. */
public static abstract class Reducer
{
/**
* @return true if Out is the same as In for the case of a single source iterator
*/
public boolean trivialReduceIsTrivial()
{
return false;
}
/**
* combine this object with the previous ones.
* intermediate state is up to your implementation.
*/
public abstract void reduce(In current);
/** @return The last object computed by reduce */
protected abstract Out getReduced();
/**
* Called at the beginning of each new key, before any reduce is called.
* To be overridden by implementing classes.
*/
protected void onKeyChange() {}
/**
* May be overridden by implementations that require cleaning up after use
*/
public void close() {}
}
private static class OneToOne extends MergeIterator
{
private final Iterator source;
public OneToOne(List> sources, Reducer reducer)
{
super(sources, reducer);
source = sources.get(0);
}
protected Out computeNext()
{
if (!source.hasNext())
return endOfData();
reducer.onKeyChange();
reducer.reduce(source.next());
return reducer.getReduced();
}
}
private static class TrivialOneToOne extends MergeIterator
{
private final Iterator source;
public TrivialOneToOne(List> sources, Reducer reducer)
{
super(sources, reducer);
source = sources.get(0);
}
@SuppressWarnings("unchecked")
protected Out computeNext()
{
if (!source.hasNext())
return endOfData();
return (Out) source.next();
}
}
}