Lower MatrixColsTable (hail-is#4968)

* broken checkpoint * Lower MatrixColsTable * Fix sorting * fix bug * use selectFields * fix fatal
bcajes · Dec 20, 2018 · 110f78f · 110f78f
1 parent ceff6db
commit 110f78f
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 10 deletions.
diff --git a/hail/src/main/scala/is/hail/expr/ir/IRBuilder.scala b/hail/src/main/scala/is/hail/expr/ir/IRBuilder.scala
@@ -17,6 +17,8 @@ object IRBuilder {
 
   implicit def intToProxy(i: Int): IRProxy = I32(i)
 
+  implicit def booleanToProxy(b: Boolean): IRProxy = if (b) True() else False()
+
   implicit def ref(s: Symbol): IRProxy = (env: E) =>
     Ref(s.name, env.lookup(s.name))
 
@@ -33,6 +35,9 @@ object IRBuilder {
   def irRange(start: IRProxy, end: IRProxy, step: IRProxy = 1): IRProxy = (env: E) =>
     ArrayRange(start(env), end(env), step(env))
 
+  def irArrayLen(a: IRProxy): IRProxy = (env: E) => ArrayLen(a(env))
+
+
   def irIf(cond: IRProxy)(cnsq: IRProxy)(altr: IRProxy): IRProxy = (env: E) =>
     If(cond(env), cnsq(env), altr(env))
 
@@ -49,6 +54,8 @@ object IRBuilder {
 
     def typ: TableType = tir.typ
 
+    def getGlobals: IR = TableGetGlobals(tir)
+
     def mapGlobals(newGlobals: IRProxy): TableIR =
       TableMapGlobals(tir, newGlobals(globalEnv))
 
@@ -120,10 +127,14 @@ object IRBuilder {
       ArrayMap(ir(env), f.s.name, f.body(env.bind(f.s.name -> eltType)))
     }
 
+    def sort(ascending: IRProxy, onKey: Boolean = false): IRProxy = (env: E) => ArraySort(ir(env), ascending(env), onKey)
+
     def groupByKey: IRProxy = (env: E) => GroupByKey(ir(env))
 
     def toArray: IRProxy = (env: E) => ToArray(ir(env))
 
+    def parallelize(nPartitions: Option[Int] = None): TableIR = TableParallelize(ir(Env.empty), nPartitions)
+
     private[ir] def apply(env: E): IR = ir(env)
   }
 

diff --git a/hail/src/main/scala/is/hail/expr/ir/LowerMatrixIR.scala b/hail/src/main/scala/is/hail/expr/ir/LowerMatrixIR.scala
@@ -1,8 +1,8 @@
 package is.hail.expr.ir
 
 import is.hail.expr.types._
-import is.hail.expr.types.virtual.TArray
-import is.hail.utils.FastSeq
+import is.hail.expr.types.virtual.{TArray, TInt32}
+import is.hail.utils._
 
 object LowerMatrixIR {
   val entriesFieldName = MatrixType.entriesIdentifier
@@ -32,7 +32,8 @@ object LowerMatrixIR {
 
   private[this] def lower(tir: TableIR): TableIR = {
     val lowered = tableRules.applyOrElse(tir, (tir: TableIR) => lowerChildren(tir).asInstanceOf[TableIR])
-    assert(lowered.typ == tir.typ)
+    if(lowered.typ != tir.typ)
+      fatal(s"lowering changed type:\n  before: ${tir.typ}\n  after: ${lowered.typ}")
     lowered
   }
 
@@ -196,5 +197,27 @@ object LowerMatrixIR {
       lower(child)
         .mapGlobals('global.dropFields(colsField))
         .mapRows('row.dropFields(entriesField))
+
+    case MatrixColsTable(child) =>
+      val colKey = child.typ.colKey
+      let(__cols_and_globals = lower(child).getGlobals) {
+        val sortedCols = if (colKey.isEmpty)
+          '__cols_and_globals (colsField)
+        else
+          irRange(0, irArrayLen('__cols_and_globals (colsField)), 1)
+            .map {
+              'i ~> let(__cols_element = '__cols_and_globals (colsField)('i)) {
+                makeStruct(
+                  // key struct
+                  '_1 -> '__cols_element.selectFields(colKey: _*),
+                  '_2 -> '__cols_element)
+              }
+            }
+            .sort(true, onKey = true)
+            .map {
+              'elt ~> 'elt ('_2)
+            }
+        makeStruct('rows -> sortedCols, 'global -> '__cols_and_globals.dropFields(colsField))
+      }.parallelize(None).keyBy(child.typ.colKey)
   }
 }
diff --git a/hail/src/main/scala/is/hail/expr/ir/TableIR.scala b/hail/src/main/scala/is/hail/expr/ir/TableIR.scala
@@ -900,13 +900,6 @@ case class MatrixColsTable(child: MatrixIR) extends TableIR {
   }
 
   val typ: TableType = child.typ.colsTableType
-
-  protected[ir] override def execute(hc: HailContext): TableValue = {
-    val mv = child.execute(hc)
-    val ctv = mv.colsTableValue
-    assert(ctv.typ == typ)
-    ctv
-  }
 }
 
 case class MatrixEntriesTable(child: MatrixIR) extends TableIR {