com.basho.riak.client.query.MapReduce.java Source code

Java tutorial

Introduction

Here is the source code for com.basho.riak.client.query.MapReduce.java

Source

/*
 * This file is provided 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.basho.riak.client.query;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Collection;
import java.util.LinkedList;

import com.basho.riak.client.IRiakClient;
import com.basho.riak.client.RiakException;
import com.basho.riak.client.operations.RiakOperation;
import com.basho.riak.client.query.functions.Function;
import com.basho.riak.client.query.serialize.FunctionToJson;
import com.basho.riak.client.raw.RawClient;
import com.basho.riak.client.raw.query.MapReduceSpec;
import com.fasterxml.jackson.core.JsonEncoding;
import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * An operation for defining and runnig a Map/Reduce query on Riak.
 * 
 * <p>
 * See <a href="http://wiki.basho.com/MapReduce.html">Map/Reduce</a> for details.
 * </p>
 * @author russell
 * 
 * @see IRiakClient#mapReduce()
 * @see IRiakClient#mapReduce(String)
 */
public abstract class MapReduce implements RiakOperation<MapReduceResult> {

    private final RawClient client;

    private Collection<MapReducePhase> phases = new LinkedList<MapReducePhase>();
    private Long timeout;

    /**
     * Create the MapRedcue operation with the {@link RawClient} to delegate to.
     * 
     * @param client
     *            a {@link RawClient}
     * 
     * @see IRiakClient#mapReduce()
     * @see IRiakClient#mapReduce(String)
     */
    public MapReduce(RawClient client) {
        this.client = client;
    }

    /**
     * Run the Map/Reduce job against the {@link RawClient} the operation was
     * constructed with.
     * 
     * @return a {@link MapReduceResult} containing the results of the query.
     * @throws RiakException
     * @throws InvalidMapReduceException
     */
    public MapReduceResult execute() throws RiakException {
        validate();
        final String strSpec = writeSpec();
        MapReduceSpec spec = new MapReduceSpec(strSpec);
        try {
            return client.mapReduce(spec);
        } catch (IOException e) {
            throw new RiakException(e);
        }
    }

    /**
     * Check that this map/reduce job is valid
     * @throws InvalidMapReduceException
     */
    protected void validate() {
        // Changed to no-op; phase no longer required
        // Left in place so as not to break inheritence
    }

    /**
     * Creates the JSON string of the M/R job for submitting to the
     * {@link RawClient}
     * 
     * Uses Jackson to write out the JSON string. I'm not very happy with this
     * method, it is a candidate for change.
     * 
     * TODO re-evaluate this method, look for something smaller and more elegant.
     * 
     * @return a String of JSON
     * @throws RiakException
     *             if, for some reason, we can't create a JSON string.
     */
    private String writeSpec() throws RiakException {

        final ByteArrayOutputStream out = new ByteArrayOutputStream();

        try {
            JsonGenerator jg = new JsonFactory().createJsonGenerator(out, JsonEncoding.UTF8);
            jg.setCodec(new ObjectMapper());

            jg.writeStartObject();

            jg.writeFieldName("inputs");
            writeInput(jg);

            jg.writeFieldName("query");
            jg.writeStartArray();

            writePhases(jg);

            jg.writeEndArray();
            if (timeout != null) {
                jg.writeNumberField("timeout", timeout);
            }

            jg.writeEndObject();
            jg.flush();

            return out.toString("UTF8");
        } catch (IOException e) {
            throw new RiakException(e);
        }
    }

    /**
     * Write the collection of phases to the json output generator
     * @param jg
     *            a {@link JsonGenerator}
     */
    private void writePhases(JsonGenerator jg) throws IOException {
        writeMapReducePhases(jg);
    }

    /**
     * @param jg
     */
    private void writeMapReducePhases(JsonGenerator jg) throws IOException {
        int cnt = 0;
        synchronized (phases) {
            final int lastPhase = phases.size();
            for (MapReducePhase phase : phases) {
                cnt++;
                jg.writeStartObject();
                jg.writeFieldName(phase.getType().toString());
                jg.writeStartObject();

                switch (phase.getType()) {
                case MAP:
                case REDUCE:
                    MapPhase mapPhase = (MapPhase) phase;
                    FunctionToJson.newWriter(mapPhase.getPhaseFunction(), jg).write();
                    if (mapPhase.getArg() != null) {
                        jg.writeObjectField("arg", mapPhase.getArg());
                    }
                    break;
                case LINK:
                    jg.writeStringField("bucket", ((LinkPhase) phase).getBucket());
                    jg.writeStringField("tag", ((LinkPhase) phase).getTag());
                    break;
                }

                //the final phase results should be returned, unless specifically set otherwise
                if (cnt == lastPhase) {
                    jg.writeBooleanField("keep", isKeepResult(true, phase.isKeep()));
                } else {
                    jg.writeBooleanField("keep", isKeepResult(false, phase.isKeep()));
                }

                jg.writeEndObject();
                jg.writeEndObject();
            }
        }
    }

    /**
     * Decide if a map/reduce phase result should be kept (returned) or not.
     *
     * @param isLastPhase
     *            is the phase being considered the last phase in an m/r job?
     * @param phaseKeepValue
     *            the Boolean value from a {@link MapPhase} (null|true|false)
     * @return <code>phaseKeepValue</code> if not null, otherwise
     *         <code>true</code> if <code>isLastPhase</code> is true, false
     *         otherwise.
     */
    private boolean isKeepResult(boolean isLastPhase, Boolean phaseKeepValue) {
        if (phaseKeepValue != null) {
            return phaseKeepValue;
        } else {
            return isLastPhase;
        }
    }

    /**
     * Set the operations timeout
     * @param timeout
     * @return this
     */
    public MapReduce timeout(long timeout) {
        this.timeout = timeout;
        return this;
    }

    /**
     * Add {@link MapPhase} to the query
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @param keep
     *            keep the results and return them with the query results?
     * @return this
     */
    public MapReduce addMapPhase(Function phaseFunction, boolean keep) {
        synchronized (phases) {
            phases.add(new MapPhase(phaseFunction, keep));
        }

        return this;
    }

    /**
     * Add a MapPhase
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @param arg
     *            an argument that will be passed to the phase verbatim
     *            (Object#toString)
     * @param keep
     *            if the result should be returned or merely provide input for
     *            the next phase.
     * @return this
     */
    public MapReduce addMapPhase(Function phaseFunction, Object arg, boolean keep) {
        synchronized (phases) {
            phases.add(new MapPhase(phaseFunction, arg, keep));
        }

        return this;
    }

    /**
     * Add a MapPhase
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @param arg
     *            an argument that will be passed to the phase verbatim
     *            (Object#toString)
     * @return this
     */
    public MapReduce addMapPhase(Function phaseFunction, Object arg) {
        synchronized (phases) {
            phases.add(new MapPhase(phaseFunction, arg));
        }

        return this;
    }

    /**
     * Add a MapPhase
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @return this
     */
    public MapReduce addMapPhase(Function phaseFunction) {
        synchronized (phases) {
            phases.add(new MapPhase(phaseFunction));
        }

        return this;
    }

    /**
     * Add {@link ReducePhase} to the query
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @param keep
     *            keep the results and return them with the query results?
     * @return this
     */
    public MapReduce addReducePhase(Function phaseFunction, boolean keep) {
        synchronized (phases) {
            phases.add(new ReducePhase(phaseFunction, keep));
        }

        return this;
    }

    /**
     * Add a {@link ReducePhase}
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @param arg
     *            an argument that will be passed to the phase verbatim
     *            (Object#toString)
     * @param keep
     *            if the result should be returned or merely provide input for
     *            the next phase.
     * @return this
     */
    public MapReduce addReducePhase(Function phaseFunction, Object arg, boolean keep) {
        synchronized (phases) {
            phases.add(new ReducePhase(phaseFunction, arg, keep));
        }

        return this;
    }

    /**
     * Add a {@link ReducePhase}
     * 
     * @param phaseFunction
     *            the {@link Function}
     * @param arg
     *            an argument that will be passed to the phase verbatim
     * @return this
     */
    public MapReduce addReducePhase(Function phaseFunction, Object arg) {
        synchronized (phases) {
            phases.add(new ReducePhase(phaseFunction, arg));
        }

        return this;
    }

    /**
     * Add a {@link ReducePhase}
     * 
     * @param phaseFunction
     * @return this
     */
    public MapReduce addReducePhase(Function phaseFunction) {
        synchronized (phases) {
            phases.add(new ReducePhase(phaseFunction));
        }

        return this;
    }

    /**
     * Add a Link Phase that points to <code>bucket</code> / <code>tag</code>
     * .
     * 
     * @param bucket
     *            the bucket at the end of the link (or "_" or "" for wildcard)
     * @param tag
     *            the tag (or ("_", or "" for wildcard)
     * @param keep
     *            to keep the result of this phase and return it at the end of
     *            the operation
     */
    public MapReduce addLinkPhase(String bucket, String tag, boolean keep) {
        synchronized (phases) {
            phases.add(new LinkPhase(bucket, tag, keep));
        }

        return this;
    }

    /**
     * Create a Link Phase that points to <code>bucket</code> / <code>tag</code>
     * <code>keep</code> will be <code>false</code>
     * 
     * @param bucket
     *            the bucket at the end of the link (or "_" or "" for wildcard)
     * @param tag
     *            the tag (or ("_", or "" for wildcard)
     */
    public MapReduce addLinkPhase(String bucket, String tag) {
        synchronized (phases) {
            phases.add(new LinkPhase(bucket, tag));
        }

        return this;
    }

    /**
     * Override to write the input specification of the M/R job.
     * 
     * @param jsonGenerator a Jackson {@link JsonGenerator} to write the input spec to
     * @throws IOException
     */
    protected abstract void writeInput(JsonGenerator jsonGenerator) throws IOException;
}