com.datastax.oss.driver.internal.core.cql.PagingIterableSpliterator Maven / Gradle / Ivy
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 com.datastax.oss.driver.internal.core.cql;
import com.datastax.oss.driver.api.core.PagingIterable;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.util.Objects;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.function.Consumer;
import net.jcip.annotations.NotThreadSafe;
/**
* A Spliterator for {@link PagingIterable} instances that splits the stream in chunks of equal
* size.
*
* @param The element type of the underlying stream.
*/
@NotThreadSafe
public class PagingIterableSpliterator implements Spliterator {
@NonNull
public static Builder builder(@NonNull PagingIterable iterable) {
return new Builder<>(iterable);
}
/** The default chunk size for {@link PagingIterableSpliterator}. */
public static final int DEFAULT_CHUNK_SIZE = 128;
private final PagingIterable iterable;
private long estimatedSize;
private final int chunkSize;
private final int characteristics;
/**
* Creates a new {@link PagingIterableSpliterator} for the given iterable, with unknown size and
* default chunk size ({@value #DEFAULT_CHUNK_SIZE}).
*
* @param iterable The {@link PagingIterable} to create a spliterator for.
*/
public PagingIterableSpliterator(@NonNull PagingIterable iterable) {
this(iterable, Long.MAX_VALUE, DEFAULT_CHUNK_SIZE);
}
private PagingIterableSpliterator(
@NonNull PagingIterable iterable, long estimatedSize, int chunkSize) {
this.iterable = Objects.requireNonNull(iterable, "iterable cannot be null");
this.estimatedSize = estimatedSize;
this.chunkSize = chunkSize;
if (estimatedSize < Long.MAX_VALUE) {
characteristics =
Spliterator.ORDERED
| Spliterator.IMMUTABLE
| Spliterator.NONNULL
| Spliterator.SIZED
| Spliterator.SUBSIZED;
} else {
characteristics = Spliterator.ORDERED | Spliterator.IMMUTABLE | Spliterator.NONNULL;
}
}
@Override
public boolean tryAdvance(Consumer action) {
Objects.requireNonNull(action, "action cannot be null");
ElementT row = iterable.one();
if (row == null) {
return false;
}
action.accept(row);
return true;
}
@Override
@Nullable
public Spliterator trySplit() {
if (estimatedSize != Long.MAX_VALUE && estimatedSize <= chunkSize) {
// There is no point in splitting if the number of remaining elements is below the chunk size
return null;
}
ElementT row = iterable.one();
if (row == null) {
return null;
}
Object[] array = new Object[chunkSize];
int i = 0;
do {
array[i++] = row;
if (i < chunkSize) {
row = iterable.one();
} else {
break;
}
} while (row != null);
if (estimatedSize != Long.MAX_VALUE) {
estimatedSize -= i;
}
// Splits will also report SIZED and SUBSIZED as well.
return Spliterators.spliterator(array, 0, i, characteristics());
}
@Override
public void forEachRemaining(Consumer action) {
iterable.iterator().forEachRemaining(action);
}
@Override
public long estimateSize() {
return estimatedSize;
}
@Override
public int characteristics() {
return characteristics;
}
public static class Builder {
private final PagingIterable iterable;
private long estimatedSize = Long.MAX_VALUE;
private int chunkSize = DEFAULT_CHUNK_SIZE;
Builder(@NonNull PagingIterable iterable) {
this.iterable = iterable;
}
@NonNull
public Builder withEstimatedSize(long estimatedSize) {
Preconditions.checkArgument(estimatedSize >= 0, "estimatedSize must be >= 0");
this.estimatedSize = estimatedSize;
return this;
}
@NonNull
public Builder withChunkSize(int chunkSize) {
Preconditions.checkArgument(chunkSize > 0, "chunkSize must be > 0");
this.chunkSize = chunkSize;
return this;
}
@NonNull
public PagingIterableSpliterator build() {
return new PagingIterableSpliterator<>(iterable, estimatedSize, chunkSize);
}
}
}