<commit_before>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);
        }
    }
}
<commit_msg>Change the behavior to choose the fields to keep on the output.<commit_after>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.ArrayList;
import java.util.List;
import java.util.regex.Pattern;

/**
 * Extended MultiStorage class that allows to choose the fields to keep on the output file.
 */
public class ExtendedMultiStorage extends MultiStorage {

    private static Pattern COMMA_PATTERN = Pattern.compile(",");

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

    protected int[] indexesToOutput;

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

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

    public ExtendedMultiStorage(String parentPathStr, String splitFieldIndex, String compression, String fieldDel, String indexesToOutput) {
        super(parentPathStr, splitFieldIndex, compression, fieldDel);
        if (indexesToOutput != null && !indexesToOutput.isEmpty()) {
            final String[] indexes = COMMA_PATTERN.split(indexesToOutput);
            if (indexes.length > 0) {
                this.indexesToOutput = new int[indexes.length];
                for (int i = 0; i < indexes.length; i++) {
                    this.indexesToOutput[i] = Integer.parseInt(indexes[i]);
                }
            }
        }
    }

    @Override
    public void prepareToWrite(RecordWriter writer) throws IOException {
        super.prepareToWrite(indexesToOutput == null ? 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 = new ArrayList<>();
            for (int indexToKeep : indexesToOutput) {
                fields.add(value.get(indexToKeep));
            }

            final Tuple newTuple = TupleFactory.getInstance().newTupleNoCopy(fields);
            originalWriter.write(key, newTuple);
        }

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

{'fraction_non_alphanumeric': 0.0784177654406662, 'fraction_numerical': 0.00034698126301179735, 'mean_word_length': 3.500390320062451, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '827953', 'n_tokens_mistral': 1480, 'n_tokens_neox': 1421, 'n_words': 391}