Initial import into git.

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

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

import java.io.*;
import java.util.*;

import com.yahoo.pig.AlgebraicAtomEvalFunc;
import com.yahoo.pig.AlgebraicBagEvalFunc;
import com.yahoo.pig.AlgebraicTupleEvalFunc;
import com.yahoo.pig.EvalFunc;
import com.yahoo.pig.data.*;
import com.yahoo.pig.impl.PigContext;
import com.yahoo.pig.impl.eval.func.*;
import com.yahoo.pig.impl.logicalLayer.parser.ParseException;
import com.yahoo.pig.impl.logicalLayer.parser.QueryParser;
import com.yahoo.pig.impl.logicalLayer.schema.SchemaItem;
import com.yahoo.pig.impl.logicalLayer.schema.SchemaItemList;

public class EvalItemList extends EvalSpec {
    public List<EvalItem> columns = null;
    boolean               isJustStar;
    int                   driver;
    PigContext pigContext;
    LinkedList<CrossProductItem> pendingCrossProducts = new LinkedList<CrossProductItem>();
    
    /**
     * Creates an empty EvalItemList
     */
    public EvalItemList(PigContext pigContext) {
        this.pigContext = pigContext;
        columns = new ArrayList<EvalItem>();
    }

    /**
     * Creates an EvalSpec from an input string, by parsing it
     */
    public EvalItemList(PigContext pigContext, String specStr) throws IOException, ParseException {
        this.pigContext = pigContext;
        ByteArrayInputStream in = new ByteArrayInputStream(specStr.getBytes());
        QueryParser parser = new QueryParser(in, pigContext, null);
        columns = parser.PEvalItemList().columns;
    }

    /**
     * add an item to projection list
     */
    public void add(EvalItem item) {
        columns.add(item);
    }

    public boolean isSimple(){
        selectDriver();
        if (driver==-1)
                return true;
        else
                return false;
    }
    
    /**
     * creates a data collector that executes this EvalItemList
     */
    public DataCollector collector(final DataCollector output) {
        isJustStar = isJustStar();
        selectDriver();

        return new DataCollector(output) {

            public void add(Tuple t) throws IOException {
                if (t != null){
                        if (pendingCrossProducts.peek()!=null)
                                System.err.println(Thread.currentThread().getName() + ": Before eval, head of queue was " + pendingCrossProducts.peek() + " with input " + pendingCrossProducts.peek().cpiInput.hashCode());
                        eval(t, output);
                        if (pendingCrossProducts.peek()!=null)
                                System.err.println(Thread.currentThread().getName() + ": After eval, head of queue was " + pendingCrossProducts.peek() + " with input " + pendingCrossProducts.peek().cpiInput.hashCode());

                }else{
                        for (EvalItem item: columns){
                                if (driver==-1 || columns.get(driver)!=item)
                                        item.finish();
                        }
                        while (!pendingCrossProducts.isEmpty()){
                                CrossProductItem cpi = pendingCrossProducts.remove();
                        System.err.println(Thread.currentThread().getName() + ": Dequeuing " + cpi + "having input " + cpi.cpiInput.hashCode());
                                cpi.waitToBeReady();
                                cpi.exec();
                        }
                        if (driver!=-1)
                                columns.get(driver).finish();
                        synchronized(output){
                                while (output.isStale()){
                                        try{
                                                output.wait();
                                        }catch(InterruptedException e){}
                                }
                        }
                        output.add(null); // pass on EOF
                }
            }
        };
    }

    /**
     * only handles cases with no nested projection or subcols or functions
     */
    public Tuple simpleEval(Tuple input) throws IOException {
        // make it go fast in the case of '*'
        if (isJustStar)
            return input;

        ArrayList<EvalItem> items = expandStars(columns, input.arity());

        Tuple output = new Tuple(items.size());

        for (int i = 0; i < items.size(); i++) {
            output.setField(i, items.get(i).simpleEval(input));
        }

        return output;
    }

    private class CrossProductItem{
        DataBag[] toBeCrossed;
        final Tuple cpiInput;
        DataCollector output, assembleTuples;
        
        public CrossProductItem(Tuple input, DataCollector output) throws IOException{
                
                this.cpiInput = new Tuple(input.arity());
                for (int i=0; i<input.arity(); i++){
                        cpiInput.setField(i, input.getField(i));
                }
                
                this.output = output;
                
                // materialize data for all to-be-crossed items
            // (except driver, which is done in streaming fashion)
            toBeCrossed = new DataBag[columns.size()];
            for (int i = 0; i < columns.size(); i++) {
                if (i == driver)
                    continue;

                EvalItem item = columns.get(i);
                if (!item.isSimple()) {
                    toBeCrossed[i] = BagFactory.getInstance().getNewBag();
                    item.eval(input, toBeCrossed[i]);
                }
            }
        }
        
        public void initAssembler(){
                if (assembleTuples!=null)
                        return;
                
                final int numItems = columns.size();
                
                assembleTuples = new DataCollector(output) {

               public void add(Tuple t) throws IOException {
                
                   if (t == null) { // EOF
                       output.add(null);
                       return;
                   }

                   // create one iterator per to-be-crossed bag
                   Iterator[] its = new Iterator[numItems];
                   int numIts = 0;
                   for (int i = 0; i < numItems; i++) {
                       if (toBeCrossed[i] != null) {
                           numIts++;
                           its[i] = toBeCrossed[i].content();
                           if (!its[i].hasNext())
                               return; // one of inputs is empty, so cross-prod yields empty result
                       }
                   }

                   Tuple[] lastOutput = null;
                   Tuple[] outData = new Tuple[numItems];

                   boolean done = false;
                   while (!done) {
                       if (lastOutput == null) { // we're generating our first output
                           for (int i = 0; i < numItems; i++) {
                               if (i == driver)
                                   outData[i] = t;
                               else if (toBeCrossed[i] == null)
                                   outData[i] = new Tuple(columns.get(i).simpleEval(cpiInput));
                               else
                                   outData[i] = (Tuple) its[i].next();
                           }
                       } else {
                           boolean needToAdvance = true;

                           for (int i = 0; i < numItems; i++) {
                               if (its[i] != null && needToAdvance) {
                                   if (its[i].hasNext()) {
                                       outData[i] = (Tuple) its[i].next();
                                       needToAdvance = false;
                                   } else {
                                       its[i] = toBeCrossed[i].content(); // rewind iterator
                                       outData[i] = (Tuple) its[i].next();
                                       // still need to advance some other input..
                                   }
                               } else {
                                   outData[i] = lastOutput[i]; // use same value as last time
                               }
                           }
                       }

                       // check for completion:
                       done = true;
                       if (numIts > 0) { // done iff all iterators empty
                           for (int i = 0; i < numItems; i++) {
                               if (its[i] != null && its[i].hasNext()) {
                                   done = false;
                                   break;
                               }
                           }
                       }

                       output.add(new Tuple(outData));

                       lastOutput = outData;
                   }
               }
           };
                
        }
        
        public boolean isReady(){
                for (int i=0; i<toBeCrossed.length; i++){
                        if (toBeCrossed[i]!=null && toBeCrossed[i].isStale())
                                return false;
                }
                
                return true;
        }
        
        public void waitToBeReady(){
                for (int i=0; i<toBeCrossed.length; i++){
                        if (toBeCrossed[i]!=null){
                                synchronized(toBeCrossed[i]){
                                        while (toBeCrossed[i].isStale()){
                                                try{
                                                        toBeCrossed[i].wait();
                                                }catch (InterruptedException e){}
                                        }
                                }
                        }
                }
                
        }
        
        public void exec() throws IOException{
            initAssembler();
            columns.get(driver).eval(cpiInput, assembleTuples);
            System.err.println(Thread.currentThread().getName() + ": Executing driver on " + cpiInput);
            output.markStale(false);
        }
        
    }
    /**
     * Process the actual eval of this EvalItemList
     * @param input
     * @param output
     * @throws IOException
     */
    private void eval(Tuple input, DataCollector output) throws IOException {
        // make it go fast in the case of single item
        if (columns.size() == 1) {
            columns.get(0).eval(input, output);
            return;
        }

        if (driver == -1) { // all items are simple, so do simpleEval()
            output.add(simpleEval(input));
            return;
        }

        // general case (use driver method):
        CrossProductItem cpi = new CrossProductItem(input,output);
        
        pendingCrossProducts.addLast(cpi);

        
        //Since potentially can return without filling output, mark output as stale
        //the exec method of CrossProductItem will mark output as not stale
        output.markStale(true);
        while (pendingCrossProducts.size() > 0 && pendingCrossProducts.peek().isReady()){
                pendingCrossProducts.remove().exec();
        }
    }

    private void selectDriver() {
        driver = -1;

        for (int i = 0; i < columns.size(); i++) {
            EvalItem item = columns.get(i);
            if (!item.isSimple()) {
                if (item instanceof BagFuncEvalItem) { // trumps 'em all
                    driver = i;
                    return;
                } else {
                    driver = i; // we'll use this as the driver, unless something better comes along
                }
            }
        }
    }

    /**
     * convert each star column into a list of regular columns.
     */
    public ArrayList<EvalItem> expandStars(List<EvalItem> items, int nCols) {
        ArrayList<EvalItem> output = new ArrayList<EvalItem>();

        for (Iterator<EvalItem> it = items.iterator(); it.hasNext();) {
            EvalItem item = it.next();

            if (item instanceof StarEvalItem) {
                for (int col = 0; col < nCols; col++)
                    output.add(new ColEvalItem(pigContext, col));
            } else {
                output.add(item);
            }
        }
        return output;
    }

    /**
     * Determine if this instance of EvalItems is a candiate for algebraic
     * evaluation. This means it contains an Algebraic Function, and does not
     * contain repeated references to column used by the algebraic function.
     * @return
     */
    public boolean amenableToCombiner() {
        for (Iterator<EvalItem> it = columns.iterator(); it.hasNext();) {
            EvalItem item = it.next();
            if (item instanceof FuncEvalItem) {
                EvalFunc func = ((FuncEvalItem) item).func;
                if ((func instanceof AlgebraicBagEvalFunc || func instanceof AlgebraicAtomEvalFunc || func instanceof AlgebraicTupleEvalFunc)
                        && item.isSimple()) {
                    for (EvalItem arg : ((FuncEvalItem) item).args.columns) {
                        if (arg instanceof ColEvalItem) {
                            if (!checkForUnique((FuncEvalItem) item, (ColEvalItem) arg)) {
                                System.out.println("Column is not unique!");
                                return false; // repeated column reference, no dice
                            }
                        } else {
                            return false;
                        }
                    }
                    return true;
                } else if (item instanceof StarEvalItem) {
                    System.out.println("Star precludes algebraic!");
                    return false;
                }
            }
        }
        return false;
    }

    /**
     * Helper function for amenableToCombiner(). This will recurse a function
     * eval spec and ensure that input columns are not repeatedly reference. 
     * @param func
     * @param col
     * @return
     */
    public boolean checkForUnique(FuncEvalItem func, ColEvalItem col) {
        for (EvalItem item : columns) {
            if (item instanceof FuncEvalItem) {
                FuncEvalItem fi = (FuncEvalItem) item;
                for (EvalItem ei : fi.args.columns) {
                    if (ei instanceof ColEvalItem) {
                        // You will need to make sure you don't disqualify
                        // due to matching against the arguments
                        if (fi == func && ei == col) {
                            continue;
                        }
                        if (!fi.args.checkForUnique(func, col)) {
                            return false;
                        }
                    }
                }
            } else if (item instanceof ColEvalItem) {
                ColEvalItem ci = (ColEvalItem) item;
                if (ci.colNum == col.colNum) {
                    return false;
                }
            } else if (item instanceof StarEvalItem) {
                return false;
            }
        }
        return true;
    }

    public String toString() {
        StringBuffer sb = new StringBuffer();
        for (Iterator<EvalItem> it = columns.iterator(); it.hasNext();) {
            sb.append(it.next());
            if (it.hasNext())
                sb.append(", ");
        }

        return sb.toString();
    }

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

    public boolean isJustStar() {
        return (columns.size() == 1 && columns.get(0) instanceof StarEvalItem);
    }

    public SchemaItem mapInputSchema(SchemaItem input) {
        SchemaItemList output = new SchemaItemList();
        
        Set<String> aliases = new HashSet<String>();
        Set<String> conflictedAliases = new HashSet<String>();
        
        for (EvalItem item: columns) {
                SchemaItem sItem = item.mapInputSchema(input);
                                
                if (item instanceof NestableEvalItem && ((NestableEvalItem)item).subColSpec!=null){
                        for(SchemaItem flattenedItem: sItem.flatten()){
                                
                                if (aliases.contains(flattenedItem.alias))
                                        conflictedAliases.add(flattenedItem.alias);
                                else
                                        aliases.add(flattenedItem.alias);
                        }
                }else{
                        if (aliases.contains(sItem.alias))
                                conflictedAliases.add(sItem.alias);
                }
        }
        
        for (EvalItem item: columns) {
                SchemaItem sItem = item.mapInputSchema(input);
                
                if (item instanceof NestableEvalItem && ((NestableEvalItem)item).subColSpec!=null){
                        for (SchemaItem flattenedItem: sItem.flatten()){
                                if (conflictedAliases.contains(flattenedItem.alias)){
                                        if (sItem.alias!=null){
                                                flattenedItem.alias = sItem.alias + "::" + flattenedItem.alias;
                                        }
                                }
                                output.add(flattenedItem);
                        }
                }else{
                        output.add(sItem);
                }
        }
        return output;
    }

    public Comparator<Tuple> getComparator() {
        return new Comparator<Tuple>() {
            private List<EvalItem> cols = columns;

            public int compare(Tuple t1, Tuple t2) {
                for (Iterator<EvalItem> it = cols.iterator(); it.hasNext();) {
                    EvalItem itm = it.next();

                    int comp = 0;

                    if (itm instanceof StarEvalItem) {
                        comp = t1.compareTo(t2);
                    } else if (itm instanceof ColEvalItem) {
                        int colNum = ((ColEvalItem) itm).colNum;
                        try {
                            comp = t1.getField(colNum).compareTo(t2.getField(colNum));
                        } catch (IOException e) {
                            throw new RuntimeException(e.getMessage());
                        }
                    } else if (itm instanceof ConstEvalItem) {
                        comp = 0;
                    } else {
                        throw new RuntimeException("Internal error: unexpected EvalItem type.");
                    }

                    if (comp != 0)
                        return comp;
                }

                return 0;
            }

        };
    }

}