com.datastax.oss.driver.internal.core.cql.PagingIterableSpliterator Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of java-driver-core
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);
    }
  }
}