com.palantir.common.base.BatchingVisitables Maven / Gradle / Ivy
Show all versions of atlasdb-commons Show documentation
/*
* (c) Copyright 2018 Palantir Technologies Inc. All rights reserved.
*
* 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.palantir.common.base;
import com.google.common.base.Function;
import com.google.common.base.Functions;
import com.google.common.base.Predicate;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.Ordering;
import com.palantir.common.annotation.Inclusive;
import com.palantir.logsafe.Preconditions;
import com.palantir.logsafe.SafeArg;
import com.palantir.util.Mutable;
import com.palantir.util.Mutables;
import com.palantir.util.paging.SimpleTokenBackedResultsPage;
import com.palantir.util.paging.TokenBackedBasicResultsPage;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import javax.annotation.Nullable;
public final class BatchingVisitables {
private BatchingVisitables() {
/**/
}
public static final int DEFAULT_BATCH_SIZE = 1000;
@SuppressWarnings("unused") // Used in internal legacy code
public static final int KEEP_ALL_BATCH_SIZE = 100000;
public static BatchingVisitableView emptyBatchingVisitable() {
return BatchingVisitableView.of(new BatchingVisitable() {
@Override
public boolean batchAccept(int batchSize, AbortingVisitor, K> v)
throws K {
return true;
}
});
}
public static BatchingVisitableView singleton(final T t) {
return BatchingVisitableView.of(new BatchingVisitable() {
@Override
public boolean batchAccept(int batchSize, AbortingVisitor, K> v)
throws K {
return v.visit(ImmutableList.of(t));
}
});
}
public static long count(BatchingVisitable visitable) {
return count(visitable, 1000);
}
public static long count(BatchingVisitable visitable, int batchSize) {
return countInternal(visitable, batchSize);
}
private static long countInternal(BatchingVisitable visitable, int batchSize) {
final long[] count = new long[1];
visitable.batchAccept(batchSize, AbortingVisitors.batching(item -> {
count[0]++;
return true;
}));
return count[0];
}
public static boolean isEmpty(BatchingVisitable v) {
return v.batchAccept(1, AbortingVisitors.batching(AbortingVisitors.alwaysFalse()));
}
/**
* Returns a visitable containing elements from the start of the visitable
* for which the predicate holds true. Once boundingPredicate returns false for an
* element of visitable, no more elements will be included in the returned visitable.
*/
public static BatchingVisitableView visitWhile(
final BatchingVisitable visitable, final Predicate condition) {
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, final ConsistentVisitor v)
throws K {
visitable.batchAccept(batchSizeHint, batch -> {
for (T t : batch) {
if (!condition.apply(t)) {
return false;
}
boolean keepGoing = v.visitOne(t);
if (!keepGoing) {
return false;
}
}
return true;
});
}
});
}
/**
* Returns the first element or null if the visitable is empty
*/
@Nullable
public static T getFirst(BatchingVisitable visitable) {
return getFirst(visitable, null);
}
@Nullable
public static T getFirst(BatchingVisitable visitable, @Nullable T defaultElement) {
final Mutable ret = Mutables.newMutable(defaultElement);
visitable.batchAccept(1, AbortingVisitors.batching(item -> {
ret.set(item);
return false;
}));
return ret.get();
}
public static boolean isEqual(BatchingVisitable v, final Iterable it) {
return isEqual(v, it.iterator());
}
public static boolean isEqual(BatchingVisitable v, final Iterator it) {
boolean ret = v.batchAccept(DEFAULT_BATCH_SIZE, batch -> {
Iterator toMatch = Iterators.limit(it, batch.size());
return Iterators.elementsEqual(toMatch, batch.iterator());
});
if (it.hasNext()) {
return false;
}
return ret;
}
/**
* This method will throw if there are nulls in the visitable.
*
* @return null if the visitable is empty, otherwise return the smallest
* value.
*/
public static > T getMax(BatchingVisitable v) {
return getMax(v, Ordering.natural(), null);
}
/**
* This will return the first maximal element in the visitable. This method
* takes a default element and will return that if the visitable is empty.
* If the visitable is non-empty it will return the largest value in the
* visitable.
*
* A common way to use this would be to pass null as the
* defaultElement and have the ordering throw on null elements so you know
* the visitable doesn't have any nulls.
*/
public static T getMax(BatchingVisitable v, final Ordering o, @Nullable T defaultElement) {
final Mutable ret = Mutables.newMutable(defaultElement);
v.batchAccept(
DEFAULT_BATCH_SIZE,
AbortingVisitors.batching(new AbortingVisitor() {
boolean hasSeenFirst = false;
@Override
public boolean visit(T item) throws RuntimeException {
if (hasSeenFirst) {
ret.set(o.max(ret.get(), item));
} else {
// Call o.max here so it will throw if item is null and this
// ordering hates on nulls.
ret.set(o.max(item, item));
hasSeenFirst = true;
}
return true;
}
}));
return ret.get();
}
/**
* This method will throw if there are nulls in the visitable.
*
* @return null if the visitable is empty, otherwise return the smallest value.
*/
public static > T getMin(BatchingVisitable v) {
return getMin(v, Ordering.natural(), null);
}
/**
* This will return the first maximal element in the visitable. This method
* takes a default element and will return that if the visitable is empty.
* If the visitable is non-empty it will return the largest value in the
* visitable.
*
* A common way to use this would be to pass null as the
* defaultElement and have the ordering throw on null elements so you know
* the visitable doesn't have any nulls.
*/
public static T getMin(BatchingVisitable v, final Ordering o, @Nullable T defaultElement) {
final Mutable ret = Mutables.newMutable(defaultElement);
v.batchAccept(
DEFAULT_BATCH_SIZE,
AbortingVisitors.batching(new AbortingVisitor() {
boolean hasSeenFirst = false;
@Override
public boolean visit(T item) throws RuntimeException {
if (hasSeenFirst) {
ret.set(o.min(ret.get(), item));
} else {
// Call o.max here so it will throw if item is null and this
// ordering hates on nulls.
ret.set(o.min(item, item));
hasSeenFirst = true;
}
return true;
}
}));
return ret.get();
}
@Nullable
public static T getLast(BatchingVisitable visitable) {
return getLast(visitable, null);
}
@Nullable
public static T getLast(BatchingVisitable visitable, @Nullable T defaultElement) {
final Mutable ret = Mutables.newMutable(defaultElement);
visitable.batchAccept(DEFAULT_BATCH_SIZE, AbortingVisitors.batching(item -> {
ret.set(item);
return true;
}));
return ret.get();
}
public static BatchingVisitableView filter(BatchingVisitable visitable, final Predicate pred) {
Preconditions.checkNotNull(pred);
return transformBatch(visitable, input -> ImmutableList.copyOf(Iterables.filter(input, pred)));
}
public static BatchingVisitableView transform(
BatchingVisitable visitable, final Function f) {
return transformBatch(visitable, from -> Lists.transform(from, f));
}
public static BatchingVisitableView transformBatch(
final BatchingVisitable visitable, final Function, ? extends List> f) {
Preconditions.checkNotNull(visitable);
Preconditions.checkNotNull(f);
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, final ConsistentVisitor v)
throws K {
visitable.batchAccept(batchSizeHint, batch -> v.visit(f.apply(batch)));
}
});
}
public static BatchingVisitableView limit(final BatchingVisitable visitable, final long limit) {
Preconditions.checkNotNull(visitable);
Preconditions.checkArgument(limit >= 0);
if (limit == 0) {
return emptyBatchingVisitable();
}
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, final ConsistentVisitor v)
throws K {
if (batchSizeHint > limit) {
batchSizeHint = (int) limit;
}
visitable.batchAccept(batchSizeHint, new AbortingVisitor, K>() {
long visited = 0;
@Override
public boolean visit(List batch) throws K {
for (T item : batch) {
if (!v.visitOne(item)) {
return false;
}
visited++;
if (visited >= limit) {
// Stop visiting early by returning false.
// The underlying ConsistentVisitor will still cause #batchAccept
// to return true.
return false;
}
}
return true;
}
});
}
});
}
public static BatchingVisitableView skip(final BatchingVisitable visitable, final long toSkip) {
Preconditions.checkNotNull(visitable);
Preconditions.checkArgument(toSkip >= 0);
if (toSkip == 0) {
return BatchingVisitableView.of(visitable);
}
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, final ConsistentVisitor v)
throws K {
visitable.batchAccept(batchSizeHint, new AbortingVisitor, K>() {
long visited = 0;
@Override
public boolean visit(List batch) throws K {
for (T item : batch) {
if (visited < toSkip) {
visited++;
continue;
}
if (!v.visitOne(item)) {
return false;
}
}
return true;
}
});
}
});
}
/**
* This is similar to the unix uniq command.
*
* If there is a sequence of identical values, it will be replaced by one value.
*
* For example "AAABABBB" will become "ABAB"
*
* If the list passed in is sorted, then it will actually result in a list of unique elements
*
* This uses {@link Objects#equal(Object, Object)} to do comparisons.
*
* null is supported bug discouraged
*/
public static BatchingVisitableView unique(final BatchingVisitable visitable) {
return uniqueOn(visitable, Functions.identity());
}
public static BatchingVisitableView uniqueOn(
final BatchingVisitable visitable, final Function function) {
Preconditions.checkNotNull(visitable);
Preconditions.checkNotNull(function);
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, final ConsistentVisitor v)
throws K {
visitable.batchAccept(batchSizeHint, new AbortingVisitor, K>() {
boolean hasVisitedFirst = false;
Object lastVisited = null;
@Override
public boolean visit(List batch) throws K {
for (T item : batch) {
Object itemKey = function.apply(item);
if (!hasVisitedFirst || !Objects.equals(itemKey, lastVisited)) {
if (!v.visitOne(item)) {
return false;
}
hasVisitedFirst = true;
lastVisited = itemKey;
}
}
return true;
}
});
}
});
}
public static List copyToList(BatchingVisitable v) {
return BatchingVisitableView.of(v).immutableCopy();
}
public static List take(BatchingVisitable v, int howMany) {
return take(v, howMany, true);
}
public static List take(BatchingVisitable v, final int howMany, final boolean includeFirst) {
BatchingVisitableView visitable = BatchingVisitableView.of(v);
if (!includeFirst) {
visitable = visitable.skip(1);
}
return visitable.limit(howMany).immutableCopy();
}
public static TokenBackedBasicResultsPage getFirstPage(
BatchingVisitable v, int numToVisitArg, Function tokenExtractor) {
Preconditions.checkArgument(
numToVisitArg >= 0, "numToVisit cannot be negative.", SafeArg.of("numToVisit", numToVisitArg));
if (numToVisitArg == Integer.MAX_VALUE) {
// prevent issue with overflow
numToVisitArg--;
}
final int numToVisit = numToVisitArg + 1;
ImmutableList list = BatchingVisitableView.of(v).limit(numToVisit).immutableCopy();
Preconditions.checkState(list.size() <= numToVisit);
if (list.size() >= numToVisit) {
TOKEN token = tokenExtractor.apply(list.get(list.size() - 1));
list = list.subList(0, numToVisit - 1);
return new SimpleTokenBackedResultsPage(token, list, true);
}
return new SimpleTokenBackedResultsPage(null, list, false);
}
public static TokenBackedBasicResultsPage getFirstPage(BatchingVisitable v, int numToVisitArg) {
return getFirstPage(v, numToVisitArg, Functions.identity());
}
/**
* This method will wrap the passed visitable so it is called with the passed pageSize when it is called.
*
* This can be used to make the performance of batching visitables better. One example of where this is useful
* is if I just visit the results one at a time, but I know that I will visit 100 results, then I can save
* a lot of potential round trips by hinting a page size of 100.
*/
public static BatchingVisitableView hintPageSize(final BatchingVisitable bv, final int pageSize) {
Preconditions.checkArgument(pageSize > 0);
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, ConsistentVisitor v)
throws K {
bv.batchAccept(pageSize, v);
}
});
}
/**
* This will wrap the passed visitor and return a more generic type. Since Visitables just produce values,
* their type can be made more generic because a Visitable<Long> can safely be cast to a Visitable<Object>.
*
* @see BatchingVisitableView for more helper methods
*/
public static BatchingVisitable wrap(final BatchingVisitable visitable) {
if (visitable == null) {
return null;
}
return new BatchingVisitable() {
@Override
public boolean batchAccept(int batchSize, AbortingVisitor, K> v)
throws K {
// This cast is safe because an abortingVisitor can consume more specific values
@SuppressWarnings("unchecked")
AbortingVisitor, K> specificVisitor = (AbortingVisitor, K>) v;
return visitable.batchAccept(batchSize, specificVisitor);
}
};
}
public static BatchingVisitableView concat(BatchingVisitable... inputs) {
return concat(ImmutableList.copyOf(inputs));
}
public static BatchingVisitableView concat(final Iterable> inputs) {
Preconditions.checkNotNull(inputs);
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, ConsistentVisitor v)
throws K {
for (BatchingVisitable bv : inputs) {
// This is safe because cast is never passed to anything and it's function
// batchAccept is covariant
@SuppressWarnings("unchecked")
BatchingVisitable cast = (BatchingVisitable) bv;
if (!cast.batchAccept(batchSizeHint, v)) {
return;
}
}
}
});
}
public static BatchingVisitableView flatten(
final int outerBatchHint, final BatchingVisitable> inputs) {
Preconditions.checkNotNull(inputs);
return BatchingVisitableView.of(new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(
final int batchSizeHint, final ConsistentVisitor v) throws K {
inputs.batchAccept(
outerBatchHint, (AbortingVisitor>, K>) bvs -> {
for (BatchingVisitable bv : bvs) {
// This is safe because cast is never passed to anything and it's function
// batchAccept is covariant
@SuppressWarnings("unchecked")
BatchingVisitable cast = (BatchingVisitable) bv;
if (!cast.batchAccept(batchSizeHint, v)) {
return false;
}
}
return true;
});
}
});
}
public static BatchingVisitable getOrderedVisitableUsingSublists(
final OrderedSublistProvider sublistProvider, @Inclusive final long startId) {
return BatchingVisitables.getOrderedVisitableUsingSublists(
sublistProvider, startId, Functions.identity());
}
public static BatchingVisitable getOrderedVisitableUsingSublists(
final OrderedSublistProvider sublistProvider,
@Inclusive final long startId,
final Function idFunction) {
BatchingVisitable visitableWithDuplicates = new AbstractBatchingVisitable() {
@Override
protected void batchAcceptSizeHint(int batchSizeHint, ConsistentVisitor v)
throws K {
long nextId = startId;
boolean abortedOrFinished = false;
while (!abortedOrFinished) {
List resultBatchWithDuplicates = sublistProvider.getBatchAllowDuplicates(nextId, batchSizeHint);
boolean abortedByVisitor = !v.visit(resultBatchWithDuplicates);
nextId = nextIdStart(resultBatchWithDuplicates, batchSizeHint, idFunction);
boolean finishedAllBatches = nextId == -1L;
abortedOrFinished = abortedByVisitor || finishedAllBatches;
}
}
};
return BatchingVisitables.unique(visitableWithDuplicates);
}
private static long nextIdStart(List recentlyVisited, long batchSizeHint, Function idFunction) {
if (recentlyVisited.size() < batchSizeHint) {
return -1L;
}
long cur = idFunction.apply(recentlyVisited.get(recentlyVisited.size() - 1));
if (cur == Long.MAX_VALUE) {
return -1L;
}
return cur + 1L;
}
}