WIP:Performance long reads

src/main/scala/org/biodatageeks/sequila/pileup/PileupRunner.scala

@@ -3,6 +3,7 @@ package org.biodatageeks.sequila.pileup
 import java.io.{OutputStreamWriter, PrintWriter}
 
 import org.apache.spark.sql.{SequilaSession, SparkSession}
+import org.bdgenomics.utils.instrumentation.{Metrics, MetricsListener, RecordedMetrics}
 import org.biodatageeks.sequila.utils.{InternalParams, SequilaRegister}
 
 object PileupRunner {
@@ -14,8 +15,7 @@ object PileupRunner {
       .config("spark.driver.memory","48g")
       .config( "spark.serializer", "org.apache.spark.serializer.KryoSerializer" )
       .config("spark.kryo.registrator", "org.biodatageeks.sequila.pileup.serializers.CustomKryoRegistrator")
-      .config("spark.driver.maxResultSize","5g")
-      //      .config("spark.kryoserializer.buffer.max", "512m")
+//      .config("spark.kryoserializer.buffer.max", "512m")
       .enableHiveSupport()
       .getOrCreate()
 
@@ -62,8 +62,7 @@ object PileupRunner {
                  |FROM  pileup('$tableNameBAM', 'NA12878.proper.wgs.md', '${referencePath}', false)
                """.stripMargin
 
-
-    //    ss.sqlContext.setConf("spark.biodatageeks.readAligment.method", "disq")
+//    ss.sqlContext.setConf("spark.biodatageeks.readAligment.method", "disq")
     ss.sqlContext.setConf("spark.biodatageeks.bam.useGKLInflate","true")
     ss.sparkContext.setLogLevel("WARN")
     val results = ss.sql(query)
@@ -73,9 +72,7 @@ object PileupRunner {
         .write.orc("/tmp/wes_seq_pileup_2")
       //        .show()
     }
-//    val writer = new PrintWriter(new OutputStreamWriter(System.out, "UTF-8"))
-//    Metrics.print(writer, Some(metricsListener.metrics.sparkMetrics.stageTimes))
-//    writer.close()
+
 
     ss.stop()
   }

src/main/scala/org/biodatageeks/sequila/pileup/model/CigarDerivedConf.scala

@@ -38,7 +38,7 @@ case class CigarDerivedConf(
 object CigarDerivedConf {
   def create(start: Int, cigar:Cigar) ={
     val hasClip = cigar.isLeftClipped
-    val softClipLength = if(hasClip) {
+    val softClipLength = if(hasClip && cigar.getFirstCigarElement.getOperator != CigarOperator.HARD_CLIP) {
       val cigarElement = cigar.getFirstCigarElement
       cigarElement.getLength
     } else 0

src/main/scala/org/biodatageeks/sequila/pileup/model/Read.scala

@@ -7,8 +7,12 @@ import org.biodatageeks.sequila.pileup.conf.Conf
 import org.biodatageeks.sequila.pileup.conf.QualityConstants.{OUTER_QUAL_SIZE, QUAL_INDEX_SHIFT, REF_SYMBOL}
 import org.biodatageeks.sequila.pileup.model.Quals._
 import org.biodatageeks.sequila.pileup.model.Alts._
+import org.biodatageeks.sequila.pileup.timers.PileupTimers.{AnalyzeReadsCalculateAltsLoopTimer, AnalyzeReadsCalculateQualsUpdateAltsTimer, AnalyzeReadscalculatePositionInReadSeq}
 
 import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+
+case class PositionInReadState(position: Int, numInsertions: Int)
 
 object ReadOperations {
 
@@ -27,7 +31,7 @@ case class ExtendedReads(read: SAMRecord) {
                  ): Unit = {
     val start = read.getStart
     val cigar = read.getCigar
-    val bQual = read.getBaseQualities
+    val bQual = if (conf.value.includeBaseQualities) read.getBaseQualities else null
     val isPositiveStrand = ! read.getReadNegativeStrandFlag
     calculateEvents(contig, agg, contigMaxReadLen, start, cigar)
     val foundAlts = calculateAlts(agg, agg.qualityCache, start, cigar, bQual, isPositiveStrand, conf)
@@ -74,53 +78,106 @@ case class ExtendedReads(read: SAMRecord) {
   }
 
 
-  def calculatePositionInReadSeq(mdPosition: Int, cigar: Cigar): Int = {
-    if (!cigar.containsOperator(CigarOperator.INSERTION))
-      return mdPosition
-
-    var numInsertions = 0
+//  def calculatePositionInReadSeq(mdPosition: Int, cigar: Cigar): Int = {
+//    if (!cigar.containsOperator(CigarOperator.INSERTION))
+//      return mdPosition
+//
+//    var numInsertions = 0
+//    val cigarIterator = cigar.iterator()
+//    var position = 0
+//
+//    while (cigarIterator.hasNext) {
+//      if (position > mdPosition + numInsertions)
+//        return mdPosition + numInsertions
+//      val cigarElement = cigarIterator.next()
+//      val cigarOpLength = cigarElement.getLength
+//      val cigarOp = cigarElement.getOperator
+//
+//      if (cigarOp == CigarOperator.INSERTION)
+//        numInsertions += cigarOpLength
+//      else if (cigarOp != CigarOperator.HARD_CLIP)
+//        position = position + cigarOpLength
+//    }
+//    mdPosition + numInsertions
+//  }
+  def calculatePositionInReadSeq(mdPosition: Int, insertions: Array[Int], positionInReadState: PositionInReadState): PositionInReadState = {
+    if (insertions.isEmpty)
+      PositionInReadState(mdPosition, 0)
+    else {
+      var numInsertions = positionInReadState.numInsertions
+      var position = 0
+      var idx = positionInReadState.numInsertions
+      while (idx < insertions.length ){
+        position = insertions(idx)
+        val posWithIns = mdPosition + numInsertions
+        if (position > posWithIns)
+          return PositionInReadState(posWithIns, numInsertions)
+        else
+          numInsertions += 1
+        idx += 1
+      }
+      PositionInReadState(mdPosition + numInsertions, numInsertions)
+    }
+  }
+  def calculateInsertions(cigar: Cigar): Array[Int] = {
+    val insertions = new ArrayBuffer[Int]()
     val cigarIterator = cigar.iterator()
     var position = 0
 
     while (cigarIterator.hasNext) {
-      if (position > mdPosition + numInsertions)
-        return mdPosition + numInsertions
       val cigarElement = cigarIterator.next()
       val cigarOpLength = cigarElement.getLength
       val cigarOp = cigarElement.getOperator
 
       if (cigarOp == CigarOperator.INSERTION)
-        numInsertions += cigarOpLength
+        insertions.append(position)
       else if (cigarOp != CigarOperator.HARD_CLIP)
         position = position + cigarOpLength
     }
-    mdPosition + numInsertions
+    insertions.toArray
   }
 
   def calculateAlts(aggregate: ContigAggregate, qualityCache: QualityCache, start: Int,
                     cigar: Cigar, bQual: Array[Byte], 
                     isPositiveStrand:Boolean, conf: Broadcast[Conf]): scala.collection.Set[Int] = {
     var position = start
     val ops = MDTagParser.parseMDTag(read.getAttribute("MD").toString)
-
+    val insertions = calculateInsertions(cigar)
     var delCounter = 0
     var altsPositions = mutable.Set.empty[Int]
     val clipLen =
       if (cigar.getCigarElement(0).getOperator == CigarOperator.SOFT_CLIP)
         cigar.getCigarElement(0).getLength else 0
-
+    val readString = read.getReadString
     position += clipLen
+    val zeroByte = 0.toByte
+    var positionInReadState = PositionInReadState(0, 0)
 
+    AnalyzeReadsCalculateAltsLoopTimer.time {
     for (mdtag <- ops) {
       if (mdtag.isDeletion) {
         delCounter += 1
         position += 1
       } else if (mdtag.base != 'S') {
         position += 1
 
-        val indexInSeq = calculatePositionInReadSeq(position - start - delCounter, cigar)
-        val altBase = if (isPositiveStrand) read.getReadString.charAt(indexInSeq - 1).toUpper else read.getReadString.charAt(indexInSeq - 1).toLower
-        val altBaseQual = if (conf.value.isBinningEnabled) (bQual(indexInSeq - 1)/conf.value.binSize).toByte else bQual(indexInSeq - 1)
+        val indexInSeq = AnalyzeReadscalculatePositionInReadSeq.time {
+          val pos = position - start - delCounter
+          if(insertions.isEmpty)
+            pos
+          else {
+            positionInReadState = calculatePositionInReadSeq(pos, insertions, positionInReadState)
+            positionInReadState.position
+          }
+        }
+        val altBase = if (isPositiveStrand) readString.charAt(indexInSeq - 1).toUpper else readString.charAt(indexInSeq - 1).toLower
+        val altBaseQual = {
+          if (bQual == null) zeroByte
+          else if (conf.value.isBinningEnabled)
+            (bQual(indexInSeq - 1) / conf.value.binSize).toByte
+          else
+            bQual(indexInSeq - 1)
+        }
         val altPosition = position - clipLen - 1
 
         if (conf.value.includeBaseQualities) {
@@ -129,14 +186,15 @@ case class ExtendedReads(read: SAMRecord) {
           else
             aggregate.quals.updateQuals(altPosition, altBase, altBaseQual, firstUpdate = true, updateMax = true, conf)
         }
-        aggregate.alts.updateAlts(altPosition, altBase)
+        AnalyzeReadsCalculateQualsUpdateAltsTimer.time { aggregate.alts.updateAlts(altPosition, altBase) }
         if (conf.value.includeBaseQualities)
           aggregate.altsKeyCache.add(altPosition)
         altsPositions += altPosition
       }
       else if (mdtag.base == 'S')
         position += mdtag.length
     }
+  }
     altsPositions
   }
 

src/test/resources/long_reads/rel5-guppy-0.3.0-chunk10k.chr1.bed

@@ -0,0 +1 @@
+chr1	1	100000

src/test/resources/long_reads/rel5-guppy-0.3.0-chunk10k.chr1.fasta.fai

@@ -0,0 +1 @@
+chr1	99999	6	99999	100000

src/test/scala/org/biodatageeks/sequila/tests/pileup/PileupLongReadsTestSuite.scala

@@ -0,0 +1,33 @@
+package org.biodatageeks.sequila.tests.pileup
+
+import org.apache.spark.sql.SequilaSession
+import org.biodatageeks.sequila.utils.{Columns, SequilaRegister}
+
+class PileupLongReadsTestSuite extends PileupTestBase{
+  test("Pileup without BQ") {
+    val ss = SequilaSession(spark)
+    SequilaRegister.register(ss)
+    val query =
+      s"""
+         |SELECT ${Columns.CONTIG}, ${Columns.START}, ${Columns.END}, ${Columns.REF}, ${Columns.COVERAGE},${Columns.ALTS}
+         |FROM  pileup('${tableNameLongReads}', '${sampleLongReadsId}' , '$referenceLongReadsPath', false)
+                 """.stripMargin
+
+    ss.sql(query).count
+
+  }
+
+  test("Pileup with BQ") {
+    val ss = SequilaSession(spark)
+    SequilaRegister.register(ss)
+    val query =
+      s"""
+         |SELECT ${Columns.CONTIG}, ${Columns.START}, ${Columns.END}, ${Columns.REF}, ${Columns.COVERAGE},${Columns.ALTS}
+         |FROM  pileup('${tableNameLongReads}', '${sampleLongReadsId}' , '$referenceLongReadsPath', true)
+                 """.stripMargin
+
+    ss.sql(query).count
+
+  }
+
+}

src/test/scala/org/biodatageeks/sequila/tests/pileup/PileupTestBase.scala

@@ -25,6 +25,12 @@ class PileupTestBase extends FunSuite
   val tableName = "reads_bam"
   val tableNameCRAM = "reads_cram"
 
+
+  val sampleLongReadsId = "rel5-guppy-0.3.0-chunk10k.chr1"
+  val referenceLongReadsPath: String = getClass.getResource(s"/long_reads/${sampleLongReadsId}.fasta").getPath
+  val bamLongReadsPath: String = getClass.getResource(s"/long_reads/${sampleLongReadsId}.bam").getPath
+  val tableNameLongReads = "reads_long"
+
   val schema: StructType = StructType(
     List(
       StructField("contig", StringType, nullable = true),
@@ -59,6 +65,16 @@ class PileupTestBase extends FunSuite
          |
       """.stripMargin)
 
+
+    spark.sql(s"DROP TABLE IF EXISTS $tableNameLongReads")
+    spark.sql(
+      s"""
+         |CREATE TABLE $tableNameLongReads
+         |USING org.biodatageeks.sequila.datasources.BAM.BAMDataSource
+         |OPTIONS(path "$bamLongReadsPath", refPath "$referenceLongReadsPath" )
+         |
+      """.stripMargin)
+
   }
 
 }

src/test/scala/org/biodatageeks/sequila/tests/pileup/processing/PositionInReadTest.scala

@@ -1,7 +1,7 @@
 package org.biodatageeks.sequila.tests.pileup.processing
 
 import htsjdk.samtools.{Cigar, CigarElement, CigarOperator}
-import org.biodatageeks.sequila.pileup.model.ExtendedReads
+import org.biodatageeks.sequila.pileup.model.{ExtendedReads, PositionInReadState}
 import org.scalatest.FunSuite
 
 import scala.collection.JavaConversions.seqAsJavaList
@@ -15,10 +15,13 @@ class PositionInReadTest extends FunSuite{
     val cElement3 = new CigarElement(23, CigarOperator.M)
     val c = new Cigar(seqAsJavaList(List(cElement1,cElement2,cElement3)))
 
+
     val testRead = ExtendedReads(null)
-    val ind = testRead.calculatePositionInReadSeq(10,c)
+    val insertions = testRead.calculateInsertions(c)
+    val posisitonState = PositionInReadState(0,0)
+    val ind = testRead.calculatePositionInReadSeq(10, insertions, posisitonState).position
     assert(ind==11)
-    val ind2 = testRead.calculatePositionInReadSeq(20,c)
+    val ind2 = testRead.calculatePositionInReadSeq(20,insertions, posisitonState).position
     assert(ind2==21)
   }
 
@@ -30,9 +33,11 @@ class PositionInReadTest extends FunSuite{
     val c = new Cigar(seqAsJavaList(List(cElement1,cElement2,cElement3)))
 
     val testRead = ExtendedReads(null)
-    val ind = testRead.calculatePositionInReadSeq(10,c)
+    val insertions = testRead.calculateInsertions(c)
+    val posisitonState = PositionInReadState(0,0)
+    val ind = testRead.calculatePositionInReadSeq(10, insertions, posisitonState).position
     assert(ind==10)
-    val ind2 = testRead.calculatePositionInReadSeq(23,c)
+    val ind2 = testRead.calculatePositionInReadSeq(23, insertions, posisitonState).position
     assert(ind2==23)
   }
 
@@ -45,9 +50,11 @@ class PositionInReadTest extends FunSuite{
     print (c.getReadLength)
 
     val testRead = ExtendedReads(null)
-    val ind = testRead.calculatePositionInReadSeq(10,c)
+    val insertions = testRead.calculateInsertions(c)
+    val posisitonState = PositionInReadState(0,0)
+    val ind = testRead.calculatePositionInReadSeq(10,insertions, posisitonState).position
     assert(ind==11)
-    val ind2 = testRead.calculatePositionInReadSeq(23,c)
+    val ind2 = testRead.calculatePositionInReadSeq(23,insertions, posisitonState).position
     assert(ind2==24)
   }