<commit_before>package org.opencb.opencga.storage.hadoop.variant.stats.me;

import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
import org.apache.hadoop.hbase.mapreduce.TableInputFormat;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.mapreduce.Job;
import org.apache.phoenix.schema.types.PArrayDataType;
import org.apache.phoenix.schema.types.PVarchar;
import org.opencb.biodata.models.core.Region;
import org.opencb.biodata.models.feature.Genotype;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.tools.pedigree.MendelianError;
import org.opencb.opencga.storage.core.exceptions.StorageEngineException;
import org.opencb.opencga.storage.core.metadata.VariantStorageMetadataManager;
import org.opencb.opencga.storage.core.metadata.models.TaskMetadata;
import org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam;
import org.opencb.opencga.storage.hadoop.variant.AbstractVariantsTableDriver;
import org.opencb.opencga.storage.hadoop.variant.GenomeHelper;
import org.opencb.opencga.storage.hadoop.variant.adaptors.VariantHBaseQueryParser;
import org.opencb.opencga.storage.hadoop.variant.adaptors.phoenix.VariantPhoenixHelper;
import org.opencb.opencga.storage.hadoop.variant.adaptors.phoenix.VariantPhoenixKeyFactory;
import org.opencb.opencga.storage.hadoop.variant.index.annotation.mr.VariantTableSampleIndexOrderMapper;
import org.opencb.opencga.storage.hadoop.variant.index.sample.HBaseToSampleIndexConverter;
import org.opencb.opencga.storage.hadoop.variant.index.sample.SampleIndexDBLoader;
import org.opencb.opencga.storage.hadoop.variant.index.sample.SampleIndexToHBaseConverter;
import org.opencb.opencga.storage.hadoop.variant.mr.VariantAlignedInputFormat;
import org.opencb.opencga.storage.hadoop.variant.mr.VariantMapReduceUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Created on 04/03/19.
 *
 * @author <NAME> &lt;<EMAIL>&gt;
 */
public class MendelianErrorDriver extends AbstractVariantsTableDriver {

    private static final Logger LOGGER = LoggerFactory.getLogger(MendelianErrorDriver.class);

    public static final String TRIOS = "trios";
    public static final String TRIOS_LIST = "MendelianErrorDriver.trios_list";
    private List<Integer> sampleIds;
    private boolean partial;
    private String region;

    @Override
    protected Class<MendelianErrorMapper> getMapperClass() {
        return MendelianErrorMapper.class;
    }

    @Override
    protected Map<String, String> getParams() {
        Map<String, String> params = new HashMap<>();
        params.put("--" + TRIOS, "(father,mother,child;)**");
        params.put("--" + VariantQueryParam.REGION.key(), "<region>");
        return params;
    }

    @Override
    protected void parseAndValidateParameters() throws IOException {
        super.parseAndValidateParameters();


        VariantStorageMetadataManager metadataManager = getMetadataManager();
        String trios = getParam(TRIOS);
        if (StringUtils.isNotEmpty(trios)) {
            sampleIds = new LinkedList<>();
            List<Integer> trioList = new ArrayList<>(3);
            for (String trio : trios.split(";")) {
                for (String sample : trio.split(",")) {
                    Integer sampleId;
                    if (sample.equals("-")) {
                        sampleId = -1;
                    } else {
                        sampleId = metadataManager.getSampleId(getStudyId(), sample);
                        if (sampleId == null) {
                            throw new IllegalArgumentException("Sample '" + sample + "' not found.");
                        }
                        trioList.add(sampleId);
                    }
                    sampleIds.add(sampleId);
                }
                if (trioList.size() != 3) {
                    throw new IllegalArgumentException("Found trio with " + trioList.size() + " members, instead of 3: " + trioList);
                }
                LOGGER.info("Trio: " + trio + " -> " + trioList);
                trioList.clear();
            }
        } else {
            throw new IllegalArgumentException("Missing list of trios!");
        }

        if (sampleIds.isEmpty()) {
            throw new IllegalArgumentException("Missing list of trios!");
        }

        if (sampleIds.size() % 3 != 0) {
            throw new IllegalArgumentException("Wrong number of samples in trios!");
        }

        region = getParam(VariantQueryParam.REGION.key(), "");

        partial = StringUtils.isNotEmpty(region);
    }

    @Override
    protected Job setupJob(Job job, String archiveTable, String variantTable) throws IOException {

        Scan scan = new Scan();

        if (StringUtils.isNotEmpty(region)) {
            LOGGER.info("Calculate Mendelian Errors from region " + region);
            VariantHBaseQueryParser.addRegionFilter(scan, new Region(region));
        }
        LOGGER.info("Calculate Mendelian Errors for " + (sampleIds.size() / 3) + " trios");


        job.getConfiguration().set(TRIOS_LIST, sampleIds.stream().map(Objects::toString).collect(Collectors.joining(",")));

        for (Integer sampleId : sampleIds) {
            scan.addColumn(getHelper().getColumnFamily(), VariantPhoenixHelper.buildSampleColumnKey(getStudyId(), sampleId));
        }
//        scan.addColumn(getHelper().getColumnFamily(), VariantPhoenixHelper.VariantColumn.FULL_ANNOTATION.bytes());


        VariantMapReduceUtil.configureMapReduceScan(scan, getConf());

        VariantMapReduceUtil.initTableMapperJob(job, variantTable,
                scan, getMapperClass(), VariantAlignedInputFormat.class);
        VariantAlignedInputFormat.setDelegatedInputFormat(job, TableInputFormat.class);
        VariantAlignedInputFormat.setBatchSize(job, SampleIndexDBLoader.BATCH_SIZE);

        VariantMapReduceUtil.setOutputHBaseTable(job, getTableNameGenerator().getSampleIndexTableName(getStudyId()));

        VariantMapReduceUtil.setNoneReduce(job);
        return job;
    }

    @Override
    protected void postExecution(boolean succeed) throws IOException, StorageEngineException {
        super.postExecution(succeed);
        if (succeed && !partial) {
            VariantStorageMetadataManager metadataManager = getMetadataManager();
            for (int i = 0; i < sampleIds.size(); i += 3) {
                Integer child = sampleIds.get(i + 2);
                metadataManager.updateSampleMetadata(getStudyId(), child, sampleMetadata -> {
                    sampleMetadata.setMendelianErrorStatus(TaskMetadata.Status.READY);
                    return sampleMetadata;
                });
            }

        }
    }

    @Override
    protected String getJobOperationName() {
        return "calculate_mendelian_errors";
    }

    public static class MendelianErrorMapper extends VariantTableSampleIndexOrderMapper<ImmutableBytesWritable, Put> {

        private Map<Integer, ByteArrayOutputStream> mendelianErrorsMap = new HashMap<>();
        private Map<Integer, Map<String, Integer>> genotypeCount = new HashMap<>();
        private List<List<Integer>> trios;
        private byte[] family;
        private SampleIndexToHBaseConverter converter;

        @Override
        protected void setup(Context context) throws IOException, InterruptedException {
            family = new GenomeHelper(context.getConfiguration()).getColumnFamily();
            converter = new SampleIndexToHBaseConverter(family);

            int[] sampleIds = context.getConfiguration().getInts(TRIOS_LIST);
            trios = new ArrayList<>(sampleIds.length / 3);

            for (int i = 0; i < sampleIds.length; i += 3) {
                trios.add(Arrays.asList(
                        sampleIds[i],
                        sampleIds[i + 1],
                        sampleIds[i + 2]));
            }

            for (List<Integer> trio : trios) {
                Integer child = trio.get(2);
                mendelianErrorsMap.put(child, new ByteArrayOutputStream());
                genotypeCount.put(child, new HashMap<>());
            }

        }

        @Override
        protected void map(ImmutableBytesWritable key, Result value, Context context) throws IOException, InterruptedException {
            Variant variant = VariantPhoenixKeyFactory.extractVariantFromVariantRowKey(value.getRow());
            String chromosome = variant.getChromosome();

            Map<Integer, String> gtMap = new HashMap<>();
            for (Cell cell : value.rawCells()) {
                Integer sampleId = VariantPhoenixHelper
                        .extractSampleId(cell.getQualifierArray(), cell.getQualifierOffset(), cell.getQualifierLength());
                if (sampleId != null) {
                    ImmutableBytesWritable ptr = new ImmutableBytesWritable(
                            cell.getValueArray(),
                            cell.getValueOffset(),
                            cell.getValueLength());
                    PArrayDataType.positionAtArrayElement(ptr, 0, PVarchar.INSTANCE, null);
                    String gt = Bytes.toString(ptr.get(), ptr.getOffset(), ptr.getLength());
                    gtMap.put(sampleId, gt);
                }
            }


            for (List<Integer> trio : trios) {
                Integer father = trio.get(0);
                Integer mother = trio.get(1);
                Integer child = trio.get(2);

                String fatherGtStr = gtMap.get(father);
                String motherGtStr = gtMap.get(mother);
                String childGtStr = gtMap.get(child);

                int idx = genotypeCount.get(child).merge(childGtStr, 1, Integer::sum) - 1;
                if ((fatherGtStr != null || father == -1) && (motherGtStr != null || mother == -1) && childGtStr != null) {
                    Genotype fatherGt;
                    Genotype motherGt;
                    Genotype childGt;
                    try {
                        fatherGt = fatherGtStr == null ? null : new Genotype(fatherGtStr);
                        motherGt = motherGtStr == null ? null : new Genotype(motherGtStr);
                        childGt = new Genotype(childGtStr);
                    } catch (IllegalArgumentException e) {
                        // Skip malformed GT
                        context.getCounter(COUNTER_GROUP_NAME, "wrong_gt").increment(1);
                        continue;
                    }
                    Integer me = MendelianError.compute(fatherGt, motherGt, childGt, chromosome);
                    context.getCounter(COUNTER_GROUP_NAME, "me_" + me).increment(1);
                    if (me > 0) {
                        ByteArrayOutputStream stream = mendelianErrorsMap.get(child);
                        converter.serializeMendelianError(stream, variant, childGtStr, idx, me);
                    }
                }
            }

        }

        @Override
        public void flush(Context context, String chromosome, int position) throws IOException, InterruptedException {

            for (Map.Entry<Integer, ByteArrayOutputStream> entry : mendelianErrorsMap.entrySet()) {
                Integer sampleId = entry.getKey();
                ByteArrayOutputStream value = entry.getValue();


                byte[] row = HBaseToSampleIndexConverter.toRowKey(sampleId, chromosome, position);

                if (value.size() > 0) {
                    Put put = new Put(row);
                    // Copy value, as the ByteArrayOutputStream is erased
                    put.addColumn(family, HBaseToSampleIndexConverter.toMendelianErrorColumn(), value.toByteArray());

                    context.write(new ImmutableBytesWritable(row), put);
                    value.reset();
                }
            }

            genotypeCount.values().forEach(Map::clear);
        }
    }

    public static void main(String[] args) throws Exception {
        try {
            System.exit(new MendelianErrorDriver().privateMain(args, null));
        } catch (Exception e) {
            LOGGER.error("Error executing " + MendelianErrorDriver.class, e);
            System.exit(1);
        }
    }
}
<commit_msg>storage: Allow read trios to calculate from file #1206<commit_after>package org.opencb.opencga.storage.hadoop.variant.stats.me;

import org.apache.commons.io.FileUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
import org.apache.hadoop.hbase.mapreduce.TableInputFormat;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.mapreduce.Job;
import org.apache.phoenix.schema.types.PArrayDataType;
import org.apache.phoenix.schema.types.PVarchar;
import org.opencb.biodata.models.core.Region;
import org.opencb.biodata.models.feature.Genotype;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.tools.pedigree.MendelianError;
import org.opencb.opencga.storage.core.exceptions.StorageEngineException;
import org.opencb.opencga.storage.core.metadata.VariantStorageMetadataManager;
import org.opencb.opencga.storage.core.metadata.models.TaskMetadata;
import org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam;
import org.opencb.opencga.storage.hadoop.variant.AbstractVariantsTableDriver;
import org.opencb.opencga.storage.hadoop.variant.GenomeHelper;
import org.opencb.opencga.storage.hadoop.variant.adaptors.VariantHBaseQueryParser;
import org.opencb.opencga.storage.hadoop.variant.adaptors.phoenix.VariantPhoenixHelper;
import org.opencb.opencga.storage.hadoop.variant.adaptors.phoenix.VariantPhoenixKeyFactory;
import org.opencb.opencga.storage.hadoop.variant.index.annotation.mr.VariantTableSampleIndexOrderMapper;
import org.opencb.opencga.storage.hadoop.variant.index.sample.HBaseToSampleIndexConverter;
import org.opencb.opencga.storage.hadoop.variant.index.sample.SampleIndexDBLoader;
import org.opencb.opencga.storage.hadoop.variant.index.sample.SampleIndexToHBaseConverter;
import org.opencb.opencga.storage.hadoop.variant.mr.VariantAlignedInputFormat;
import org.opencb.opencga.storage.hadoop.variant.mr.VariantMapReduceUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Created on 04/03/19.
 *
 * @author <NAME> &lt;<EMAIL>&gt;
 */
public class MendelianErrorDriver extends AbstractVariantsTableDriver {

    private static final Logger LOGGER = LoggerFactory.getLogger(MendelianErrorDriver.class);

    public static final String TRIOS = "trios";
    public static final String TRIOS_FILE = "triosFile";
    public static final String TRIOS_LIST = "MendelianErrorDriver.trios_list";
    private List<Integer> sampleIds;
    private boolean partial;
    private String region;

    @Override
    protected Class<MendelianErrorMapper> getMapperClass() {
        return MendelianErrorMapper.class;
    }

    @Override
    protected Map<String, String> getParams() {
        Map<String, String> params = new HashMap<>();
        params.put("--" + TRIOS, "(father,mother,child;)**");
        params.put("--" + VariantQueryParam.REGION.key(), "<region>");
        return params;
    }

    @Override
    protected void parseAndValidateParameters() throws IOException {
        super.parseAndValidateParameters();


        VariantStorageMetadataManager metadataManager = getMetadataManager();

        String triosFile = getParam(TRIOS_FILE);
        String trios;
        if (StringUtils.isNotEmpty(triosFile)) {
            trios = FileUtils.readFileToString(new File(triosFile));
            trios = trios.replaceAll("\n", ";");
        } else {
            trios = getParam(TRIOS);
        }

        if (StringUtils.isNotEmpty(trios)) {
            sampleIds = new LinkedList<>();
            List<Integer> trioList = new ArrayList<>(3);
            for (String trio : trios.split(";")) {
                for (String sample : trio.split(",")) {
                    Integer sampleId;
                    if (sample.equals("-")) {
                        sampleId = -1;
                    } else {
                        sampleId = metadataManager.getSampleId(getStudyId(), sample);
                        if (sampleId == null) {
                            throw new IllegalArgumentException("Sample '" + sample + "' not found.");
                        }
                    }
                    trioList.add(sampleId);
                    sampleIds.add(sampleId);
                }
                if (trioList.size() != 3) {
                    throw new IllegalArgumentException("Found trio with " + trioList.size() + " members, instead of 3: " + trioList);
                }
                LOGGER.info("Trio: " + trio + " -> " + trioList);
                trioList.clear();
            }
        } else {
            throw new IllegalArgumentException("Missing list of trios!");
        }

        if (sampleIds.isEmpty()) {
            throw new IllegalArgumentException("Missing list of trios!");
        }

        if (sampleIds.size() % 3 != 0) {
            throw new IllegalArgumentException("Wrong number of samples in trios!");
        }

        region = getParam(VariantQueryParam.REGION.key(), "");

        partial = StringUtils.isNotEmpty(region);
    }

    @Override
    protected Job setupJob(Job job, String archiveTable, String variantTable) throws IOException {

        Scan scan = new Scan();

        if (StringUtils.isNotEmpty(region)) {
            LOGGER.info("Calculate Mendelian Errors from region " + region);
            VariantHBaseQueryParser.addRegionFilter(scan, new Region(region));
        }
        LOGGER.info("Calculate Mendelian Errors for " + (sampleIds.size() / 3) + " trios");


        job.getConfiguration().set(TRIOS_LIST, sampleIds.stream().map(Objects::toString).collect(Collectors.joining(",")));

        for (Integer sampleId : sampleIds) {
            scan.addColumn(getHelper().getColumnFamily(), VariantPhoenixHelper.buildSampleColumnKey(getStudyId(), sampleId));
        }
//        scan.addColumn(getHelper().getColumnFamily(), VariantPhoenixHelper.VariantColumn.FULL_ANNOTATION.bytes());


        VariantMapReduceUtil.configureMapReduceScan(scan, getConf());

        VariantMapReduceUtil.initTableMapperJob(job, variantTable,
                scan, getMapperClass(), VariantAlignedInputFormat.class);
        VariantAlignedInputFormat.setDelegatedInputFormat(job, TableInputFormat.class);
        VariantAlignedInputFormat.setBatchSize(job, SampleIndexDBLoader.BATCH_SIZE);

        VariantMapReduceUtil.setOutputHBaseTable(job, getTableNameGenerator().getSampleIndexTableName(getStudyId()));

        VariantMapReduceUtil.setNoneReduce(job);
        return job;
    }

    @Override
    protected void postExecution(boolean succeed) throws IOException, StorageEngineException {
        super.postExecution(succeed);
        if (succeed && !partial) {
            VariantStorageMetadataManager metadataManager = getMetadataManager();
            for (int i = 0; i < sampleIds.size(); i += 3) {
                Integer child = sampleIds.get(i + 2);
                metadataManager.updateSampleMetadata(getStudyId(), child, sampleMetadata -> {
                    sampleMetadata.setMendelianErrorStatus(TaskMetadata.Status.READY);
                    return sampleMetadata;
                });
            }

        }
    }

    @Override
    protected String getJobOperationName() {
        return "calculate_mendelian_errors";
    }

    public static class MendelianErrorMapper extends VariantTableSampleIndexOrderMapper<ImmutableBytesWritable, Put> {

        private Map<Integer, ByteArrayOutputStream> mendelianErrorsMap = new HashMap<>();
        private Map<Integer, Map<String, Integer>> genotypeCount = new HashMap<>();
        private List<List<Integer>> trios;
        private byte[] family;
        private SampleIndexToHBaseConverter converter;

        @Override
        protected void setup(Context context) throws IOException, InterruptedException {
            family = new GenomeHelper(context.getConfiguration()).getColumnFamily();
            converter = new SampleIndexToHBaseConverter(family);

            int[] sampleIds = context.getConfiguration().getInts(TRIOS_LIST);
            trios = new ArrayList<>(sampleIds.length / 3);

            for (int i = 0; i < sampleIds.length; i += 3) {
                trios.add(Arrays.asList(
                        sampleIds[i],
                        sampleIds[i + 1],
                        sampleIds[i + 2]));
            }

            for (List<Integer> trio : trios) {
                Integer child = trio.get(2);
                mendelianErrorsMap.put(child, new ByteArrayOutputStream());
                genotypeCount.put(child, new HashMap<>());
            }

        }

        @Override
        protected void map(ImmutableBytesWritable key, Result value, Context context) throws IOException, InterruptedException {
            Variant variant = VariantPhoenixKeyFactory.extractVariantFromVariantRowKey(value.getRow());
            String chromosome = variant.getChromosome();

            Map<Integer, String> gtMap = new HashMap<>();
            for (Cell cell : value.rawCells()) {
                Integer sampleId = VariantPhoenixHelper
                        .extractSampleId(cell.getQualifierArray(), cell.getQualifierOffset(), cell.getQualifierLength());
                if (sampleId != null) {
                    ImmutableBytesWritable ptr = new ImmutableBytesWritable(
                            cell.getValueArray(),
                            cell.getValueOffset(),
                            cell.getValueLength());
                    PArrayDataType.positionAtArrayElement(ptr, 0, PVarchar.INSTANCE, null);
                    String gt = Bytes.toString(ptr.get(), ptr.getOffset(), ptr.getLength());
                    gtMap.put(sampleId, gt);
                }
            }


            for (List<Integer> trio : trios) {
                Integer father = trio.get(0);
                Integer mother = trio.get(1);
                Integer child = trio.get(2);

                String fatherGtStr = gtMap.get(father);
                String motherGtStr = gtMap.get(mother);
                String childGtStr = gtMap.get(child);

                int idx = genotypeCount.get(child).merge(childGtStr, 1, Integer::sum) - 1;
                if ((fatherGtStr != null || father == -1) && (motherGtStr != null || mother == -1) && childGtStr != null) {
                    Genotype fatherGt;
                    Genotype motherGt;
                    Genotype childGt;
                    try {
                        fatherGt = fatherGtStr == null ? null : new Genotype(fatherGtStr);
                        motherGt = motherGtStr == null ? null : new Genotype(motherGtStr);
                        childGt = new Genotype(childGtStr);
                    } catch (IllegalArgumentException e) {
                        // Skip malformed GT
                        context.getCounter(COUNTER_GROUP_NAME, "wrong_gt").increment(1);
                        continue;
                    }
                    Integer me = MendelianError.compute(fatherGt, motherGt, childGt, chromosome);
                    context.getCounter(COUNTER_GROUP_NAME, "me_" + me).increment(1);
                    if (me > 0) {
                        ByteArrayOutputStream stream = mendelianErrorsMap.get(child);
                        converter.serializeMendelianError(stream, variant, childGtStr, idx, me);
                    }
                }
            }

        }

        @Override
        public void flush(Context context, String chromosome, int position) throws IOException, InterruptedException {

            for (Map.Entry<Integer, ByteArrayOutputStream> entry : mendelianErrorsMap.entrySet()) {
                Integer sampleId = entry.getKey();
                ByteArrayOutputStream value = entry.getValue();


                byte[] row = HBaseToSampleIndexConverter.toRowKey(sampleId, chromosome, position);

                if (value.size() > 0) {
                    Put put = new Put(row);
                    // Copy value, as the ByteArrayOutputStream is erased
                    put.addColumn(family, HBaseToSampleIndexConverter.toMendelianErrorColumn(), value.toByteArray());

                    context.write(new ImmutableBytesWritable(row), put);
                    value.reset();
                }
            }

            genotypeCount.values().forEach(Map::clear);
        }
    }

    public static void main(String[] args) throws Exception {
        try {
            System.exit(new MendelianErrorDriver().privateMain(args, null));
        } catch (Exception e) {
            LOGGER.error("Error executing " + MendelianErrorDriver.class, e);
            System.exit(1);
        }
    }
}

{'fraction_non_alphanumeric': 0.1032741617357002, 'fraction_numerical': 0.003076923076923077, 'mean_word_length': 2.3479926043317487, 'pattern_counts': {'":': 0, '<': 59, '<?xml version=': 0, '>': 63, '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': '1897685', 'n_tokens_mistral': 7020, 'n_tokens_neox': 6600, 'n_words': 1276}