[query] move block matrix read out of HailContext

Author: ehigham

hail/hail/src/is/hail/HailContext.scala

@@ -1,22 +1,15 @@
 package is.hail
 
-import is.hail.backend.{Backend, ExecuteContext}
+import is.hail.backend.Backend
 import is.hail.backend.spark.SparkBackend
 import is.hail.expr.ir.functions.IRFunctionRegistry
 import is.hail.io.fs.FS
 import is.hail.utils._
 
-import scala.reflect.ClassTag
-
-import java.io.InputStream
 import java.util.Properties
 
 import org.apache.log4j.{LogManager, PropertyConfigurator}
 import org.apache.spark._
-import org.apache.spark.executor.InputMetrics
-import org.apache.spark.rdd.RDD
-
-case class FilePartition(index: Int, file: String) extends Partition
 
 object HailContext {
 
@@ -115,32 +108,6 @@ object HailContext {
       theContext.backend.close()
       theContext = null
     }
-
-  def readPartitions[T: ClassTag](
-    ctx: ExecuteContext,
-    path: String,
-    partFiles: IndexedSeq[String],
-    read: (Int, InputStream, InputMetrics) => Iterator[T],
-    optPartitioner: Option[Partitioner] = None,
-  ): RDD[T] = {
-    val nPartitions = partFiles.length
-
-    val fsBc = ctx.fsBc
-
-    new RDD[T](ctx.backend.asSpark.sc, Nil) {
-      def getPartitions: Array[Partition] =
-        Array.tabulate(nPartitions)(i => FilePartition(i, partFiles(i)))
-
-      override def compute(split: Partition, context: TaskContext): Iterator[T] = {
-        val p = split.asInstanceOf[FilePartition]
-        val filename = path + "/parts/" + p.file
-        val in = fsBc.value.open(filename)
-        read(p.index, in, context.taskMetrics().inputMetrics)
-      }
-
-      @transient override val partitioner: Option[Partitioner] = optPartitioner
-    }
-  }
 }
 
 class HailContext private (

hail/hail/src/is/hail/linalg/BlockMatrix.scala

@@ -1,6 +1,5 @@
 package is.hail.linalg
 
-import is.hail._
 import is.hail.annotations._
 import is.hail.backend.{BroadcastValue, ExecuteContext, HailStateManager}
 import is.hail.backend.spark.{SparkBackend, SparkTaskContext}
@@ -28,7 +27,6 @@ import breeze.numerics.{abs => breezeAbs, log => breezeLog, pow => breezePow, sq
 import breeze.stats.distributions.RandBasis
 import org.apache.commons.lang3.StringUtils
 import org.apache.spark._
-import org.apache.spark.executor.InputMetrics
 import org.apache.spark.mllib.linalg.distributed.{GridPartitioner => _, _}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
@@ -216,16 +214,27 @@ object BlockMatrix {
       readMetadata(ctx.fs, uri)
 
     val gp = GridPartitioner(blockSize, nRows, nCols, maybeFiltered)
+    val nPartitions = partFiles.length
+    val fsBc = ctx.fsBc
 
-    def readBlock(pi: Int, is: InputStream, metrics: InputMetrics)
-      : Iterator[((Int, Int), BDM[Double])] = {
-      val block = RichDenseMatrixDouble.read(is, bufferSpec)
-      is.close()
+    val blocks =
+      new RDD[((Int, Int), BDM[Double])](ctx.backend.asSpark.sc, Nil) {
 
-      Iterator.single(gp.partCoordinates(pi) -> block)
-    }
+        case class FilePartition(index: Int, file: String) extends Partition
+
+        override lazy val getPartitions: Array[Partition] =
+          Array.tabulate(nPartitions)(i => FilePartition(i, partFiles(i)))
 
-    val blocks = HailContext.readPartitions(ctx, uri, partFiles, readBlock, Some(gp))
+        override def compute(split: Partition, context: TaskContext)
+          : Iterator[((Int, Int), BDM[Double])] =
+          using(fsBc.value.open(uri + "/parts/" + split.asInstanceOf[FilePartition].file)) { in =>
+            val block = RichDenseMatrixDouble.read(in, bufferSpec)
+            Iterator.single(gp.partCoordinates(split.index) -> block)
+          }
+
+        @transient override val partitioner: Option[Partitioner] =
+          Some(gp)
+      }
 
     new BlockMatrix(blocks, blockSize, nRows, nCols)
   }