Initial import into git.

[galago.git] / java / pig-galago / src / com / yahoo / pig / PigServer.java

/*
 * Copyright (c) 2007 Yahoo! Inc.  All rights reserved.
 * See accompanying LICENSE file.
 */
package com.yahoo.pig;

import com.yahoo.pig.impl.galago.POGalago;
import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import org.apache.hadoop.dfs.DistributedFileSystem;
import org.apache.hadoop.fs.Path;

import com.yahoo.pig.builtin.BinStorage;
import com.yahoo.pig.builtin.PigStorage;
import com.yahoo.pig.data.*;
import com.yahoo.pig.impl.PigContext;
import com.yahoo.pig.impl.io.FileLocalizer;
import com.yahoo.pig.impl.logicalLayer.LORead;
import com.yahoo.pig.impl.logicalLayer.LOStore;
import com.yahoo.pig.impl.logicalLayer.LogicalOperator;
import com.yahoo.pig.impl.logicalLayer.LogicalPlan;
import com.yahoo.pig.impl.logicalLayer.LogicalPlanBuilder;
import com.yahoo.pig.impl.logicalLayer.parser.ParseException;
import com.yahoo.pig.impl.mapreduceExec.MapReduceLauncher;
import com.yahoo.pig.impl.physicalLayer.IntermedResult;
import com.yahoo.pig.impl.physicalLayer.POMapreduce;
import com.yahoo.pig.impl.physicalLayer.PhysicalPlan;
import com.yahoo.pig.impl.physicalLayer.POStore;
import com.yahoo.pig.impl.pigbodyExec.PigBodyLauncher;

/**
 * 
 * This class is the program's connection to Pig. Typically a program will create a PigServer
 * instance and call UseMapreduce(). The programmer then registers queries using registerQuery() and
 * retrieves results using openIterator() or store().
 * 
 */
public class PigServer {
    
    /**
     * The type of query execution
     */
    static public enum ExecType {
        /**
         * Run everything on the local machine
         */
        LOCAL,
        /**
         * Use the Hadoop Map/Reduce framework
         */
        MAPREDUCE,
        /**
         * Use the Experimental Hadoop framework; not available yet.
         */
        PIG,
        /**
         * Use the Galago TupleFlow framework
         */
        GALAGO;
    }
    
    private static ExecType parseExecType(String str) throws IOException {
        String normStr = str.toLowerCase();
        
        if (normStr.equals("local")) return ExecType.LOCAL;
        if (normStr.equals("mapreduce")) return ExecType.MAPREDUCE;
        if (normStr.equals("mapred")) return ExecType.MAPREDUCE;
        if (normStr.equals("pig")) return ExecType.PIG;
        if (normStr.equals("pigbody")) return ExecType.PIG;
        if (normStr.equals("galago")) return ExecType.GALAGO;
   
        throw new IOException("Unrecognized exec type: " + str);
    }
        /**
     *  a table mapping intermediate results to physical plans that read them (the state
     *  of how much has been read is maintained in PORead 
     */
    Map<IntermedResult, PhysicalPlan> physicalPlans = new HashMap<IntermedResult, PhysicalPlan>();
    /**
     * a table mapping ID's to intermediate results
     */
    Map<String, IntermedResult> queryResults = new HashMap<String, IntermedResult>();

    PigContext pigContext = null;
    
    public PigServer(String execType) throws IOException {
        this(parseExecType(execType));
    }
    
    public PigServer(String execType, String filesystemLocation, String jobtrackerLocation) throws IOException {
        this(parseExecType(execType), filesystemLocation, jobtrackerLocation);
    }
    
    public PigServer(ExecType execType) {
        this(execType, null, null);
    }
    
    public PigServer(ExecType execType, String filesystemLocation, String jobtrackerLocation) {
        pigContext = new PigContext(execType);
        
        // propagate the two locations to the configuration
        if (filesystemLocation != null) MapReduceLauncher.setFilesystemLocation(filesystemLocation);
        if (jobtrackerLocation != null) MapReduceLauncher.setJobtrackerLocation(jobtrackerLocation);
        if (filesystemLocation != null) PigBodyLauncher.setFilesystemLocation(filesystemLocation);
        if (jobtrackerLocation != null) PigBodyLauncher.setJobtrackerLocation(jobtrackerLocation);
    }

    public void debugOn() {
        MapReduceLauncher.debug = true;
    }
    
    public void debugOff() {
        MapReduceLauncher.debug = false;
    }
    
    /**
     * Defines an alias for the given function spec. This
     * is useful for functions that require arguments to the 
     * constructor.
     * 
     * @param alias - the new function alias to define.
     * @param functionSpec - the name of the function and any arguments.
     * It should have the form: classname('arg1', 'arg2', ...)
     */
    public void registerFunction(String function, String functionSpec) {
        pigContext.registerFunction(function, functionSpec);
    }
    
    public void registerJar(String path) {
        MapReduceLauncher.addJar(path);
    }
    
    /**
     * Register a query with the Pig runtime. The query is parsed and registered, but it is not
     * executed until it is needed.
     * 
     * @param query
     *            a Pig Latin expression to be evaluated.
     * @return a handle to the query.
     * @throws IOException
     */
    public void registerQuery(String query) throws IOException {
        // Bugzilla Bug 1006706 -- ignore empty queries
        //=============================================
        if(query != null) {
            query = query.trim();
            if(query.length() == 0) return;
        }else {
            return;
        }
            
        // parse the query into a logical plan
        LogicalPlan lp = null;
        try {
            lp = (new LogicalPlanBuilder(pigContext).parse(query, queryResults));
            // System.out.println(lp.toString());
        } catch (ParseException e) {
            throw (IOException) new IOException(e.getMessage()).initCause(e);
        }

        if (lp.alias != null)
            queryResults.put(lp.alias, new IntermedResult(lp));
    }
    
    public void registerSQL(String query) throws IOException {
        registerQuery(sqlToPigLatin(query));
    }
    
    private String sqlToPigLatin(String sql) throws IOException {
        // TODO
        throw new IOException("SQL support not yet implemented.");
    }
    
    public void newRelation(String id) {
        queryResults.put(id, new IntermedResult());
    }
    
    public void insertTuple(String id, Tuple t) throws IOException {
        if (!queryResults.containsKey(id)) throw new IOException("Attempt to insert tuple into nonexistent relation.");
        queryResults.get(id).add(t);
    }

    
    public Iterator<Tuple> openIterator(String id, boolean continueFromLast) throws IOException{
        if (!continueFromLast)
                return openIterator(id);
        
        if (pigContext.getExecType()!= ExecType.LOCAL){
                System.err.println("Streaming execution not supported in non-local mode.");
                System.exit(-1);
        }
                
        if (!queryResults.containsKey(id))
            throw new IOException("Invalid alias: " + id);

        IntermedResult readFrom = (IntermedResult) queryResults.get(id);
        
        if (readFrom.getOutputType()==LogicalOperator.FIXED){
                //Its not a continuous plan, just let it go through the non-continuous channel
                return openIterator(id);
        }
        
        PhysicalPlan pp = null;
        
        if (!physicalPlans.containsKey(readFrom)){
                
                //Not trying to do sharing, so won't compile and exec it as in the other cases
                //First check if some other operator tried to execute this intermediate result
                if (readFrom.executed() || readFrom.compiled()){
                        //Compile it again, so that POSlaves are not added in between query execution
                        readFrom.compile(queryResults);
                }
                LogicalPlan lp = new LogicalPlan(new LORead(pigContext, readFrom), pigContext);
                pp = new PhysicalPlan(lp,queryResults);
                physicalPlans.put(readFrom, pp);
        }else{
                pp = physicalPlans.get(readFrom);
        }
        
        return pp.exec(continueFromLast).content();
        
    }
    
    /**
     * Forces execution of query (and all queries from which it reads), in order to materialize
     * result
     */
    public Iterator<Tuple> openIterator(String id) throws IOException {
        if (!queryResults.containsKey(id))
            throw new IOException("Invalid alias: " + id);

        IntermedResult readFrom = (IntermedResult) queryResults.get(id);

        readFrom.compile(queryResults);
        readFrom.exec();
        if (pigContext.getExecType() == ExecType.LOCAL)
            return readFrom.read().content();
        final BinStorage p = new BinStorage();
        InputStream is = FileLocalizer.open(pigContext.getExecType(), readFrom.file);
        p.bindTo(is, 0, Long.MAX_VALUE);
        return new Iterator<Tuple>() {
            Tuple   t;
            boolean atEnd;

            public boolean hasNext() {
                if (atEnd)
                    return false;
                try {
                    if (t == null)
                        t = p.getNext();
                    if (t == null)
                        atEnd = true;
                } catch (Exception e) {
                    e.printStackTrace();
                    t = null;
                    atEnd = true;
                }
                return !atEnd;
            }

            public Tuple next() {
                Tuple next = t;
                if (next != null) {
                    t = null;
                    return next;
                }
                try {
                    next = p.getNext();
                } catch (Exception e) {
                    e.printStackTrace();
                }
                if (next == null)
                    atEnd = true;
                return next;
            }

            public void remove() {
                throw new RuntimeException("Removal not supported");
            }

        };
        
        }
    
    /**
     * Store an alias into a file
     * @param id: The alias to store
     * @param filename: The file to which to store to
     * @throws IOException
     */
                        
        public void store(String id, String filename) throws IOException {
                store(id, filename, PigStorage.class.getName() + "()");   // SFPig is the default store function
        }
        
        /**
         * Continuous case of store: store the updates to particular alias into the given file
         * @param id
         * @param filename
         * @throws IOException
         */
        
        public void update(String id, String filename, boolean append) throws IOException {
                update(id, filename, PigStorage.class.getName() + "()", append);
        }
        
        /**
         *  forces execution of query (and all queries from which it reads), in order to store result in file
         */
        public void store(String id, String filename, String func) throws IOException {
                filename = removeQuotes(filename);
                
        if (!queryResults.containsKey(id))
            throw new IOException("Invalid alias: " + id);

        IntermedResult readFrom = (IntermedResult) queryResults.get(id);
        
        if (!readFrom.executed()) {
            if (pigContext.getExecType() == ExecType.GALAGO) {
                LogicalOperator root = readFrom.lp.root();
                LogicalOperator store = new LOStore( pigContext, root, filename, func, false );
                LogicalPlan lp = new LogicalPlan( store, pigContext );
                readFrom = new IntermedResult( lp );
            }
            readFrom.compile(queryResults);
            if (pigContext.getExecType() == ExecType.MAPREDUCE) {
                POMapreduce pom = (POMapreduce)readFrom.pp.root;
                if (func == null) {
                        pom.outputFile = filename;
                } else {
                        pom.outputFile = filename;
                        pom.storeFunc = func;
                }
            }
            readFrom.exec();
            return;
        } 
        //else if (func != null) {
        //    //XXX We need to fix this!
        //    throw new IOException("Cannot currently restore an evaluated result");
        //}
        
        // TODO: following code assumes certain things about store/load functions
        
        if (pigContext.getExecType() != ExecType.LOCAL) {
            if (readFrom.file == null) {
                readFrom.toDFSFile(filename, func);
                return;
            }
            String src = readFrom.file.substring(FileLocalizer.HADOOP_PREFIX.length());
            String dst = filename;
            boolean islocal = false;
            if (dst.startsWith(FileLocalizer.LOCAL_PREFIX)) {
                dst = dst.substring(FileLocalizer.LOCAL_PREFIX.length());
                islocal = true;
            }
            if (dst.startsWith(FileLocalizer.HADOOP_PREFIX)) {
                dst = dst.substring(FileLocalizer.HADOOP_PREFIX.length());
            }
            if (readFrom.isTemporary()) {    // TODO: assumes we want to use same store func that was used last time we stored
                if (islocal) {
                    MapReduceLauncher.copy(src, dst, islocal);
                } else {
                    MapReduceLauncher.rename(src, dst);
                }
                readFrom.setPermanentFilename(dst);
            } else {
                if (islocal) {
                    MapReduceLauncher.copy(src, dst, islocal);
                }
            }
        } else {
            // generate a simple logical plan to store the results
            LogicalPlan lp = new LogicalPlan(new LOStore(pigContext, new LORead(pigContext, readFrom), filename, func,false), pigContext);

            // build a physical plan from the lp and then exec it,
            // this will write the file to disk
            (new PhysicalPlan(lp, queryResults)).exec(false);
        }               
    }

        /**
         * Continuous version of the store function above
         *
         */
        public void update(String id, String filename, String func, boolean append) throws IOException{
                
                
                if (pigContext.getExecType()!=ExecType.LOCAL){
                System.err.println("Streaming execution not supported in non-local mode.");
                System.exit(-1);
        }
                
                filename = removeQuotes(filename);
                
        if (!queryResults.containsKey(id))
            throw new IOException("Invalid alias: " + id);

        IntermedResult readFrom = (IntermedResult) queryResults.get(id);
        
        if (readFrom.getOutputType()==LogicalOperator.FIXED){
                //Its not a continuous plan, just let it go through the non-continuous channel
                store(id,filename,func);
                return;
        }
        
        PhysicalPlan pp = null;
        
        if (!physicalPlans.containsKey(readFrom)){
                
                //Not trying to do sharing, so won't compile and exec it as in the other cases
                //First check if some other operator tried to execute this intermediate result
                if (readFrom.executed() || readFrom.compiled()){
                        //Compile it again, so that POSlaves are not added in between query execution
                        readFrom.compile(queryResults);
                }
                LogicalPlan lp = new LogicalPlan(new LORead(pigContext, readFrom), pigContext);
                pp = new PhysicalPlan(lp,queryResults);
                physicalPlans.put(readFrom, pp);
        }else{
                pp = physicalPlans.get(readFrom);
        }
        new PhysicalPlan(new POStore(pp.root,filename,func,append)).exec(true);
        }



    /**
     * Returns the unused byte capacity of an HDFS filesystem. This value does
     * not take into account a replication factor, as that can vary from file
     * to file. Thus if you are using this to determine if you data set will fit
     * in the HDFS, you need to divide the result of this call by your specific replication
     * setting. 
     * @return
     * @throws IOException
     */
    public long capacity() throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("capacity only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            return dfs.getRawCapacity() - dfs.getRawUsed();
        }
    }

    /**
     * Returns the length of a file in bytes which exists in the HDFS (accounts for replication).
     * @param filename
     * @return
     * @throws IOException
     */
    public long fileSize(String filename) throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("stat only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            Path p = new Path(filename);
            long len = dfs.getLength(p);
            long replication = dfs.getDefaultReplication(); // did not work, for some reason: dfs.getReplication(p);
            return len * replication;
        }
    }
    
    public boolean existsFile(String filename) throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("existsFile only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            return dfs.exists(new Path(filename));
        }
    }
    
    public boolean deleteFile(String filename) throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("deleteFile only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            return dfs.delete(new Path(filename));
        }
    }
    
    public boolean renameFile(String source, String target) throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("renameFile only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            return dfs.rename(new Path(source), new Path(target));
        }
    }
    
    public boolean mkdirs(String dirs) throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("renameFile only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            return dfs.mkdirs(new Path(dirs));
        }
    }
    
    public String[] listPaths(String dir) throws IOException {
        if (pigContext.getExecType() == ExecType.LOCAL) {
            throw new IOException("renameFile only supported for non-local execution");
        } else {
            DistributedFileSystem dfs = (DistributedFileSystem) MapReduceLauncher.getDfs();
            Path paths[] = dfs.listPaths(new Path(dir));
            String strPaths[] = new String[paths.length];
            for (int i = 0; i < paths.length; i++) {
                strPaths[i] = paths[i].toString();
            }
            return strPaths;
        }
    }
    
    public long totalHadoopTimeSpent() {
        return MapReduceLauncher.totalHadoopTimeSpent;
    }

//    public static class PStat {
//        // These are the members of the unix fstat and stat structure
//        // we currently don't implement all of them, but try to provide
//        // a mirror for the function and values for HDFS files
//        // ====================================================
//        // struct stat {
//        // dev_t st_dev; /* device inode resides on */
//        // ino_t st_ino; /* inode's number */
//        // mode_t st_mode; /* inode protection mode */
//        // nlink_t st_nlink; /* number or hard links to the file */
//        // uid_t st_uid; /* user-id of owner */
//        // gid_t st_gid; /* group-id of owner */
//        // dev_t st_rdev; /* device type, for special file inode */
//        // struct timespec st_atimespec; /* time of last access */
//        // struct timespec st_mtimespec; /* time of last data modification */
//        // struct timespec st_ctimespec; /* time of last file status change */
//        // off_t st_size; /* file size, in bytes */
//        // quad_t st_blocks; /* blocks allocated for file */
//        // u_long st_blksize;/* optimal file sys I/O ops blocksize */
//        // u_long st_flags; /* user defined flags for file */
//        // u_long st_gen; /* file generation number */
//        // };
//        public long st_size;
//        public long st_blocks;
//    }

    private String removeQuotes(String str) {
        if (str.startsWith("\'") && str.endsWith("\'"))
            return str.substring(1, str.length() - 1);
        else
            return str;
    }
}