HBaseIterator.java

package eu.fbk.knowledgestore.datastore.hbase;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

import javax.annotation.Nullable;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;

import org.openrdf.model.URI;

import eu.fbk.knowledgestore.data.Record;
import eu.fbk.knowledgestore.datastore.hbase.utils.AbstractHBaseUtils;

/**
 * Performs batch lookup of records by ID from a given HBase table, optionally returning only a
 * subset of record properties.
 */
public class HBaseIterator extends AbstractIterator<Record> {

    /** The batch size. */
    private static final int BATCH_SIZE = 100;

    /** Object referencing transactional layer. */
    private final AbstractHBaseUtils hbaseUtils;

    /** HBase table name to be used */
    private final String tableName;

    /** Properties to be looked up. */
    private final URI[] properties;

    /** An iterator over the IDs of the records to lookup. */
    private final Iterator<URI> idIterator;

    /** An iterator over records buffered in the last batch looked up from HBase. */
    private Iterator<Record> recordIterator;

    /**
     * Creates a new {@code HBaseStream} for the parameters supplied.
     * 
     * @param hbaseUtils
     *            the {@code AbstractHBaseUtils} object for accessing HBase, not null
     * @param tableName
     *            the name of the HBase table to access, not null
     * @param ids
     *            the IDs of records to fetch from HBase, not null
     * @param properties
     *            the properties of records to return, null if all properties should be returned
     */
    @SuppressWarnings("unchecked")
    public HBaseIterator(final AbstractHBaseUtils hbaseUtils, final String tableName,
            final Set<? extends URI> ids, @Nullable final Set<? extends URI> properties) {

        Preconditions.checkNotNull(hbaseUtils);
        Preconditions.checkNotNull(tableName);
        Preconditions.checkNotNull(ids);

        this.hbaseUtils = hbaseUtils;
        this.tableName = tableName;
        this.properties = properties == null ? null : Iterables.toArray(properties, URI.class);
        this.idIterator = (Iterator<URI>) ImmutableList.copyOf(ids).iterator();
        this.recordIterator = Collections.emptyIterator();
    }

    @Override
    protected Record computeNext() {
        while (true) {
            // Return a record previously buffered, if available
            if (this.recordIterator.hasNext()) {
                return this.recordIterator.next();
            }

            // Otherwise, retrieve next batch of IDs
            final List<URI> ids = Lists.newArrayListWithCapacity(BATCH_SIZE);
            while (this.idIterator.hasNext() && ids.size() < BATCH_SIZE) {
                ids.add(this.idIterator.next());
            }

            // EOF reached if there are no more IDs to retrieve
            if (ids.isEmpty()) {
                return endOfData();
            }

            // Retrieve next batch of records corresponding to IDs batch
            final List<Record> records;
            try {
                records = this.hbaseUtils.get(this.tableName, ids);
            } catch (final IOException ex) {
                throw Throwables.propagate(ex);
            }

            // Perform client-side projection, if requested
            if (this.properties != null) {
                for (int i = 0; i < records.size(); ++i) {
                    records.set(i, records.get(i).retain(this.properties));
                }
            }

            // Store fetched record in record iterator and return first one
            this.recordIterator = records.iterator();
        }
    }

}