package pg.hadoop.pig.piggybank;

import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.piggybank.storage.MultiStorage;

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

/**
 * Extended MultiStorage class that allows to remove the key field from the output files.
 */
public class ExtendedMultiStorage extends MultiStorage {


    protected static final String DEFAULT_COMPRESSION = "none";
    protected static final String DEFAULT_DELIMITER = "\\t";

    protected final int keyIndex;
    protected final boolean includeFieldIndex;

    public ExtendedMultiStorage(String parentPathStr, String splitFieldIndex, String includeFieldIndex) {
        this(parentPathStr, splitFieldIndex, DEFAULT_COMPRESSION, includeFieldIndex);
    }

    public ExtendedMultiStorage(String parentPathStr, String splitFieldIndex, String compression, String includeFieldIndex) {
        this(parentPathStr, splitFieldIndex, compression, DEFAULT_DELIMITER, includeFieldIndex);
    }

    public ExtendedMultiStorage(String parentPathStr, String splitFieldIndex, String compression, String fieldDel, String includeFieldIndex) {
        super(parentPathStr, splitFieldIndex, compression, fieldDel);
        this.keyIndex = Integer.parseInt(splitFieldIndex);
        this.includeFieldIndex = Boolean.parseBoolean(includeFieldIndex);
    }

    @Override
    public void prepareToWrite(RecordWriter writer) throws IOException {
        super.prepareToWrite(includeFieldIndex ? writer : new RecordWriterWrapper(writer));
    }

    /**
     * Simple wrapper that intercepts the write call and remove the unnecessary field from the tuple before the write.
     */
    private class RecordWriterWrapper extends RecordWriter<String, Tuple> {

        private RecordWriter<String, Tuple> originalWriter;

        private RecordWriterWrapper(final RecordWriter<String, Tuple> originalWriter) {
            this.originalWriter = originalWriter;
        }

        @Override
        public void write(String key, Tuple value) throws IOException, InterruptedException {
            final List<Object> fields = value.getAll();
            fields.remove(keyIndex);
            final Tuple newTuple = TupleFactory.getInstance().newTupleNoCopy(fields);
            originalWriter.write(key, newTuple);
        }

        @Override
        public void close(TaskAttemptContext context) throws IOException, InterruptedException {
            originalWriter.close(context);
        }
    }
}

@@ -310,13 +310,73 @@
 til.
+Array
 List;
+%0Aimport java.util.List;%0Aimport java.util.regex.Pattern;
 %0A%0A/*
@@ -427,33 +427,36 @@
  to 
-remov
+choos
 e the 
-key 
 field
- from
+s to keep on
  the
@@ -467,17 +467,16 @@
 put file
-s
 .%0A */%0Apu
@@ -531,16 +531,81 @@
 rage %7B%0A%0A
+    private static Pattern COMMA_PATTERN = Pattern.compile(%22,%22);%0A
 %0A    pro
@@ -741,73 +741,29 @@
 ted 
-final int keyIndex;%0A    protected final boolean includeFieldIndex
+int%5B%5D indexesToOutput
 ;%0A%0A 
@@ -844,39 +844,37 @@
 x, String in
-cludeFieldIndex
+dexesToOutput
 ) %7B%0A        
@@ -929,39 +929,37 @@
 PRESSION, in
-cludeFieldIndex
+dexesToOutput
 );%0A    %7D%0A%0A  
@@ -1059,39 +1059,37 @@
 n, String in
-cludeFieldIndex
+dexesToOutput
 ) %7B%0A        
@@ -1159,31 +1159,29 @@
 ITER, in
-cludeFieldIndex
+dexesToOutput
 );%0A    %7D
@@ -1306,31 +1306,29 @@
 tring in
-cludeFieldIndex
+dexesToOutput
 ) %7B%0A    
@@ -1405,132 +1405,422 @@
     
-this.keyIndex = Integer.parseInt(splitFieldIndex);%0A        this.includeFieldIndex = Boolean.parseBoolean(includeFieldIndex);
+if (indexesToOutput != null && !indexesToOutput.isEmpty()) %7B%0A            final String%5B%5D indexes = COMMA_PATTERN.split(indexesToOutput);%0A            if (indexes.length %3E 0) %7B%0A                this.indexesToOutput = new int%5Bindexes.length%5D;%0A                for (int i = 0; i %3C indexes.length; i++) %7B%0A                    this.indexesToOutput%5Bi%5D = Integer.parseInt(indexes%5Bi%5D);%0A                %7D%0A            %7D%0A        %7D
 %0A   
@@ -1945,23 +1945,29 @@
 e(in
-cludeFieldIndex
+dexesToOutput == null
  ? w
@@ -2593,24 +2593,82 @@
 s = 
-value.getAll();%0A
+new ArrayList%3C%3E();%0A            for (int indexToKeep : indexesToOutput) %7B%0A 
     
@@ -2679,32 +2679,61 @@
     
+   
 fields.
-remove(keyIndex);
+add(value.get(indexToKeep));%0A            %7D%0A
 %0A   

Change the behavior to choose the fields to keep on the output.