package com.thinkaurelius.faunus.mapreduce.sideeffect;
import com.thinkaurelius.faunus.FaunusVertex;
import com.thinkaurelius.faunus.Tokens;
import com.thinkaurelius.faunus.mapreduce.util.EmptyConfiguration;
import com.thinkaurelius.faunus.mapreduce.util.SafeMapperOutputs;
import com.thinkaurelius.faunus.tinkerpop.gremlin.FaunusGremlinScriptEngine;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Mapper;
import javax.script.ScriptEngine;
import java.io.IOException;
import java.io.InputStreamReader;
/**
* @author Marko A. Rodriguez (http://markorodriguez.com)
*/
public class ScriptMap {
public static final String CLASS = Tokens.makeNamespace(ScriptMap.class) + ".class";
public static final String SCRIPT_PATH = Tokens.makeNamespace(ScriptMap.class) + ".scriptPath";
public static final String SCRIPT_ARGS = Tokens.makeNamespace(ScriptMap.class) + ".scriptArgs";
private static final String ARGS = "args";
private static final String V = "v";
private static final String SETUP_ARGS = "setup(args)";
private static final String MAP_V_ARGS = "map(v,args)";
private static final String CLEANUP_ARGS = "cleanup(args)";
public static Configuration createConfiguration(final String scriptUri, final String... args) {
Configuration configuration = new EmptyConfiguration();
configuration.set(SCRIPT_PATH, scriptUri);
configuration.setStrings(SCRIPT_ARGS, args);
return configuration;
}
public static class Map extends Mapper<NullWritable, FaunusVertex, NullWritable, FaunusVertex> {
private final ScriptEngine engine = new FaunusGremlinScriptEngine();
private SafeMapperOutputs outputs;
private Text textWritable = new Text();
@Override
public void setup(final Mapper.Context context) throws IOException, InterruptedException {
final FileSystem fs = FileSystem.get(context.getConfiguration());
try {
this.engine.eval(new InputStreamReader(fs.open(new Path(context.getConfiguration().get(SCRIPT_PATH)))));
this.engine.put(ARGS, context.getConfiguration().getStrings(SCRIPT_ARGS));
this.engine.eval(SETUP_ARGS);
} catch (Exception e) {
throw new InterruptedException(e.getMessage());
}
this.outputs = new SafeMapperOutputs(context);
}
@Override
public void map(final NullWritable key, final FaunusVertex value, final Mapper<NullWritable, FaunusVertex, NullWritable, FaunusVertex>.Context context) throws IOException, InterruptedException {
if (value.hasPaths()) {
final Object result;
try {
this.engine.put(V, value);
result = engine.eval(MAP_V_ARGS);
} catch (Exception e) {
throw new InterruptedException(e.getMessage());
}
this.textWritable.set((null == result) ? Tokens.NULL : result.toString());
this.outputs.write(Tokens.SIDEEFFECT, NullWritable.get(), this.textWritable);
}
this.outputs.write(Tokens.GRAPH, NullWritable.get(), value);
}
@Override
public void cleanup(final Mapper<NullWritable, FaunusVertex, NullWritable, FaunusVertex>.Context context) throws IOException, InterruptedException {
try {
this.engine.eval(CLEANUP_ARGS);
} catch (Exception e) {
throw new InterruptedException(e.getMessage());
}
this.outputs.close();
}
}
}