Refactor BlockMatrixSparsity

Author: patrick-schultz

hail/hail/src-2.12/is/hail/utils/compat/immutable/ArraySeq.scala

@@ -23,4 +23,7 @@ object ArraySeq {
 
   implicit def canBuildFrom[T: ClassTag]: CanBuildFrom[ArraySeq[_], T, ArraySeq[T]] =
     A.canBuildFrom
+
+  def tabulate[T: ClassTag](n: Int)(f: Int => T): ArraySeq[T] =
+    A.unsafeWrapArray(Array.tabulate(n)(f))
 }

hail/hail/src-2.12/is/hail/utils/compat/package.scala

@@ -3,9 +3,20 @@ package is.hail.utils
 import is.hail.utils.compat.immutable.ArraySeq
 
 import scala.collection.compat.Factory
+import scala.collection.mutable
 import scala.reflect.ClassTag
 
 package object compat {
   implicit def arraySeqbf[A: ClassTag](ob: ArraySeq.type): Factory[A, ArraySeq[A]] =
     ob.canBuildFrom[A]
+
+  implicit class ArrayOps[A](private val a: Array[A]) extends AnyVal {
+    def sortInPlace[B >: A]()(implicit ct: ClassTag[B], ord: Ordering[B]): mutable.WrappedArray[A] = {
+      scala.util.Sorting.stableSort(a.asInstanceOf[Array[B]])
+      mutable.WrappedArray.make(a)
+    }
+
+    def sortInPlaceBy[B](f: A => B)(implicit ord: Ordering[B], ct: ClassTag[A]): mutable.WrappedArray[A] =
+      sortInPlace()(ct, ord on f)
+  }
 }

hail/hail/src/is/hail/expr/ir/BlockMatrixIR.scala

@@ -8,7 +8,7 @@ import is.hail.expr.ir.defs.{MetadataWriter, Str, UUID4, WriteMetadata, WriteVal
 import is.hail.expr.ir.lowering.{BlockMatrixStage2, LowererUnsupportedOperation}
 import is.hail.io.{StreamBufferSpec, TypedCodecSpec}
 import is.hail.io.fs.FS
-import is.hail.linalg.{BlockMatrix, BlockMatrixMetadata}
+import is.hail.linalg.{BlockMatrix, BlockMatrixMetadata, MatrixSparsity}
 import is.hail.types.TypeWithRequiredness
 import is.hail.types.encoded.{EBlockMatrixNDArray, ENumpyBinaryNDArray, EType}
 import is.hail.types.virtual._
@@ -135,8 +135,11 @@ case class BlockMatrixNativeMetadataWriter(
     cb: EmitCodeBuilder,
     region: Value[Region],
   ): Unit = {
-    val metaHelper =
-      BMMetadataHelper(path, typ.blockSize, typ.nRows, typ.nCols, typ.linearizedDefinedBlocks)
+    val partIdxToBlockIdx = typ.sparsity match {
+      case _: MatrixSparsity.Dense => None
+      case x: MatrixSparsity.Sparse => Some(x.definedBlocksColMajorLinear)
+    }
+    val metaHelper = BMMetadataHelper(path, typ.blockSize, typ.nRows, typ.nCols, partIdxToBlockIdx)
 
     val pc = writeAnnotations.getOrFatal(cb, "write annotations can't be missing!").asIndexable
     val partFiles = cb.newLocal[Array[String]]("partFiles")

hail/hail/src/is/hail/expr/ir/BlockMatrixWriter.scala

@@ -512,6 +512,8 @@ package defs {
     def <=(other: IR): IR = ApplyComparisonOp(LTEQ, self, other)
 
     def >=(other: IR): IR = ApplyComparisonOp(GTEQ, self, other)
+
+    def log(messages: AnyRef*): IR = logIR(self, messages: _*)
   }
 
   object ErrorIDs {

hail/hail/src/is/hail/expr/ir/IR.scala

@@ -14,7 +14,7 @@ import is.hail.io.gen.{BgenWriter, ExportGen}
 import is.hail.io.index.StagedIndexWriter
 import is.hail.io.plink.{BitPacker, ExportPlink}
 import is.hail.io.vcf.{ExportVCF, TabixVCF}
-import is.hail.linalg.BlockMatrix
+import is.hail.linalg.{BlockMatrix, MatrixSparsity}
 import is.hail.rvd.{IndexSpec, RVDPartitioner, RVDSpecMaker}
 import is.hail.types._
 import is.hail.types.encoded.{EBaseStruct, EBlockMatrixNDArray, EType}
@@ -2340,7 +2340,7 @@ case class MatrixBlockMatrixWriter(
 
     val countColumnsIR = ArrayLen(GetField(ts.getGlobals(), colsFieldName))
     val numCols: Int = CompileAndEvaluate[Int](ctx, countColumnsIR)
-    val numBlockCols: Int = (numCols - 1) / blockSize + 1
+    val numBlockCols: Int = BlockMatrixType.numBlocks(numCols.toLong, blockSize)
     val lastBlockNumCols = (numCols - 1) % blockSize + 1
 
     val rowCountIR = ts.mapCollect("matrix_block_matrix_writer_partition_counts")(paritionIR =>
@@ -2353,7 +2353,7 @@ case class MatrixBlockMatrixWriter(
     val inputPartStops = inputPartStartsPlusLast.tail
 
     val numRows = inputPartStartsPlusLast.last
-    val numBlockRows: Int = (numRows.toInt - 1) / blockSize + 1
+    val numBlockRows: Int = BlockMatrixType.numBlocks(numRows, blockSize)
 
     // Zip contexts with partition starts and ends
     val zippedWithStarts = ts.mapContexts { oldContextsStream =>
@@ -2510,7 +2510,7 @@ case class MatrixBlockMatrixWriter(
       numRows,
       numCols.toLong,
       blockSize,
-      BlockMatrixSparsity.dense,
+      MatrixSparsity.dense(numBlockRows, numBlockCols),
     )
     RelationalWriter.scoped(path, overwrite, None)(WriteMetadata(
       flatPaths,

hail/hail/src/is/hail/expr/ir/MatrixWriter.scala

@@ -311,6 +311,48 @@ object ArrayFunctions extends RegistryFunctions {
       ToArray(flatMapIR(ToStream(a))(ToStream(_)))
     }
 
+    registerSCode3t(
+      "scatter",
+      Array(tv("T")),
+      TArray(tv("T")),
+      TArray(TInt32),
+      TInt32,
+      TArray(tv("T")),
+      (_, a, _, _) => PCanonicalArray(a.asInstanceOf[SContainer].elementType.storageType()).sType,
+    ) {
+      case (
+            er,
+            cb,
+            _,
+            rt: SIndexablePointer,
+            elts: SIndexableValue,
+            indices: SIndexableValue,
+            len: SInt32Value,
+            errorID,
+          ) =>
+        val pt = rt.pType.asInstanceOf[PCanonicalArray]
+        val (push, finish) =
+          pt.constructFromIndicesUnsafe(cb, er.region, len.value, deepCopy = false)
+        cb.if_(
+          elts.loadLength.cne(indices.loadLength),
+          cb._fatalWithError(errorID, "scatter: values and indices arrays have different lengths"),
+        )
+        indices.forEachDefined(cb) { case (cb, pos, idx: SInt32Value) =>
+          cb.if_(
+            idx.value < 0 || idx.value >= len.value,
+            cb._fatalWithError(
+              errorID,
+              "scatter: indices array contained index ",
+              idx.value.toS,
+              ", which is greater than result length ",
+              len.value.toS,
+            ),
+          )
+          push(cb, idx.value, elts.loadElement(cb, pos))
+        }
+        finish(cb)
+    }
+
     registerSCode4(
       "lowerBound",
       TArray(tv("T")),

hail/hail/src/is/hail/expr/ir/functions/ArrayFunctions.scala

@@ -816,6 +816,31 @@ abstract class RegistryFunctions {
       case (r, cb, _, rt, Array(a1, a2, a3), errorID) => impl(r, cb, rt, a1, a2, a3, errorID)
     }
 
+  def registerSCode3t(
+    name: String,
+    typeParams: Array[Type],
+    mt1: Type,
+    mt2: Type,
+    mt3: Type,
+    rt: Type,
+    pt: (Type, SType, SType, SType) => SType,
+  )(
+    impl: (
+      EmitRegion,
+      EmitCodeBuilder,
+      Seq[Type],
+      SType,
+      SValue,
+      SValue,
+      SValue,
+      Value[Int],
+    ) => SValue
+  ): Unit =
+    registerSCode(name, Array(mt1, mt2, mt3), rt, unwrappedApply(pt), typeParams) {
+      case (r, cb, typeParams, rt, Array(a1, a2, a3), errorID) =>
+        impl(r, cb, typeParams, rt, a1, a2, a3, errorID)
+    }
+
   def registerSCode4(
     name: String,
     mt1: Type,