/*
* 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<In,Out> extends AbstractIterator<Out> implements IMergeIterator<In, Out>
{
protected final Reducer<In,Out> reducer;
protected final List<? extends Iterator<In>> iterators;
protected MergeIterator(List<? extends Iterator<In>> iters, Reducer<In, Out> reducer)
{
this.iterators = iters;
this.reducer = reducer;
}
public static <In, Out> IMergeIterator<In, Out> get(List<? extends Iterator<In>> sources,
Comparator<In> comparator,
Reducer<In, Out> reducer)
{
if (sources.size() == 1)
{
return reducer.trivialReduceIsTrivial()
? new TrivialOneToOne<>(sources, reducer)
: new OneToOne<>(sources, reducer);
}
return new ManyToOne<>(sources, comparator, reducer);
}
public Iterable<? extends Iterator<In>> iterators()
{
return iterators;
}
public void close()
{
for (Iterator<In> 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<In,Out> extends MergeIterator<In,Out>
{
// a queue for return: all candidates must be open and have at least one item
protected final PriorityQueue<Candidate<In>> 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<Candidate<In>> candidates;
public ManyToOne(List<? extends Iterator<In>> iters, Comparator<In> comp, Reducer<In, Out> reducer)
{
super(iters, reducer);
this.queue = new PriorityQueue<>(Math.max(1, iters.size()));
for (Iterator<In> iter : iters)
{
Candidate<In> 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<In> 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<In> 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<In> implements Comparable<Candidate<In>>
{
private final Iterator<In> iter;
private final Comparator<In> comp;
private In item;
public Candidate(Iterator<In> iter, Comparator<In> 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<In> 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<In,Out>
{
/**
* @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<In, Out> extends MergeIterator<In, Out>
{
private final Iterator<In> source;
public OneToOne(List<? extends Iterator<In>> sources, Reducer<In, Out> 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<In, Out> extends MergeIterator<In, Out>
{
private final Iterator<In> source;
public TrivialOneToOne(List<? extends Iterator<In>> sources, Reducer<In, Out> reducer)
{
super(sources, reducer);
source = sources.get(0);
}
@SuppressWarnings("unchecked")
protected Out computeNext()
{
if (!source.hasNext())
return endOfData();
return (Out) source.next();
}
}
}