Pig 重写加载函数和存储函数UDF

pig自带的pigstorage不能指定行分隔符，所以自己重写了一个简单的UDF类，可以指定列和行的分隔符，之前研究过的简单的，

http://blog.csdn.net/ruishenh/article/details/12048067

但是弊端大，所以这次重写一下。

操作步骤打好包上传到服务器，

grunt> register /home/pig/pig-0.11.0/udflib/myStorage.jar

grunt> cat student;
1,xiaohouzi,25/2,xiaohouzi2,24/3,xiaohouzi3,23

grunt> a = load 'student' using com.hcr.hadoop.pig.MyStorage(',','/');

grunt> dump a;

(1,xiaohouzi,25)
(2,xiaohouzi2,24)
(3,xiaohouzi3,23)

grunt> store a into 'myStorageOut' using com.hcr.hadoop.pig.MyStorage(',','/');

执行提示成功后查看

grunt> cat myStorageOut
1,xiaohouzi,25/2,xiaohouzi2,24/3,xiaohouzi3,23/

源码类

package com.hcr.hadoop.pig;

import java.io.DataOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.util.ArrayList;
import java.util.List;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.io.compress.BZip2Codec;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.GzipCodec;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.LineRecordReader;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.lib.output.TextOutputFormat;
import org.apache.hadoop.util.ReflectionUtils;
import org.apache.pig.Expression;
import org.apache.pig.LoadFunc;
import org.apache.pig.LoadMetadata;
import org.apache.pig.PigException;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceSchema.ResourceFieldSchema;
import org.apache.pig.ResourceStatistics;
import org.apache.pig.StoreFunc;
import org.apache.pig.StoreFuncInterface;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigSplit;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema;
import org.apache.pig.impl.util.StorageUtil;

public class MyStorage extends LoadFunc implements StoreFuncInterface,LoadMetadata {

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

	private static final String utf8 = "UTF-8";

	private static String fieldDel = "\t";

	private static String recordDel = "\n";

	protected RecordReader recordReader = null;

	protected RecordWriter writer = null;

	public MyStorage() {
	}

	public MyStorage(String fieldDel) {
		this(fieldDel, "\n");
	}

	public MyStorage(String fieldDel, String recordDel) {
		this.fieldDel = fieldDel;
		this.recordDel = recordDel;
	}

	@Override
	public void setLocation(String s, Job job) throws IOException {
		FileInputFormat.setInputPaths(job, s);
	}

	@Override
	public InputFormat getInputFormat() throws IOException {
		return new MyStorageInputFormat(recordDel);
	}

	@Override
	public void prepareToRead(RecordReader recordReader, PigSplit pigSplit)
			throws IOException {
		this.recordReader = recordReader;
	}

	@Override
	public Tuple getNext() throws IOException {
		try {
			boolean flag = recordReader.nextKeyValue();
			if (!flag) {
				return null;
			}
			Text value = (Text) recordReader.getCurrentValue();
			String[] strArray = value.toString().split(fieldDel);
			List lst = new ArrayList<String>();
			int i = 0;
			for (String singleItem : strArray) {
				lst.add(i++, singleItem);
			}
			return TupleFactory.getInstance().newTuple(lst);
		} catch (InterruptedException e) {
			throw new ExecException("Read data error",
					PigException.REMOTE_ENVIRONMENT, e);
		}
	}

	/**
	 * */
	@Override
	public String relToAbsPathForStoreLocation(String location, Path curDir)
			throws IOException {
		return LoadFunc.getAbsolutePath(location, curDir);
	}

	@Override
	public OutputFormat getOutputFormat() throws IOException {
		return new MyStorageOutputFormat(StorageUtil.parseFieldDel(fieldDel),
				this.recordDel);
	}

	@Override
	public void setStoreLocation(String location, Job job) throws IOException {
		job.getConfiguration().set("mapred.textoutputformat.separator", "");
		FileOutputFormat.setOutputPath(job, new Path(location));
		if ("true".equals(job.getConfiguration().get(
				"output.compression.enabled"))) {
			FileOutputFormat.setCompressOutput(job, true);
			String codec = job.getConfiguration().get(
					"output.compression.codec");
			try {
				FileOutputFormat.setOutputCompressorClass(job,
						(Class<? extends CompressionCodec>) Class
								.forName(codec));
			} catch (ClassNotFoundException e) {
				throw new RuntimeException("Class not found: " + codec);
			}
		} else {
			// This makes it so that storing to a directory ending with ".gz" or
			// ".bz2" works.
			setCompression(new Path(location), job);
		}

	}

	private void setCompression(Path path, Job job) {
		String location = path.getName();
		if (location.endsWith(".bz2") || location.endsWith(".bz")) {
			FileOutputFormat.setCompressOutput(job, true);
			FileOutputFormat.setOutputCompressorClass(job, BZip2Codec.class);
		} else if (location.endsWith(".gz")) {
			FileOutputFormat.setCompressOutput(job, true);
			FileOutputFormat.setOutputCompressorClass(job, GzipCodec.class);
		} else {
			FileOutputFormat.setCompressOutput(job, false);
		}
	}

	@Override
	public void checkSchema(ResourceSchema s) throws IOException {
		// TODO Auto-generated method stub

	}

	@Override
	public void prepareToWrite(RecordWriter writer) throws IOException {
		this.writer = writer;
	}

	@Override
	public void putNext(Tuple t) throws IOException {
		try {
			writer.write(null, t);
		} catch (InterruptedException e) {
			throw new IOException(e);
		}
	}

	@Override
	public void setStoreFuncUDFContextSignature(String signature) {
		// TODO Auto-generated method stub

	}

	@Override
	public void cleanupOnFailure(String location, Job job) throws IOException {
		StoreFunc.cleanupOnFailureImpl(location, job);
	}

	@Override
	public void cleanupOnSuccess(String location, Job job) throws IOException {
		// TODO Auto-generated method stub

	}

	@Override
	public ResourceSchema getSchema(String location, Job job)
			throws IOException {
			ResourceSchema rs=new ResourceSchema();
            FieldSchema c1 = new FieldSchema("c1", DataType.INTEGER);  
            FieldSchema c2 = new FieldSchema("c2", DataType.INTEGER);
            FieldSchema c3 = new FieldSchema("c3", DataType.DOUBLE);
            ResourceFieldSchema fs1 =new ResourceFieldSchema(c1);
            ResourceFieldSchema fs2 =new ResourceFieldSchema(c2);
            ResourceFieldSchema fs3 =new ResourceFieldSchema(c3);
            rs.setFields(new ResourceFieldSchema[]{fs1,fs2,fs3});
            return rs;
	}

	@Override
	public ResourceStatistics getStatistics(String location, Job job)
			throws IOException {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public String[] getPartitionKeys(String location, Job job)
			throws IOException {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public void setPartitionFilter(Expression partitionFilter)
			throws IOException {
		// TODO Auto-generated method stub
		
	}
}

class MyStorageInputFormat extends TextInputFormat {

	private final String recordDel;

	public MyStorageInputFormat(String recordDel) {
		this.recordDel = recordDel;
	}

	@Override
	public RecordReader<LongWritable, Text> createRecordReader(
			InputSplit split, TaskAttemptContext context) {
		String delimiter = context.getConfiguration().get(
				"textinputformat.record.delimiter");
		if (recordDel != null) {
			delimiter = recordDel;
		}
		byte[] recordDelimiterBytes = null;
		if (null != delimiter){
			try {
				recordDelimiterBytes = decode(delimiter).getBytes("UTF-8");
			} catch (UnsupportedEncodingException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
		return new LineRecordReader(recordDelimiterBytes);
	}
	/**
	 * 工作流传过来的列分隔符,有可能是特殊字符,用八进制或者十六进制表示
	 * @throws IOException 
	 */
	public static String decode(String str) throws IOException {
		String re = str;
		if (str != null && str.startsWith("\\")) {
			str = str.substring(1, str.length());
			String[] chars = str.split("\\\\");
			byte[] bytes = new byte[chars.length];
			for (int i = 0; i < chars.length; i++) {
				if (chars[i].equals("t")) {
					bytes[i] = 9;
				} else if (chars[i].equals("r")) {
					bytes[i] = 13;
				} else if (chars[i].equals("n")) {
					bytes[i] = 10;
				} else if (chars[i].equals("b")) {
					bytes[i] = 8;
				} else {
					bytes[i] = Byte.decode(chars[i]);
				}
			}
			try {
				re = new String(bytes, "UTF-8");
			} catch (UnsupportedEncodingException e) {
				throw new IOException(str, e);
			}
		}
		return re;
	}

}

class MyStorageOutputFormat extends TextOutputFormat<WritableComparable, Tuple> {

	private final byte fieldDel;

	private final String recordDel;

	public MyStorageOutputFormat(byte delimiter) {
		this(delimiter, "\n");
	}

	public MyStorageOutputFormat(byte delimiter, String recordDel) {
		this.fieldDel = delimiter;
		this.recordDel = recordDel;
	}

	protected static class MyRecordWriter extends
			TextOutputFormat.LineRecordWriter<WritableComparable, Tuple> {

		private static byte[] newline;

		private final byte fieldDel;

		public MyRecordWriter(DataOutputStream out, byte fieldDel)
				throws UnsupportedEncodingException {
			this(out, fieldDel, "\n".getBytes("UTF-8"));
		}

		public MyRecordWriter(DataOutputStream out, byte fieldDel, byte[] record) {
			super(out);
			this.fieldDel = fieldDel;
			this.newline = record;
		}

		public synchronized void write(WritableComparable key, Tuple value)
				throws IOException {
			int sz = value.size();
			for (int i = 0; i < sz; i++) {
				StorageUtil.putField(out, value.get(i));
				if (i != sz - 1) {
					out.writeByte(fieldDel);
				}
			}
			out.write(newline);
		}
	}

	@Override
	public RecordWriter<WritableComparable, Tuple> getRecordWriter(
			TaskAttemptContext job) throws IOException, InterruptedException {
		Configuration conf = job.getConfiguration();
		boolean isCompressed = getCompressOutput(job);
		CompressionCodec codec = null;
		String extension = "";
		if (isCompressed) {
			Class<? extends CompressionCodec> codecClass = getOutputCompressorClass(
					job, GzipCodec.class);
			codec = (CompressionCodec) ReflectionUtils.newInstance(codecClass,
					conf);
			extension = codec.getDefaultExtension();
		}
		Path file = getDefaultWorkFile(job, extension);
		FileSystem fs = file.getFileSystem(conf);
		if (!isCompressed) {
			FSDataOutputStream fileOut = fs.create(file, false);
			return new MyRecordWriter(fileOut, fieldDel,
					this.recordDel.getBytes());
		} else {
			FSDataOutputStream fileOut = fs.create(file, false);
			return new MyRecordWriter(new DataOutputStream(
					codec.createOutputStream(fileOut)), fieldDel,
					this.recordDel.getBytes());
		}
	}

}

本人试验 \001 和 \002同样可以

grunt> register  /home/pig/pig-0.11.0/udflib/myStorage.jar 
grunt> cat X;
keyDataKNZKCZY:ZDKJS:616150:AFS:3842708d_20131219194420-642464756keyDataKNZKCZY:ZDKJS:616614:AFS:3843920d_20131219194420-642464756keyDataKNZKCZY:ZDKJS:616661:AFS:3844040d_20131219194420-642464756
grunt> a = load 'X' using com.hcr.hadoop.pig.MyStorage('\\001','\\002');
grunt> dump a;
(keyData,KNZKCZY:ZDKJS:616150:AFS:3842708,d_20131219194420-642464756)
(keyData,KNZKCZY:ZDKJS:616614:AFS:3843920,d_20131219194420-642464756)
(keyData,KNZKCZY:ZDKJS:616661:AFS:3844040,d_20131219194420-642464756)
grunt> 

有的时候如果加载模式不想指定具体模式（比如太多了字段，或者不够公有化）就想使用已存在的模式

实现LoadMetadata接口，然后

重写

@Override
public ResourceSchema getSchema(String location, Job job)
throws IOException {
ResourceSchema rs=new ResourceSchema();
FieldSchema c1 = new FieldSchema("c1", DataType.INTEGER);
FieldSchema c2 = new FieldSchema("c2", DataType.INTEGER);
FieldSchema c3 = new FieldSchema("c3", DataType.DOUBLE);
ResourceFieldSchema fs1 =new ResourceFieldSchema(c1);
ResourceFieldSchema fs2 =new ResourceFieldSchema(c2);
ResourceFieldSchema fs3 =new ResourceFieldSchema(c3);
rs.setFields(new ResourceFieldSchema[]{fs1,fs2,fs3});
return rs;

}

这一个简单的例子中就返回了直接使用模式的形式

grunt> register /home/pig/pig-0.11.0/udflib/myStorage.jar
grunt> a = load 'student' using com.hcr.hadoop.pig.MyStorage(',','/');
grunt> describe a;
a: {c1: int,c2: int,c3: double}
grunt> b = foreach a generate c1,c2,c3;
grunt> describe b;
b: {c1: int,c2: int,c3: double}