whoops

hail-is · Mar 20, 2018 · ae2272b · ae2272b
1 parent 1a06b9e
commit ae2272b
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 34 deletions.
diff --git a/src/main/scala/is/hail/methods/LDPrune.scala b/src/main/scala/is/hail/methods/LDPrune.scala
@@ -492,6 +492,6 @@ object LDPrune {
     val ((globalPrunedRDD, nVariantsFinal), globalDuration) = time(pruneGlobal(rddLP2, r2Threshold, windowSize))
     info(s"LD prune step 3 of 3: nVariantsKept=$nVariantsFinal, time=${ formatTime(globalDuration) }")
 
-    vsm.copy2(rvd = vsm.rvd.copy(rdd = vsm.rvd.orderedJoinDistinct(globalPrunedRDD, "inner").map(_.rvLeft)))
+    vsm.copy2(rvd = vsm.rvd.orderedJoinDistinct(globalPrunedRDD, "inner", _.map(_.rvLeft), vsm.rvd.typ))
   }
 }
diff --git a/src/main/scala/is/hail/rvd/KeyedOrderedRVD.scala b/src/main/scala/is/hail/rvd/KeyedOrderedRVD.scala
@@ -52,7 +52,7 @@ class KeyedOrderedRVD(val rvd: OrderedRVD, val key: Array[String]) extends Seria
     new OrderedRVD(joinedType, newPartitioner, joinedRDD)
   }
 
-  def newOrderedJoinDistinct(
+  def orderedJoinDistinct(
     right: KeyedOrderedRVD,
     joinType: String,
     joiner: Iterator[JoinedRegionValue] => Iterator[RegionValue],
@@ -81,22 +81,6 @@ class KeyedOrderedRVD(val rvd: OrderedRVD, val key: Array[String]) extends Seria
     new OrderedRVD(joinedType, newPartitioner, joinedRDD)
   }
 
-  def orderedJoinDistinct(right: KeyedOrderedRVD, joinType: String): RDD[JoinedRegionValue] = {
-    checkJoinCompatability(right)
-    val rekeyedLTyp = new OrderedRVDType(typ.partitionKey, key, typ.rowType)
-    val rekeyedRTyp = new OrderedRVDType(right.typ.partitionKey, right.key, right.typ.rowType)
-
-    val repartitionedRight = right.rvd.constrainToOrderedPartitioner(right.typ, this.rvd.partitioner)
-    val compute: (OrderedRVIterator, OrderedRVIterator) => Iterator[JoinedRegionValue] =
-      (joinType: @unchecked) match {
-        case "inner" => _.innerJoinDistinct(_)
-        case "left" => _.leftJoinDistinct(_)
-      }
-    this.rvd.rdd.zipPartitions(repartitionedRight.rdd, true){ (leftIt, rightIt) =>
-      compute(OrderedRVIterator(rekeyedLTyp, leftIt), OrderedRVIterator(rekeyedRTyp, rightIt))
-    }
-  }
-
   def orderedZipJoin(right: KeyedOrderedRVD): RDD[JoinedRegionValue] = {
     val newPartitioner = rvd.partitioner.enlargeToRange(right.rvd.partitioner.range)
 

diff --git a/src/main/scala/is/hail/rvd/OrderedRVD.scala b/src/main/scala/is/hail/rvd/OrderedRVD.scala
@@ -121,8 +121,13 @@ class OrderedRVD(
   ): OrderedRVD =
     keyBy().orderedJoin(right.keyBy(), joinType, joiner, joinedType)
 
-  def orderedJoinDistinct(right: OrderedRVD, joinType: String): RDD[JoinedRegionValue] =
-    keyBy().orderedJoinDistinct(right.keyBy(), joinType)
+  def orderedJoinDistinct(
+    right: OrderedRVD,
+    joinType: String,
+    joiner: Iterator[JoinedRegionValue] => Iterator[RegionValue],
+    joinedType: OrderedRVDType
+  ): OrderedRVD =
+    keyBy().orderedJoinDistinct(right.keyBy(), joinType, joiner, joinedType)
 
   def orderedZipJoin(right: OrderedRVD): RDD[JoinedRegionValue] =
     keyBy().orderedZipJoin(right.keyBy())

diff --git a/src/main/scala/is/hail/variant/MatrixTable.scala b/src/main/scala/is/hail/variant/MatrixTable.scala
@@ -1032,7 +1032,7 @@ class MatrixTable(val hc: HailContext, val ast: MatrixIR) {
       }
     }
 
-    val joinedRVD = this.rvd.keyBy(rowKey.take(right.typ.key.length).toArray).newOrderedJoinDistinct(
+    val joinedRVD = this.rvd.keyBy(rowKey.take(right.typ.key.length).toArray).orderedJoinDistinct(
       right.keyBy(),
       "left",
       joiner,
@@ -1756,7 +1756,7 @@ class MatrixTable(val hc: HailContext, val ast: MatrixIR) {
     val localEntriesType = matrixType.entryArrayType
     assert(right.matrixType.entryArrayType == localEntriesType)
 
-    val joined = rvd.orderedJoinDistinct(right.rvd, "inner").mapPartitions({ it =>
+    val joiner: Iterator[JoinedRegionValue] => Iterator[RegionValue] = { it =>
       val rvb = new RegionValueBuilder()
       val rv2 = RegionValue()
 
@@ -1800,13 +1800,13 @@ class MatrixTable(val hc: HailContext, val ast: MatrixIR) {
         rv2.set(lrv.region, rvb.end())
         rv2
       }
-    }, preservesPartitioning = true)
+    }
 
     val newMatrixType = matrixType.copyParts() // move entries to the end
 
     copyMT(matrixType = newMatrixType,
       colValues = colValues ++ right.colValues,
-      rvd = OrderedRVD(rvd.typ, rvd.partitioner, joined))
+      rvd = rvd.orderedJoinDistinct(right.rvd, "inner", joiner, rvd.typ))
   }
 
   def makeKT(rowExpr: String, entryExpr: String, keyNames: Array[String] = Array.empty, seperator: String = "."): Table = {

diff --git a/src/test/scala/is/hail/variant/vsm/PartitioningSuite.scala b/src/test/scala/is/hail/variant/vsm/PartitioningSuite.scala
@@ -75,7 +75,7 @@ class PartitioningSuite extends SparkSuite {
     val mt = MatrixTable.fromRowsTable(Table.range(hc, 100, "idx", partitions=Some(6)))
     val orvdType = mt.matrixType.orvdType
 
-    mt.rvd.orderedJoinDistinct(OrderedRVD.empty(hc.sc, orvdType), "left").count()
-    mt.rvd.orderedJoinDistinct(OrderedRVD.empty(hc.sc, orvdType), "inner").count()
+    mt.rvd.orderedJoinDistinct(OrderedRVD.empty(hc.sc, orvdType), "left", _.map(_._1), orvdType).count()
+    mt.rvd.orderedJoinDistinct(OrderedRVD.empty(hc.sc, orvdType), "inner", _.map(_._1), orvdType).count()
   }
 }
diff --git a/src/test/scala/is/hail/vds/JoinSuite.scala b/src/test/scala/is/hail/vds/JoinSuite.scala
@@ -67,24 +67,23 @@ class JoinSuite extends SparkSuite {
     val localRowType = left.rvRowType
 
     // Inner distinct ordered join
-    val jInner = left.rvd.orderedJoinDistinct(right.rvd, "inner")
+    val jInner = left.rvd.orderedJoinDistinct(right.rvd, "inner", _.map(_._1), left.rvd.typ)
     val jInnerOrdRDD1 = left.rdd.join(right.rdd.distinct)
 
     assert(jInner.count() == jInnerOrdRDD1.count())
-    assert(jInner.forall(jrv => jrv.rvLeft != null && jrv.rvRight != null))
-    assert(jInner.map { jrv =>
-      val ur = new UnsafeRow(localRowType, jrv.rvLeft)
+    assert(jInner.map { rv =>
+      val ur = new UnsafeRow(localRowType, rv)
       ur.getAs[Locus](0)
     }.collect() sameElements jInnerOrdRDD1.map(_._1.asInstanceOf[Row].get(0)).collect().sorted(vType.ordering.toOrdering))
 
     // Left distinct ordered join
-    val jLeft = left.rvd.orderedJoinDistinct(right.rvd, "left")
+    val jLeft = left.rvd.orderedJoinDistinct(right.rvd, "left", _.map(_._1), left.rvd.typ)
     val jLeftOrdRDD1 = left.rdd.leftOuterJoin(right.rdd.distinct)
 
     assert(jLeft.count() == jLeftOrdRDD1.count())
-    assert(jLeft.forall(jrv => jrv.rvLeft != null))
-    assert(jLeft.map { jrv =>
-      val ur = new UnsafeRow(localRowType, jrv.rvLeft)
+    assert(jLeft.rdd.forall(rv => rv != null))
+    assert(jLeft.map { rv =>
+      val ur = new UnsafeRow(localRowType, rv)
       ur.getAs[Locus](0)
     }.collect() sameElements jLeftOrdRDD1.map(_._1.asInstanceOf[Row].get(0)).collect().sorted(vType.ordering.toOrdering))
   }