ThrottledUnfilteredIterator.java
/*
* 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.db.rows;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import com.google.common.annotations.VisibleForTesting;
import org.apache.cassandra.db.DeletionTime;
import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
import org.apache.cassandra.utils.AbstractIterator;
import org.apache.cassandra.utils.CloseableIterator;
/**
* A utility class to split the given {@link#UnfilteredRowIterator} into smaller chunks each
* having at most {@link #throttle} + 1 unfiltereds.
*
* Only the first output contains partition level info: {@link UnfilteredRowIterator#partitionLevelDeletion}
* and {@link UnfilteredRowIterator#staticRow}.
*
* Besides splitting, this iterator will also ensure each chunk does not finish with an open tombstone marker,
* by closing any opened tombstone markers and re-opening on the next chunk.
*
* The lifecycle of outputed {{@link UnfilteredRowIterator} only last till next call to {@link #next()}.
*
* A subsequent {@link #next} call will exhaust the previously returned iterator before computing the next,
* effectively skipping unfiltereds up to the throttle size.
*
* Closing this iterator will close the underlying iterator.
*
*/
public class ThrottledUnfilteredIterator extends AbstractIterator<UnfilteredRowIterator> implements CloseableIterator<UnfilteredRowIterator>
{
private final UnfilteredRowIterator origin;
private final int throttle;
// internal mutable state
private UnfilteredRowIterator throttledItr;
// extra unfiltereds from previous iteration
private Iterator<Unfiltered> overflowed = Collections.emptyIterator();
@VisibleForTesting
ThrottledUnfilteredIterator(UnfilteredRowIterator origin, int throttle)
{
assert origin != null;
assert throttle > 1 : "Throttle size must be higher than 1 to properly support open and close tombstone boundaries.";
this.origin = origin;
this.throttle = throttle;
this.throttledItr = null;
}
@Override
protected UnfilteredRowIterator computeNext()
{
// exhaust previous throttled iterator
while (throttledItr != null && throttledItr.hasNext())
throttledItr.next();
// The original UnfilteredRowIterator may have only partition deletion or static column but without unfiltereds.
// Return the original UnfilteredRowIterator
if (!origin.hasNext())
{
if (throttledItr != null)
return endOfData();
return throttledItr = origin;
}
throttledItr = new WrappingUnfilteredRowIterator()
{
private int count = 0;
private boolean isFirst = throttledItr == null;
// current batch's openMarker. if it's generated in previous batch,
// it must be consumed as first element of current batch
private RangeTombstoneMarker openMarker;
// current batch's closeMarker.
// it must be consumed as last element of current batch
private RangeTombstoneMarker closeMarker = null;
@Override
public UnfilteredRowIterator wrapped()
{
return origin;
}
@Override
public boolean hasNext()
{
return (withinLimit() && origin.hasNext()) || closeMarker != null;
}
@Override
public Unfiltered next()
{
if (closeMarker != null)
{
assert count == throttle;
Unfiltered toReturn = closeMarker;
closeMarker = null;
return toReturn;
}
Unfiltered next;
assert withinLimit();
// in the beginning of the batch, there might be remaining unfiltereds from previous iteration
if (overflowed.hasNext())
next = overflowed.next();
else
next = origin.next();
recordNext(next);
return next;
}
private void recordNext(Unfiltered unfiltered)
{
count++;
if (unfiltered.isRangeTombstoneMarker())
updateMarker((RangeTombstoneMarker) unfiltered);
// when reach throttle with a remaining openMarker, we need to create corresponding closeMarker.
if (count == throttle && openMarker != null)
{
assert origin.hasNext();
closeOpenMarker(origin.next());
}
}
private boolean withinLimit()
{
return count < throttle;
}
private void updateMarker(RangeTombstoneMarker marker)
{
openMarker = marker.isOpen(isReverseOrder()) ? marker : null;
}
/**
* There 3 cases for next, 1. if it's boundaryMarker, we split it as closeMarker for current batch, next
* openMarker for next batch 2. if it's boundMakrer, it must be closeMarker. 3. if it's Row, create
* corresponding closeMarker for current batch, and create next openMarker for next batch including current
* Row.
*/
private void closeOpenMarker(Unfiltered next)
{
assert openMarker != null;
if (next.isRangeTombstoneMarker())
{
RangeTombstoneMarker marker = (RangeTombstoneMarker) next;
// if it's boundary, create closeMarker for current batch and openMarker for next batch
if (marker.isBoundary())
{
RangeTombstoneBoundaryMarker boundary = (RangeTombstoneBoundaryMarker) marker;
closeMarker = boundary.createCorrespondingCloseMarker(isReverseOrder());
overflowed = Collections.singleton((Unfiltered)boundary.createCorrespondingOpenMarker(isReverseOrder())).iterator();
}
else
{
// if it's bound, it must be closeMarker.
assert marker.isClose(isReverseOrder());
updateMarker(marker);
closeMarker = marker;
}
}
else
{
// it's Row, need to create closeMarker for current batch and openMarker for next batch
DeletionTime openDeletion = openMarker.openDeletionTime(isReverseOrder());
closeMarker = RangeTombstoneBoundMarker.exclusiveClose(isReverseOrder(), next.clustering(), openDeletion);
// for next batch
overflowed = Arrays.asList(RangeTombstoneBoundMarker.inclusiveOpen(isReverseOrder(),
next.clustering(),
openDeletion), next).iterator();
}
}
@Override
public DeletionTime partitionLevelDeletion()
{
return isFirst ? origin.partitionLevelDeletion() : DeletionTime.LIVE;
}
@Override
public Row staticRow()
{
return isFirst ? origin.staticRow() : Rows.EMPTY_STATIC_ROW;
}
@Override
public void close()
{
// no op
}
};
return throttledItr;
}
public void close()
{
if (origin != null)
origin.close();
}
/**
* Splits a {@link UnfilteredPartitionIterator} in {@link UnfilteredRowIterator} batches with size no higher than
* <b>maxBatchSize</b>
*
* @param partitionIterator
* @param maxBatchSize max number of unfiltereds in the UnfilteredRowIterator. if 0 is given, it means no throttle.
* @return
*/
public static CloseableIterator<UnfilteredRowIterator> throttle(UnfilteredPartitionIterator partitionIterator, int maxBatchSize)
{
if (maxBatchSize == 0) // opt out
return partitionIterator;
return new AbstractIterator<UnfilteredRowIterator>()
{
ThrottledUnfilteredIterator current = null;
protected UnfilteredRowIterator computeNext()
{
if (current != null && !current.hasNext())
{
current.close();
current = null;
}
if (current == null && partitionIterator.hasNext())
{
current = new ThrottledUnfilteredIterator(partitionIterator.next(), maxBatchSize);
}
if (current != null && current.hasNext())
return current.next();
return endOfData();
}
public void close()
{
if (current != null)
current.close();
}
};
}
}