src/java/org/apache/hcatalog/mapreduce/HCatBaseOutputFormat.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.hcatalog.mapreduce;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.hive.ql.metadata.Table;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.hadoop.util.ReflectionUtils;
import org.apache.hcatalog.common.ErrorType;
import org.apache.hcatalog.common.HCatConstants;
import org.apache.hcatalog.common.HCatException;
import org.apache.hcatalog.common.HCatUtil;
import org.apache.hcatalog.data.HCatRecord;
import org.apache.hcatalog.data.schema.HCatSchema;

public abstract class HCatBaseOutputFormat extends OutputFormat<WritableComparable<?>, HCatRecord> {

//  static final private Log LOG = LogFactory.getLog(HCatBaseOutputFormat.class);

    /**
     * Gets the table schema for the table specified in the HCatOutputFormat.setOutput call
     * on the specified job context.
     * @param context the context
     * @return the table schema
     * @throws IOException if HCatOutputFromat.setOutput has not been called for the passed context
     */
    public static HCatSchema getTableSchema(JobContext context) throws IOException {
        OutputJobInfo jobInfo = getJobInfo(context);
        return jobInfo.getTableInfo().getDataColumns();
    }

    /**
     * Check for validity of the output-specification for the job.
     * @param context information about the job
     * @throws IOException when output should not be attempted
     */
    @Override
    public void checkOutputSpecs(JobContext context
    ) throws IOException, InterruptedException {
        getOutputFormat(context).checkOutputSpecs(context);
    }

    /**
     * Gets the output format instance.
     * @param context the job context
     * @return the output format instance
     * @throws IOException
     */
    protected OutputFormat<WritableComparable<?>, HCatRecord> getOutputFormat(JobContext context) throws IOException {
        OutputJobInfo jobInfo = getJobInfo(context);
        HCatStorageHandler storageHandler = HCatUtil.getStorageHandler(context.getConfiguration(), jobInfo.getTableInfo().getStorerInfo());
        //why do we need this?
        configureOutputStorageHandler(context);
        return storageHandler.getOutputFormatContainer(ReflectionUtils.newInstance(storageHandler.getOutputFormatClass(), context.getConfiguration()));
    }

    /**
     * Gets the HCatOuputJobInfo object by reading the Configuration and deserializing
     * the string. If InputJobInfo is not present in the configuration, throws an
     * exception since that means HCatOutputFormat.setOutput has not been called.
     * @param jobContext the job context
     * @return the OutputJobInfo object
     * @throws IOException the IO exception
     */
    public static OutputJobInfo getJobInfo(JobContext jobContext) throws IOException {
        String jobString = jobContext.getConfiguration().get(HCatConstants.HCAT_KEY_OUTPUT_INFO);
        if (jobString == null) {
            throw new HCatException(ErrorType.ERROR_NOT_INITIALIZED);
        }

        return (OutputJobInfo) HCatUtil.deserialize(jobString);
    }

    /**
     * Configure the output storage handler
     * @param jobContext the job context
     * @throws IOException
     */
    @SuppressWarnings("unchecked")
    static void configureOutputStorageHandler(
        JobContext jobContext) throws IOException {
        configureOutputStorageHandler(jobContext, (List<String>) null);
    }

    /**
     * Configure the output storage handler with allowing specification of missing dynamic partvals
     * @param jobContext the job context
     * @param dynamicPartVals
     * @throws IOException
     */
    @SuppressWarnings("unchecked")
    static void configureOutputStorageHandler(
        JobContext jobContext, List<String> dynamicPartVals) throws IOException {
        try {
            OutputJobInfo jobInfo = (OutputJobInfo) HCatUtil.deserialize(jobContext.getConfiguration().get(HCatConstants.HCAT_KEY_OUTPUT_INFO));
            HCatStorageHandler storageHandler = HCatUtil.getStorageHandler(jobContext.getConfiguration(), jobInfo.getTableInfo().getStorerInfo());

            Map<String, String> partitionValues = jobInfo.getPartitionValues();
            String location = jobInfo.getLocation();

            if (dynamicPartVals != null) {
                // dynamic part vals specified
                List<String> dynamicPartKeys = jobInfo.getDynamicPartitioningKeys();
                if (dynamicPartVals.size() != dynamicPartKeys.size()) {
                    throw new HCatException(ErrorType.ERROR_INVALID_PARTITION_VALUES,
                        "Unable to configure dynamic partitioning for storage handler, mismatch between"
                            + " number of partition values obtained[" + dynamicPartVals.size()
                            + "] and number of partition values required[" + dynamicPartKeys.size() + "]");
                }
                for (int i = 0; i < dynamicPartKeys.size(); i++) {
                    partitionValues.put(dynamicPartKeys.get(i), dynamicPartVals.get(i));
                }

//            // re-home location, now that we know the rest of the partvals
//            Table table = jobInfo.getTableInfo().getTable();
//
//            List<String> partitionCols = new ArrayList<String>();
//            for(FieldSchema schema : table.getPartitionKeys()) {
//              partitionCols.add(schema.getName());
//            }
                jobInfo.setPartitionValues(partitionValues);
            }

            HCatUtil.configureOutputStorageHandler(storageHandler, jobContext, jobInfo);
        } catch (Exception e) {
            if (e instanceof HCatException) {
                throw (HCatException) e;
            } else {
                throw new HCatException(ErrorType.ERROR_INIT_STORAGE_HANDLER, e);
            }
        }
    }

    /**
     * Configure the output storage handler, with allowing specification
     * of partvals from which it picks the dynamic partvals
     * @param context the job context
     * @param jobInfo the output job info
     * @param fullPartSpec
     * @throws IOException
     */

    protected static void configureOutputStorageHandler(
        JobContext context, OutputJobInfo jobInfo,
        Map<String, String> fullPartSpec) throws IOException {
        List<String> dynamicPartKeys = jobInfo.getDynamicPartitioningKeys();
        if ((dynamicPartKeys == null) || (dynamicPartKeys.isEmpty())) {
            configureOutputStorageHandler(context, (List<String>) null);
        } else {
            List<String> dynKeyVals = new ArrayList<String>();
            for (String dynamicPartKey : dynamicPartKeys) {
                dynKeyVals.add(fullPartSpec.get(dynamicPartKey));
            }
            configureOutputStorageHandler(context, dynKeyVals);
        }
    }


    protected static void setPartDetails(OutputJobInfo jobInfo, final HCatSchema schema,
                                         Map<String, String> partMap) throws HCatException, IOException {
        List<Integer> posOfPartCols = new ArrayList<Integer>();
        List<Integer> posOfDynPartCols = new ArrayList<Integer>();

        // If partition columns occur in data, we want to remove them.
        // So, find out positions of partition columns in schema provided by user.
        // We also need to update the output Schema with these deletions.

        // Note that, output storage handlers never sees partition columns in data
        // or schema.

        HCatSchema schemaWithoutParts = new HCatSchema(schema.getFields());
        for (String partKey : partMap.keySet()) {
            Integer idx;
            if ((idx = schema.getPosition(partKey)) != null) {
                posOfPartCols.add(idx);
                schemaWithoutParts.remove(schema.get(partKey));
            }
        }

        // Also, if dynamic partitioning is being used, we want to
        // set appropriate list of columns for the columns to be dynamically specified.
        // These would be partition keys too, so would also need to be removed from
        // output schema and partcols

        if (jobInfo.isDynamicPartitioningUsed()) {
            for (String partKey : jobInfo.getDynamicPartitioningKeys()) {
                Integer idx;
                if ((idx = schema.getPosition(partKey)) != null) {
                    posOfPartCols.add(idx);
                    posOfDynPartCols.add(idx);
                    schemaWithoutParts.remove(schema.get(partKey));
                }
            }
        }

        HCatUtil.validatePartitionSchema(
            new Table(jobInfo.getTableInfo().getTable()), schemaWithoutParts);
        jobInfo.setPosOfPartCols(posOfPartCols);
        jobInfo.setPosOfDynPartCols(posOfDynPartCols);
        jobInfo.setOutputSchema(schemaWithoutParts);
    }
}