Initial import into git.

[galago.git] / java / pig-galago / src / com / yahoo / pig / impl / eval / EvalSpecPipe.java

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

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.*;

import com.yahoo.pig.AlgebraicAtomEvalFunc;
import com.yahoo.pig.AtomEvalFunc;
import com.yahoo.pig.EvalFunc;
import com.yahoo.pig.TupleEvalFunc;
import com.yahoo.pig.data.*;
import com.yahoo.pig.impl.PigContext;
import com.yahoo.pig.impl.eval.filter.FilterSpec;
import com.yahoo.pig.impl.eval.func.AtomFuncEvalItem;
import com.yahoo.pig.impl.eval.func.FuncEvalItem;
import com.yahoo.pig.impl.eval.func.TupleFuncEvalItem;
import com.yahoo.pig.impl.eval.sad.SADSpec;
import com.yahoo.pig.impl.logicalLayer.parser.ParseException;
import com.yahoo.pig.impl.logicalLayer.parser.QueryParser;
import com.yahoo.pig.impl.logicalLayer.schema.SchemaItem;

public class EvalSpecPipe {
    public List<EvalSpec> specs;
    PigContext pigContext;

    public EvalSpecPipe(PigContext pigContext) {
        this.pigContext = pigContext;
        specs = new ArrayList<EvalSpec>();
    }

    public EvalSpecPipe(PigContext pigContext, List<EvalSpec> specs) {
        this.pigContext = pigContext;
        this.specs = specs;
    }

    /**
     * Creates an EvalSpecPipe from an input string, by parsing it
     */
    public EvalSpecPipe(PigContext pigContext, String specStr) throws IOException {
        this.pigContext = pigContext;
        ByteArrayInputStream in = new ByteArrayInputStream(specStr.getBytes());
        QueryParser parser = new QueryParser(in, pigContext, null);
        try {
            specs = parser.PEvalSpecPipe().specs;
        } catch (ParseException e) {
            throw new IOException(e.getMessage());
        }
    }
    

    // creates a data collector that executes this EvalSpecPipe
    public DataCollector collector(DataCollector output) {

        // create datacollector pipeline
        for (int i = specs.size() - 1; i >= 0; i--) {
            output = specs.get(i).collector(output);
        }

        return output;
    }

    // creates a data collector, but ignores all SAD specs (to be used if SAD is applied to a tuple
    // rather than a bag)
    public DataCollector noSADCollector(DataCollector output) {

        // create datacollector pipeline w/out SAD's
        for (int i = specs.size() - 1; i >= 0; i--) {
            if (!(specs.get(i) instanceof SADSpec)) {
                output = specs.get(i).collector(output);
            }
        }

        return output;
    }

    public Tuple simpleEval(Tuple input) throws IOException {
        Tuple t = input;

        for (int i = 0; i < specs.size(); i++) {
            EvalSpec spec = specs.get(i);
            if (!(spec instanceof EvalItemList)) {
                throw new IOException("Internal error: Invalid use of simpleEval().");
            }

            t = ((EvalItemList) spec).simpleEval(t);
        }

        return t;
    }

    public DataBag simpleEval(DataBag b) throws IOException {

        if (specs.isEmpty()) {
            return b;
        } else {

            // create datacollector pipeline
            DataBag result = BagFactory.getInstance().getNewBag();
            DataCollector collector = collector(result);

            // push data into datacollector pipeline
            for (Iterator<Tuple> it = b.content(); it.hasNext();) {
                collector.add(it.next());
            }
            collector.add(null); // EOF

            return result;

        }
    }

    public Datum simpleEval(Datum input) throws IOException {

        if (input instanceof DataBag) {
            return simpleEval((DataBag) input);
        } else if (input instanceof Tuple) {
            return simpleEval((Tuple) input);
        } else {
            throw new IOException("Error: cannot apply nested projection to DataAtom.");
        }
    }

    public void add(EvalSpec spec) {

        // add spec to pipe, but avoid redundant "star" specs (just adds unnecessary eval overhead)

        if (specs.isEmpty()) {
            specs.add(spec);
        } else if (!spec.isJustStar()) {

            int lastSpec = specs.size() - 1;
            if (specs.get(lastSpec).isJustStar()) {
                specs.remove(lastSpec);
            }

            specs.add(spec);
        }
    }

    public void add(EvalSpecPipe otherPipe) {
        for (Iterator<EvalSpec> it = otherPipe.specs(); it.hasNext();) {
            add(it.next());
        }
    }

    public Iterator<EvalSpec> specs() {
        return specs.iterator();
    }

    public String toString() {
        StringBuffer sb = new StringBuffer();
        for (Iterator<EvalSpec> it = specs.iterator(); it.hasNext();) {
            EvalSpec es = it.next();
            sb.append("[");
            if (es instanceof EvalItemList) sb.append("EVAL ");
            else if (es instanceof FilterSpec) sb.append("FILTER ");
            sb.append(es);
            sb.append("]");
        }

        return sb.toString();
    }

    public List<String> getFuncs() {
        List<String> funcs = new ArrayList<String>();
        for (Iterator<EvalSpec> it = specs.iterator(); it.hasNext();) {
            funcs.addAll(it.next().getFuncs());
        }
        return funcs;
    }

    /**
     * This function will return a new EvalItemList which is
     * the combine portion of an algbraic evaluation of the
     * EvalSpecPipe instance. Before calling this method and
     * using the result, users should call amenableToCombiner() to
     * validate that the EvalSpecPipe *can* be used in a algebraic
     * manner.
     * @return
     */
    public EvalItemList makeIntoToCombine() {
        // This function takes the first spec in the pipe 
        // and searches it for the first instance of an
        // algebraic function to split the processing on
        // =============================================
        EvalItemList first = (EvalItemList) specs.get(0);
        EvalItemList combine = new EvalItemList(pigContext);
        boolean added = false;
        for (Iterator<EvalItem> it = first.columns.iterator(); it.hasNext();) {
            EvalItem item = it.next();
            if (item instanceof FuncEvalItem && !added) {
                EvalFunc func = ((FuncEvalItem) item).func;
                // this is our algebraic function, create a new function and break
                if (func instanceof AlgebraicAtomEvalFunc) {
                    TupleEvalFunc init = ((AlgebraicAtomEvalFunc) func).getInitial();
                    FuncEvalItem fei = new TupleFuncEvalItem(pigContext, init, ((FuncEvalItem) item).args);
                    combine.add(fei);   
                    added = true;
                    break;
                }
            } 
            //combine.add(item);
        }
        return combine;
    }

    /**
     * This function will return a new EvalItemList which is
     * the reduce portion of an algbraic evaluation of the
     * EvalSpecPipe instance. Before calling this method and
     * using the result, users should call amenableToCombiner() to
     * validate that the EvalSpecPipe *can* be used in a algebraic
     * manner.
     * @return
     */
    public EvalItemList makeIntoToReduce() {
        // This function finds the proper algrebraic function which was
        // processed by makeIntoToCombine() and then creates the proper
        // spec for the reduce portion of the job. 
        //========================================
        EvalItemList first = (EvalItemList) specs.get(0);
        EvalItemList reduce = new EvalItemList(pigContext);
        boolean firstItem = true;
        boolean added = false;
        int colPos = 1; // if group was not specificed
        
        for (EvalItem item : first.columns) {
            // We have to see if the group column was used in the eval spec
            // because it is implicitly included in the output of co-group
            // and columns need to be shifted appro.
            //======================================
            if( item instanceof ColEvalItem && firstItem) {
                ColEvalItem col = (ColEvalItem) item;
                if(col.colNum == 0 || (col.alias != null & col.alias.equals("group"))) {
                    colPos--;
                }
            } else if (item instanceof FuncEvalItem && !added) {
                EvalFunc func = ((FuncEvalItem) item).func;
                // this is our algebraic function!
                if (func instanceof AlgebraicAtomEvalFunc) { 
                    AtomEvalFunc fin = ((AlgebraicAtomEvalFunc) func).getFinal();
                    ColEvalItem cei = new ColEvalItem(pigContext, colPos);
                    EvalItemList newargs = new EvalItemList(pigContext);
                   
                    newargs.add(cei);
                    FuncEvalItem fei = new AtomFuncEvalItem(pigContext, fin, newargs);
                    
                    // don't forget to add the post-processing specs to the
                    // output of the algbraic function
                    //================================
                    fei.nestedEval = ((FuncEvalItem) item).nestedEval;
                    fei.subColSpec = ((FuncEvalItem) item).subColSpec;
                    reduce.add(fei);
                    added = true;
                    continue;
                }
            }
            reduce.add(item);
            colPos++; 
            firstItem = false;
        }
        return reduce;
    }
    

    /**
     * Process the pipe and determine if the pipe is a candidate for 
     * algebraic evaluation. This is mostly reliant on determining if
     * the EvalItemLists contain algebraic functions and have unique column
     * references.
     * @return
     */
    public boolean amenableToCombiner() {
        if (specs.size() > 0 && specs.get(0) instanceof EvalItemList) {
            EvalItemList first = (EvalItemList) specs.get(0);
            return first.amenableToCombiner();
        } else {
            return false;
        }
    }
    
    /**
     * Given an input schema, determine the output schema of this pipe
     * as it operates on input tuples with the input schema.
     * @param input
     * @return
     */
    public SchemaItem mapInputSchema(SchemaItem input) {
        SchemaItem output = input;
        for(EvalSpec spec : specs) {
                output = spec.mapInputSchema(output);
        }
        return output;
    }
    

}