Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Changes to RDD utility files for new variant schema #114

Merged
merged 5 commits into from
Feb 20, 2014
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ private[adam] class VariantContextConverter extends Serializable with Logging {
if (vc.isFiltered) {
val annotations : VariantCallingAnnotations.Builder =
VariantCallingAnnotations.newBuilder
.setVariantIsPassing(!vc.isFiltered)
.setVariantIsPassing(!vc.isFiltered.booleanValue)
.setVariantFilters(new util.ArrayList(vc.getFilters))
genotype.setVariantCallingAnnotations(annotations.build)
}
Expand All @@ -87,7 +87,7 @@ private[adam] class VariantContextConverter extends Serializable with Logging {
genotype.setPhaseQuality(g.getExtendedAttribute(VCFConstants.PHASE_QUALITY_KEY).asInstanceOf[java.lang.Integer])

if (g.hasExtendedAttribute(VCFConstants.PHASE_SET_KEY))
genotype.setPhaseSetId(g.getExtendedAttribute(VCFConstants.PHASE_SET_KEY).asInstanceOf[Integer])
genotype.setPhaseSetId(Integer.parseInt(g.getExtendedAttribute(VCFConstants.PHASE_SET_KEY).asInstanceOf[java.lang.String]))

genotype.build
}).toSeq
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,6 @@ import parquet.filter.ColumnPredicates.equalTo

class GenotypeVarFilterPASSPredicate extends UnboundRecordFilter {
def bind(readers: Iterable[ColumnReader]): RecordFilter = {
or(column("varIsFiltered", equalTo(null.asInstanceOf[Boolean])), column("varIsFiltered", equalTo(false))).bind(readers)
column("variantCallingAnnotations.variantIsPassing", equalTo(true)).bind(readers)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ import edu.berkeley.cs.amplab.adam.converters.SAMRecordConverter
import edu.berkeley.cs.amplab.adam.models._
import org.apache.hadoop.fs.FileSystem
import edu.berkeley.cs.amplab.adam.projections.{ADAMRecordField, Projection}
import edu.berkeley.cs.amplab.adam.rdd.compare.CompareAdam
import fi.tkk.ics.hadoop.bam.util.SAMHeaderReader
import fi.tkk.ics.hadoop.bam.{SAMRecordWritable, AnySAMInputFormat}
import org.apache.avro.Schema
Expand Down Expand Up @@ -232,11 +231,6 @@ class AdamContext(sc: SparkContext) extends Serializable with Logging {
}
}

def adamCompareFiles(file1Path: String, file2Path: String,
predicateFactory: (Map[Int, Int]) => (SingleReadBucket, SingleReadBucket) => Boolean) = {
CompareAdam.compareADAM(sc, file1Path, file2Path, predicateFactory)
}

/**
* Searches a path recursively, returning the names of all directories in the tree whose
* name matches the given regex.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,23 +15,21 @@
*/
package edu.berkeley.cs.amplab.adam.rdd

import edu.berkeley.cs.amplab.adam.avro.{ADAMGenotype, ADAMGenotypeAllele, ADAMPileup, ADAMRecord, ADAMVariant, ADAMVariantDomain}
import edu.berkeley.cs.amplab.adam.converters.{GenotypesToVariantsConverter, VariantContextConverter}
import edu.berkeley.cs.amplab.adam.models.{ADAMRod, ADAMVariantContext, ReferencePosition, SequenceDictionary, SequenceRecord, SingleReadBucket, SnpTable}
import edu.berkeley.cs.amplab.adam.projections.{ADAMVariantAnnotations, ADAMVariantField}
import edu.berkeley.cs.amplab.adam.rdd.AdamContext._
import edu.berkeley.cs.amplab.adam.util.ParquetLogger
import org.apache.avro.specific.SpecificRecord
import parquet.hadoop.metadata.CompressionCodecName
import org.apache.hadoop.mapreduce.Job
import org.apache.spark.Logging
import org.apache.spark.SparkContext._
import org.apache.spark.rdd.RDD
import org.broadinstitute.variant.variantcontext.{Allele, VariantContext}
import parquet.avro.{AvroParquetOutputFormat, AvroWriteSupport}
import parquet.hadoop.ParquetOutputFormat
import parquet.hadoop.metadata.CompressionCodecName
import parquet.avro.{AvroParquetOutputFormat, AvroWriteSupport}
import parquet.hadoop.util.ContextUtil
import org.apache.avro.specific.SpecificRecord
import edu.berkeley.cs.amplab.adam.avro.{ADAMPileup, ADAMRecord}
import edu.berkeley.cs.amplab.adam.models.{SequenceRecord, SequenceDictionary, SingleReadBucket, SnpTable, ReferencePosition, ADAMRod}
import org.apache.spark.rdd.RDD
import org.apache.spark.SparkContext._
import org.apache.spark.Logging
import java.io.File
import edu.berkeley.cs.amplab.adam.util.ParquetLogger
import java.util.logging.Level
import edu.berkeley.cs.amplab.adam.rich.RichADAMRecord._

class AdamRDDFunctions[T <% SpecificRecord : Manifest](rdd: RDD[T]) extends Serializable {

Expand Down Expand Up @@ -140,7 +138,7 @@ class AdamRecordRDDFunctions(rdd: RDD[ADAMRecord]) extends Serializable with Log
* @param bucketSize Size in basepairs of buckets. Larger buckets take more time per
* bucket to convert, but have lower skew. Default is 1000.
*
* @returns RDD of ADAMRods.
* @return RDD of ADAMRods.
*/
def adamRecords2Rods (bucketSize: Int = 1000): RDD[ADAMRod] = {

Expand All @@ -167,7 +165,7 @@ class AdamRecordRDDFunctions(rdd: RDD[ADAMRecord]) extends Serializable with Log

val bucketedReads = rdd.filter(_.getStart != null)
.flatMap(mapToBucket)
.groupByKey
.groupByKey()

println ("Have reads in buckets.")

Expand Down Expand Up @@ -276,107 +274,3 @@ class AdamRodRDDFunctions(rdd: RDD[ADAMRod]) extends Serializable with Logging {
totalBases.toDouble / rdd.count.toDouble
}
}

class AdamVariantContextRDDFunctions(rdd: RDD[ADAMVariantContext]) extends Serializable {

/**
* Save function for variant contexts. Disaggregates internal fields of variant context
* and saves to Parquet files.
*
* @param filePath Master file path for parquet files.
* @param blockSize Parquet block size.
* @param pageSize Parquet page size.
* @param compressCodec Parquet compression codec.
* @param disableDictionaryEncoding If true, disables dictionary encoding in Parquet.
* @return Returns the initial RDD.
*/
def adamSave(filePath: String, blockSize: Int = 128 * 1024 * 1024,
pageSize: Int = 1 * 1024 * 1024, compressCodec: CompressionCodecName = CompressionCodecName.GZIP,
disableDictionaryEncoding: Boolean = false): RDD[ADAMVariantContext] = {

// Add the Void Key
val variantToSave: RDD[ADAMVariant] = rdd.flatMap(p => p.variants)
val genotypeToSave: RDD[ADAMGenotype] = rdd.flatMap(p => p.genotypes)
val domainsToSave: RDD[ADAMVariantDomain] = rdd.flatMap(p => p.domains)

// save records
variantToSave.adamSave(filePath + ".v",
blockSize,
pageSize,
compressCodec,
disableDictionaryEncoding)
genotypeToSave.adamSave(filePath + ".g",
blockSize,
pageSize,
compressCodec,
disableDictionaryEncoding)

// check if we have domains to save or not
if (domainsToSave.count() != 0) {
val fileExtension = ADAMVariantAnnotations.fileExtensions(ADAMVariantAnnotations.ADAMVariantDomain)

domainsToSave.adamSave(filePath + fileExtension,
blockSize,
pageSize,
compressCodec,
disableDictionaryEncoding)
}

rdd
}

}

class AdamGenotypeRDDFunctions(rdd: RDD[ADAMGenotype]) extends Serializable {

/**
* Validates that an RDD of genotypes is correctly formed.
*
* @return True if RDD is correctly formed.
* @throws IllegalArgumentException Throws exception if RDD is not correctly formed.
*/
def adamValidateGenotypes(): Boolean = {
val validator = new GenotypesToVariantsConverter(true, true)
val groupedGenotypes = rdd.groupBy(g => (g.getPosition, g.getSampleId))
groupedGenotypes.map(_._2.toList).foreach(validator.validateGenotypes)

true
}

/**
* Converts an RDD of GATK Variants to an RDD of ADAMVariantContext objects
*/
def gatkVariantsToADAMVariants(gatkAlleles : RDD[VariantContext]) : RDD[ADAMVariantContext] = {
val vcc : VariantContextConverter = new VariantContextConverter
gatkAlleles.flatMap(vcc.convert(_))
}

/**
* Calculates Variants from an RDD of genotypes. This allows for on-the-fly creation of variant
* data from a subset of a population. This function also allows an RDD of variant data to be provided.
* Data can be taken from this RDD by adding to the projection set.
*
* @param variants Optional RDD of variant data to supplement genotype info.
* @param variantProjection The set of fields to copy from the variant data, if this is provided.
* @param performValidation Whether to validate that the genotype data is well formed.
* @param failOnValidationError If validation is performed and failOnValidationError is true, an exception will
* be thrown if an error is encountered.
* @return An RDD containing variant data.
*
* @throws IllegalArgumentException Throws an exception if performValidation and failOnValidationError are true
* and the RDD of genotypes has bad data.
*/
def adamConvertGenotypes(variants: Option[RDD[ADAMVariant]] = None,
variantProjection: Set[ADAMVariantField.Value] = Set[ADAMVariantField.Value](),
performValidation: Boolean = false,
failOnValidationError: Boolean = false): RDD[ADAMVariant] = {
val computer = new GenotypesToVariantsConverter(performValidation, failOnValidationError)
val groupedGenotypes = rdd.groupBy(g => g.getPosition)
val groupedGenotypesWithVariants: RDD[(java.lang.Long, (Seq[ADAMGenotype], Option[Seq[ADAMVariant]]))] = variants match {
case Some(o) => groupedGenotypes.leftOuterJoin(o.asInstanceOf[RDD[ADAMVariant]].groupBy(_.getPosition))
case None => groupedGenotypes.map(kv => (kv._1, (kv._2, None.asInstanceOf[Option[Seq[ADAMVariant]]])))
}

groupedGenotypesWithVariants.map(_._2).flatMap(vg => computer.convert(vg._1, vg._2, variantProjection))
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
/*
* Copyright (c) 2014. Mount Sinai School of Medicine
*
* Licensed 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 edu.berkeley.cs.amplab.adam.rdd.variation

import org.apache.spark.rdd.RDD
import edu.berkeley.cs.amplab.adam.models.ADAMVariantContext
import org.apache.spark.{Logging, SparkContext}
import org.apache.hadoop.mapreduce.Job
import edu.berkeley.cs.amplab.adam.converters.VariantContextConverter
import fi.tkk.ics.hadoop.bam.{VariantContextWritable, VCFInputFormat}
import org.apache.hadoop.io.LongWritable
import parquet.hadoop.util.ContextUtil


object ADAMVariationContext {
implicit def sparkContextToADAMVariationContext(sc: SparkContext): ADAMVariationContext = new ADAMVariationContext(sc)

implicit def rddToADAMVariantContextRDD(rdd: RDD[ADAMVariantContext]) = new ADAMVariantContextRDDFunctions(rdd)
}

class ADAMVariationContext(sc: SparkContext) extends Serializable with Logging {

/**
* This method will create a new RDD of VariantContext objects
* @param filePath: input VCF file to read
* @return RDD of variants
*/
def adamVCFLoad(filePath: String): RDD[ADAMVariantContext] = {
log.info("Reading VCF file from %s".format(filePath))
val job = Job.getInstance(sc.hadoopConfiguration)
val vcc = new VariantContextConverter
val records = sc.newAPIHadoopFile(
filePath,
classOf[VCFInputFormat], classOf[LongWritable], classOf[VariantContextWritable],
ContextUtil.getConfiguration(job)
)
log.info("Converted %d records".format(records.count()))
records.flatMap(p => vcc.convert(p._2.get))
}
}


Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
/*
* Copyright (c) 2013. Mount Sinai School of Medicine
*
* Licensed 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 edu.berkeley.cs.amplab.adam.rdd.variation

import org.apache.spark.Logging
import org.apache.spark.rdd.RDD
import edu.berkeley.cs.amplab.adam.models.ADAMVariantContext
import edu.berkeley.cs.amplab.adam.avro.{ADAMGenotype, ADAMDatabaseVariantAnnotation}
import org.apache.spark.SparkContext._
import edu.berkeley.cs.amplab.adam.rich.RichADAMVariant

class ADAMVariantContextRDDFunctions(rdd: RDD[ADAMVariantContext]) extends Serializable with Logging {
initLogging()

/**
* Left outer join database variant annotations
*
*/
def joinDatabaseVariantAnnotation(ann: RDD[ADAMDatabaseVariantAnnotation]): RDD[ADAMVariantContext] = {
rdd.keyBy(_.variant)
.leftOuterJoin(ann.keyBy(_.getVariant))
.values
.map { case (v:ADAMVariantContext, a) => new ADAMVariantContext(v.variant, v.genotypes, databases = a) }
}
}

class ADAMGenotypeRDDFunctions(rdd: RDD[ADAMGenotype]) extends Serializable with Logging {
initLogging()

def toADAMVariantContext(): RDD[ADAMVariantContext] = {
rdd.keyBy({ g => RichADAMVariant.variantToRichVariant(g.getVariant) })
.groupByKey
.map { case (v:RichADAMVariant, g) => new ADAMVariantContext(v, genotypes = g) }
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
/*
* Copyright (c) 2014. Mount Sinai School of Medicine
*
* Licensed 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 edu.berkeley.cs.amplab.adam.rich

import edu.berkeley.cs.amplab.adam.avro.{ADAMContig, ADAMVariant}
import java.util.Arrays

object RichADAMVariant {
implicit def variantToRichVariant(variant: ADAMVariant): RichADAMVariant = new RichADAMVariant(variant)
implicit def richVariantToVariant(variant: RichADAMVariant): ADAMVariant = variant.variant
}

class RichADAMVariant(val variant: ADAMVariant) {
// Only include the contigName in the hash
val hashObjects = Array[Object](variant.getContig.getContigName,
variant.getPosition, variant.getReferenceAllele, variant.getVariantAllele)
override def hashCode = Arrays.hashCode(hashObjects)

private def isSameContig(left: ADAMContig, right: ADAMContig): Boolean = {
left.getContigName == right.getContigName && (
left.getContigMD5 == null || right.contigMD5 == null || left.getContigMD5 == right.getContigMD5
)
}

override def equals(o: Any) = o match {
case that: RichADAMVariant => {
variant.getPosition == that.variant.getPosition &&
isSameContig(variant.getContig, that.variant.getContig) &&
variant.getReferenceAllele == that.variant.getReferenceAllele &&
variant.getVariantAllele == that.variant.getVariantAllele
}
case _ => false
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@ import edu.berkeley.cs.amplab.adam.util.{ParquetLogger, SparkFunSuite}
import org.scalatest.BeforeAndAfter
import java.util.logging.Level
import java.io.File
import edu.berkeley.cs.amplab.adam.avro.{ADAMContig, ADAMVariant, ADAMGenotype}
import edu.berkeley.cs.amplab.adam.avro.{ADAMContig,
ADAMVariant,
ADAMGenotype,
VariantCallingAnnotations}
import org.apache.spark.rdd.RDD
import edu.berkeley.cs.amplab.adam.rdd.AdamContext._
import com.google.common.io.Files
Expand All @@ -40,9 +43,17 @@ class GenotypePredicatesSuite extends SparkFunSuite with BeforeAndAfter {
.setVariantAllele("C")
.build

val passFilterAnnotation =
VariantCallingAnnotations.newBuilder().setVariantIsPassing(true).build()
val failFilterAnnotation =
VariantCallingAnnotations.newBuilder().setVariantIsPassing(false).build()

val genotypes = sc.parallelize(List(
ADAMGenotype.newBuilder().setVariant(v0).setVarIsFiltered(true).build(),
ADAMGenotype.newBuilder().setVariant(v0).build()
ADAMGenotype.newBuilder().setVariant(v0)
.setVariantCallingAnnotations(passFilterAnnotation).build(),
ADAMGenotype.newBuilder()
.setVariant(v0)
.setVariantCallingAnnotations(failFilterAnnotation).build()
))

genotypesParquetFile = new File(Files.createTempDir(), "genotypes")
Expand Down
Loading