com.tinkerpop.gremlin.giraph.hdfs.GiraphElementIterator Maven / Gradle / Ivy

Go to download
package com.tinkerpop.gremlin.giraph.hdfs;

import com.tinkerpop.gremlin.giraph.Constants;
import com.tinkerpop.gremlin.giraph.process.computer.util.ConfUtil;
import com.tinkerpop.gremlin.giraph.structure.GiraphGraph;
import com.tinkerpop.gremlin.structure.Element;
import org.apache.giraph.io.VertexInputFormat;
import org.apache.giraph.io.VertexReader;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;

import java.io.IOException;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Queue;

/**
 * @author Marko A. Rodriguez (http://markorodriguez.com)
 */
public abstract class GiraphElementIterator implements Iterator {

    protected final GiraphGraph graph;
    protected final Queue readers = new LinkedList<>();

    public GiraphElementIterator(final GiraphGraph graph, final VertexInputFormat inputFormat, final Path path) throws IOException {
        this.graph = graph;
        final Configuration configuration = ConfUtil.makeHadoopConfiguration(this.graph.variables().getConfiguration());
        for (final FileStatus status : FileSystem.get(configuration).listStatus(path, HiddenFileFilter.instance())) {
            this.readers.add(inputFormat.createVertexReader(new FileSplit(status.getPath(), 0, Integer.MAX_VALUE, new String[]{}), new TaskAttemptContext(configuration, new TaskAttemptID())));
        }
    }

    public GiraphElementIterator(final GiraphGraph graph) throws IOException {
        try {
            this.graph = graph;
            if (this.graph.variables().getConfiguration().containsKey(Constants.GREMLIN_INPUT_LOCATION)) {
                final Configuration configuration = ConfUtil.makeHadoopConfiguration(this.graph.variables().getConfiguration());
                final VertexInputFormat inputFormat = this.graph.variables().getConfiguration().getInputFormat().getConstructor().newInstance();
                for (final FileStatus status : FileSystem.get(configuration).listStatus(new Path(graph.variables().getConfiguration().getInputLocation()), HiddenFileFilter.instance())) {
                    this.readers.add(inputFormat.createVertexReader(new FileSplit(status.getPath(), 0, Integer.MAX_VALUE, new String[]{}), new TaskAttemptContext(configuration, new TaskAttemptID())));
                }
            }
        } catch (Exception e) {
            throw new IllegalStateException(e.getMessage(), e);
        }
    }
}