spork-master/contrib/piggybank/java/src/main/java/org/apache/pig/piggybank/storage/AllLoader.java
/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with this
 * work for additional information regarding copyright ownership. The ASF
 * licenses this file 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 org.apache.pig.piggybank.storage;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.UnsupportedEncodingException;
import java.util.Arrays;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import org.apache.commons.codec.binary.Base64;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;
import org.apache.log4j.Logger;
import org.apache.pig.Expression;
import org.apache.pig.FileInputLoadFunc;
import org.apache.pig.FuncSpec;
import org.apache.pig.LoadCaster;
import org.apache.pig.LoadFunc;
import org.apache.pig.LoadMetadata;
import org.apache.pig.LoadPushDown;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceSchema.ResourceFieldSchema;
import org.apache.pig.ResourceStatistics;
import org.apache.pig.StoreMetadata;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigSplit;
import org.apache.pig.builtin.Utf8StorageConverter;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.PigContext;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema;
import org.apache.pig.impl.util.UDFContext;
import org.apache.pig.piggybank.storage.allloader.LoadFuncHelper;
import org.apache.pig.piggybank.storage.partition.PathPartitionHelper;

/**
 * The AllLoader provides the ability to point pig at a folder that contains
 * files in multiple formats e.g. PlainText, Gz, Bz, Lzo, HiveRC etc and have
 * the LoadFunc(s) automatically selected based on the file extension. <br/>
 * <b>How this works:<b/><br/>
 * The file extensions are mapped in the pig.properties via the property
 * file.extension.loaders.
 * 
 * <p/>
 * <b>file.extension.loaders format</b>
 * <ul>
 * <li>[file extension]:[loader func spec]</li>
 * <li>[file-extension]:[optional path tag]:[loader func spec]</li>
 * <li>[file-extension]:[optional path tag]:[sequence file key value writer
 * class name]:[loader func spec]</li>
 * </ul>
 * 
 * <p/>
 * The file.extension.loaders property associate pig loaders with file
 * extensions, if a file does not have an extension the AllLoader will look at
 * the first three bytes of a file and try to guess its format bassed on:
 * <ul>
 * <li>[ -119, 76, 90 ] = lzo</li>
 * <li>[ 31, -117, 8 ] = gz</li>
 * <li>[ 66, 90, 104 ] = bz2</li>
 * <li>[ 83, 69, 81 ] = seq</li>
 * </ul>
 * <br/>
 * The loader associated with that extension will then be used.
 * <p/>
 * 
 * <b>Path partitioning</b> The AllLoader supports hive style path partitioning
 * e.g. /log/type1/daydate=2010-11-01<br/>
 * "daydate" will be considered a partition key and filters can be written
 * against this.<br/>
 * Note that the filter should go into the AllLoader contructor e.g.<br/>
 * a = LOAD 'input' using AllLoader('daydate<\"2010-11-01\"')<br/>
 * 
 * <b>Path tags</b> AllLoader supports configuring different loaders for the
 * same extension based on there file path.<br/>
 * E.g.<br/>
 * We have the paths /log/type1, /log/type2<br/>
 * For each of these directories we'd like to use different loaders.<br/>
 * So we use setup our loaders:<br/>
 * file.extension.loaders:gz:type1:MyType1Loader, gz:type2:MyType2Loader<br/>
 * 
 * 
 * <p/>
 * <b>Sequence files<b/> Sequence files also support using the Path tags for
 * loader selection but has an extra configuration option that relates to the
 * Key Class used to write the Sequence file.<br/>
 * E.g. for HiveRC this value is: org.apache.hadoop.hive.ql.io.RCFile so we can
 * setup our sequence file formatting:<br/>
 * file.extension.loaders:seq::org.apache.hadoop.hive.ql.io.RCFile:
 * MyHiveRCLoader, seq::DefaultSequenceFileLoader<br/>
 * 
 * <p/>
 * <b>Schema</b> The JsoneMetadata schema loader is supported and the schema
 * will be loaded using this loader.<br/>
 * In case this fails, the schema can be loaded using the default schema
 * provided.
 * 
 */
public class AllLoader extends FileInputLoadFunc implements LoadMetadata,
        StoreMetadata, LoadPushDown {

    private static final Logger LOG = Logger.getLogger(AllLoader.class);

    private static final String PROJECTION_ID = AllLoader.class.getName()
            + ".projection";

    transient LoadFunc childLoadFunc;
    transient boolean supportPushDownProjection = false;
    transient RequiredFieldList requiredFieldList;
    transient SortedSet<Integer> requiredFieldHashSet;

    transient TupleFactory tupleFactory = TupleFactory.getInstance();
    transient ResourceSchema schema;

    String signature;

    /**
     * Implements the logic for searching partition keys and applying parition
     * filtering
     */
    transient PathPartitionHelper pathPartitionerHelper = new PathPartitionHelper();
    transient Map<String, String> currentPathPartitionKeyMap;
    transient String[] partitionColumns;

    transient JsonMetadata jsonMetadata;
    transient boolean partitionKeysSet = false;

    LoadFuncHelper loadFuncHelper = null;

    transient Configuration conf;
    transient Path currentPath;

    String constructorPassedPartitionFilter;

    public AllLoader() {
        jsonMetadata = new JsonMetadata();
    }

    public AllLoader(String partitionFilter) {
        this();
        LOG.debug("PartitionFilter: " + partitionFilter.toString());

        constructorPassedPartitionFilter = partitionFilter;

    }

    @Override
    public void setLocation(String location, Job job) throws IOException {
        FileInputFormat.setInputPaths(job, location);
        // called on the front end
        conf = job.getConfiguration();
        loadFuncHelper = new LoadFuncHelper(conf);

        if (constructorPassedPartitionFilter != null) {

            pathPartitionerHelper.setPartitionFilterExpression(
                    constructorPassedPartitionFilter, AllLoader.class,
                    signature);

        }

        getPartitionKeys(location, job);
    }

    @Override
    public LoadCaster getLoadCaster() throws IOException {
        return new Utf8StorageConverter();
    }

    @Override
    public AllLoaderInputFormat getInputFormat() throws IOException {
        // this plugs the AllLoaderInputFormat into the system, which in turn
        // will plug in the AllRecordReader
        // the AllRecordReader will select and create the correct LoadFunc
        return new AllLoaderInputFormat(signature);
    }

    @Override
    public void prepareToRead(
            @SuppressWarnings("rawtypes") RecordReader reader, PigSplit split)
            throws IOException {

        AllReader allReader = (AllReader) reader;

        if (currentPath == null || !(currentPath.equals(allReader.path))) {
            currentPathPartitionKeyMap = (partitionColumns == null) ? null
                    : pathPartitionerHelper
                            .getPathPartitionKeyValues(allReader.path
                                    .toString());
            currentPath = allReader.path;
        }

        childLoadFunc = allReader.prepareLoadFuncForReading(split);

        String projectProperty = getUDFContext().getProperty(PROJECTION_ID);

        if (projectProperty != null) {

            // load the required field list from the current UDF context
            ByteArrayInputStream input = new ByteArrayInputStream(
                    Base64.decodeBase64(projectProperty.getBytes("UTF-8")));

            ObjectInputStream objInput = new ObjectInputStream(input);

            try {
                requiredFieldList = (RequiredFieldList) objInput.readObject();
            } catch (ClassNotFoundException e) {
                throw new FrontendException(e.toString(), e);
            } finally {
                IOUtils.closeStream(objInput);
            }

            if (childLoadFunc.getClass().isAssignableFrom(LoadPushDown.class)) {
                supportPushDownProjection = true;
                ((LoadPushDown) childLoadFunc)
                        .pushProjection(requiredFieldList);
            } else {
                if (requiredFieldList != null) {
                    requiredFieldHashSet = new TreeSet<Integer>();
                    for (RequiredField requiredField : requiredFieldList
                            .getFields()) {
                        requiredFieldHashSet.add(requiredField.getIndex());
                    }
                }
            }

        }

    }

    @Override
    public Tuple getNext() throws IOException {
        // delegate work to the child load func selected based on the file type
        // and other criteria
        // We do support PushDown Projection if the LoadFunc does not so
        // in this method we need to look at the childLoadFunc flag
        // (supportPushDownProjection )
        // if true we use the getNext method as is, if not we remove the fields
        // not required in the spushDownProjection.

        Tuple tuple = null;

        if (supportPushDownProjection) {
            tuple = childLoadFunc.getNext();
        } else if ((tuple = childLoadFunc.getNext()) != null) {
            // ----- If the function does not support projection we do it here

            if (requiredFieldHashSet != null) {

                Tuple projectedTuple = tupleFactory
                        .newTuple(requiredFieldHashSet.size());
                int i = 0;
                int tupleSize = tuple.size();

                for (int index : requiredFieldHashSet) {
                    if (index < tupleSize) {
                        // add the tuple columns
                        projectedTuple.set(i++, tuple.get(index));
                    } else {
                        // add the partition columns
                        projectedTuple.set(i++, currentPathPartitionKeyMap
                                .get(partitionColumns[index - tupleSize]));
                    }
                }

                tuple = projectedTuple;
            }

        }

        return tuple;
    }

    @Override
    public List<OperatorSet> getFeatures() {
        return Arrays.asList(LoadPushDown.OperatorSet.PROJECTION);
    }

    @Override
    public RequiredFieldResponse pushProjection(
            RequiredFieldList requiredFieldList) throws FrontendException {
        // save the required field list to the UDFContext properties.

        Properties properties = getUDFContext();

        ByteArrayOutputStream byteArray = new ByteArrayOutputStream();
        ObjectOutputStream objOut = null;
        try {
            objOut = new ObjectOutputStream(byteArray);
            objOut.writeObject(requiredFieldList);
        } catch (IOException e) {
            throw new FrontendException(e.toString(), e);
        } finally {
            IOUtils.closeStream(objOut);
        }

        // write out the whole required fields list as a base64 string
        try {
            properties.setProperty(PROJECTION_ID,
                    new String(Base64.encodeBase64(byteArray.toByteArray()),
                            "UTF-8"));
        } catch (UnsupportedEncodingException e) {
            throw new FrontendException(e.toString(), e);
        }

        return new RequiredFieldResponse(true);
    }

    /**
     * Tries to determine the LoadFunc by using the LoadFuncHelper to identify a
     * loader for the first file in the location directory.<br/>
     * If no LoadFunc can be determine ad FrontendException is thrown.<br/>
     * If the LoadFunc implements the LoadMetadata interface and returns a non
     * null schema this schema is returned.
     * 
     * @param location
     * @param job
     * @return
     * @throws IOException
     */
    private ResourceSchema getSchemaFromLoadFunc(String location, Job job)
            throws IOException {

        ResourceSchema schema = null;

        if (loadFuncHelper == null) {
            loadFuncHelper = new LoadFuncHelper(job.getConfiguration());
        }

        Path firstFile = loadFuncHelper.determineFirstFile(location);

        if (childLoadFunc == null) {

            // choose loader
            FuncSpec funcSpec = loadFuncHelper.determineFunction(location,
                    firstFile);

            if (funcSpec == null) {
                // throw front end exception, no loader could be determined.
                throw new FrontendException(
                        "No LoadFunction could be determined for " + location);
            }

            childLoadFunc = (LoadFunc) PigContext
                    .instantiateFuncFromSpec(funcSpec);
        }

        LOG.debug("Found LoadFunc:  " + childLoadFunc.getClass().getName());

        if (childLoadFunc instanceof LoadMetadata) {
            schema = ((LoadMetadata) childLoadFunc).getSchema(firstFile.toUri()
                    .toString(), job);
            LOG.debug("Found schema " + schema + " from loadFunc:  "
                    + childLoadFunc.getClass().getName());
        }

        return schema;
    }

    @Override
    public ResourceSchema getSchema(String location, Job job)
            throws IOException {

        if (schema == null) {
            ResourceSchema foundSchema = jsonMetadata.getSchema(location, job);

            // determine schema from files in location
            if (foundSchema == null) {
                foundSchema = getSchemaFromLoadFunc(location, job);

            }

            // only add the partition keys if the schema is not null
            // we use the partitionKeySet to only set partition keys once.
            if (!(partitionKeysSet || foundSchema == null)) {
                String[] keys = getPartitionColumns(location, job);

                if (!(keys == null || keys.length == 0)) {

                    // re-edit the pigSchema to contain the new partition keys.
                    ResourceFieldSchema[] fields = foundSchema.getFields();

                    LOG.debug("Schema: " + Arrays.toString(fields));

                    ResourceFieldSchema[] newFields = Arrays.copyOf(fields,
                            fields.length + keys.length);

                    int index = fields.length;

                    for (String key : keys) {
                        newFields[index++] = new ResourceFieldSchema(
                                new FieldSchema(key, DataType.CHARARRAY));
                    }

                    foundSchema.setFields(newFields);

                    LOG.debug("Added partition fields: " + keys
                            + " to loader schema");
                    LOG.debug("Schema is: " + Arrays.toString(newFields));
                }

                partitionKeysSet = true;

            }

            schema = foundSchema;
        }

        return schema;
    }

    @Override
    public ResourceStatistics getStatistics(String location, Job job)
            throws IOException {
        return null;
    }

    @Override
    public void storeStatistics(ResourceStatistics stats, String location,
            Job job) throws IOException {

    }

    @Override
    public void storeSchema(ResourceSchema schema, String location, Job job)
            throws IOException {
        jsonMetadata.storeSchema(schema, location, job);
    }

    /**
     * Reads the partition columns
     * 
     * @param location
     * @param job
     * @return
     */
    private String[] getPartitionColumns(String location, Job job) {

        if (partitionColumns == null) {
            // read the partition columns from the UDF Context first.
            // if not in the UDF context then read it using the PathPartitioner.

            Properties properties = getUDFContext();

            if (properties == null) {
                properties = new Properties();
            }

            String partitionColumnStr = properties
                    .getProperty(PathPartitionHelper.PARTITION_COLUMNS);

            if (partitionColumnStr == null
                    && !(location == null || job == null)) {
                // if it hasn't been written yet.
                Set<String> partitionColumnSet;

                try {
                    partitionColumnSet = pathPartitionerHelper
                            .getPartitionKeys(location, job.getConfiguration());
                } catch (IOException e) {

                    RuntimeException rte = new RuntimeException(e);
                    rte.setStackTrace(e.getStackTrace());
                    throw rte;

                }

                if (partitionColumnSet != null) {

                    StringBuilder buff = new StringBuilder();

                    int i = 0;
                    for (String column : partitionColumnSet) {
                        if (i++ != 0) {
                            buff.append(',');
                        }

                        buff.append(column);
                    }

                    String buffStr = buff.toString().trim();

                    if (buffStr.length() > 0) {

                        properties.setProperty(
                                PathPartitionHelper.PARTITION_COLUMNS,
                                buff.toString());
                    }

                    partitionColumns = partitionColumnSet
                            .toArray(new String[] {});

                }

            } else {
                // the partition columns has been set already in the UDF Context
                if (partitionColumnStr != null) {
                    String split[] = partitionColumnStr.split(",");
                    Set<String> partitionColumnSet = new LinkedHashSet<String>();
                    if (split.length > 0) {
                        for (String splitItem : split) {
                            partitionColumnSet.add(splitItem);
                        }
                    }

                    partitionColumns = partitionColumnSet
                            .toArray(new String[] {});
                }

            }

        }

        return partitionColumns;

    }

    @Override
    public String[] getPartitionKeys(String location, Job job)
            throws IOException {

        String[] partitionKeys = getPartitionColumns(location, job);

        if (partitionKeys == null) {
            throw new NullPointerException("INDUCED");
        }
        LOG.info("Get Parition Keys for: " + location + " keys: "
                + Arrays.toString(partitionKeys));

        return partitionKeys;
    }

    // --------------- Save Signature and PartitionFilter Expression
    // ----------------- //
    @Override
    public void setUDFContextSignature(String signature) {
        this.signature = signature;
        super.setUDFContextSignature(signature);
    }

    private Properties getUDFContext() {
        return UDFContext.getUDFContext().getUDFProperties(this.getClass(),
                new String[] { signature });
    }

    @Override
    public void setPartitionFilter(Expression partitionFilter)
            throws IOException {
        LOG.debug("PartitionFilter: " + partitionFilter.toString());

        pathPartitionerHelper.setPartitionFilterExpression(
                partitionFilter.toString(), AllLoader.class, signature);

    }

    /**
     * InputFormat that encapsulates the correct input format based on the file
     * type.
     * 
     */
    public static class AllLoaderInputFormat extends
            FileInputFormat<Writable, Writable> {

        transient PathPartitionHelper partitionHelper = new PathPartitionHelper();
        String udfSignature;

        public AllLoaderInputFormat(String udfSignature) {
            super();
            this.udfSignature = udfSignature;
        }

        @Override
        protected List<FileStatus> listStatus(JobContext jobContext)
                throws IOException {

            List<FileStatus> files = partitionHelper.listStatus(jobContext,
                    AllLoader.class, udfSignature);

            if (files == null)
                files = super.listStatus(jobContext);

            return files;

        }

        @Override
        public RecordReader<Writable, Writable> createRecordReader(
                InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
                throws IOException, InterruptedException {

            // this method plugs the AllReader into the system, and the
            // AllReader will when called select the correct LoadFunc
            // return new AllReader(udfSignature);
            return new AllReader(udfSignature);
        }

    }

    /**
     * This is where the logic is for selecting the correct Loader.
     * 
     */
    public static class AllReader extends RecordReader<Writable, Writable> {

        LoadFunc selectedLoadFunc;
        RecordReader<Writable, Writable> selectedReader;
        LoadFuncHelper loadFuncHelper = null;
        String udfSignature;
        Path path;

        public AllReader(String udfSignature) {
            this.udfSignature = udfSignature;
        }

        @SuppressWarnings("unchecked")
        @Override
        public void initialize(InputSplit inputSplit,
                TaskAttemptContext taskAttemptContext) throws IOException,
                InterruptedException {

            FileSplit fileSplit = (FileSplit) inputSplit;

            path = fileSplit.getPath();
            String fileName = path.toUri().toString();

            // select the correct load function and initialise
            loadFuncHelper = new LoadFuncHelper(
                    taskAttemptContext.getConfiguration());

            FuncSpec funcSpec = loadFuncHelper.determineFunction(fileName);

            if (funcSpec == null) {
                throw new IOException("Cannot determine LoadFunc for "
                        + fileName);
            }

            selectedLoadFunc = (LoadFunc) PigContext
                    .instantiateFuncFromSpec(funcSpec);

            selectedLoadFunc.setUDFContextSignature(udfSignature);
            selectedLoadFunc.setLocation(fileName,
                    new Job(taskAttemptContext.getConfiguration(),
                            taskAttemptContext.getJobName()));

            selectedReader = selectedLoadFunc.getInputFormat()
                    .createRecordReader(fileSplit, taskAttemptContext);

            selectedReader.initialize(fileSplit, taskAttemptContext);

            LOG.info("Using LoadFunc " + selectedLoadFunc.getClass().getName()
                    + " on " + fileName);

        }

        // ---------------------- all functions below this line delegate work to
        // the selectedReader ------------//

        public LoadFunc prepareLoadFuncForReading(PigSplit split)
                throws IOException {

            selectedLoadFunc.prepareToRead(selectedReader, split);
            return selectedLoadFunc;

        }

        @Override
        public boolean nextKeyValue() throws IOException, InterruptedException {
            return selectedReader.nextKeyValue();
        }

        @Override
        public Writable getCurrentKey() throws IOException,
                InterruptedException {
            return selectedReader.getCurrentKey();
        }

        @Override
        public Writable getCurrentValue() throws IOException,
                InterruptedException {
            return selectedReader.getCurrentValue();
        }

        @Override
        public float getProgress() throws IOException, InterruptedException {
            return selectedReader.getProgress();
        }

        @Override
        public void close() throws IOException {
            selectedReader.close();
        }

    }

}
