[query] Support indexed table reads in lowered execution

hail/python/hail/table.py

@@ -3309,7 +3309,7 @@ def _same(self, other, tolerance=1e-6, absolute=False):
             t = t.filter(~ _values_similar(t[left_value], t[right_value], tolerance, absolute))
             bad_rows = t.take(10)
             for r in bad_rows:
-                print(f'  Row mismatch:\n    L: {r[left_value]}\n    R: {r[right_value]}')
+                print(f'  Row mismatch at key={r._key}:\n    L: {r[left_value]}\n    R: {r[right_value]}')
             return False
 
         return True

hail/python/test/hail/table/test_table.py

@@ -783,7 +783,6 @@ def test_read_back_same_as_exported(self):
         t_read_back = hl.import_table(tmp_file, types=dict(t.row.dtype)).key_by('idx')
         self.assertTrue(t.select_globals()._same(t_read_back, tolerance=1e-4, absolute=True))
 
-    @fails_local_backend()
     def test_indexed_read(self):
         t = hl.utils.range_table(2000, 10)
         f = new_temp_file(extension='ht')
@@ -793,6 +792,7 @@ def test_indexed_read(self):
             hl.Interval(start=250, end=500, includes_start=True, includes_end=False),
         ])
         self.assertEqual(t2.n_partitions(), 2)
+        self.assertEqual(t2.count(), 350)
         self.assertTrue(t.filter((t.idx >= 150) & (t.idx < 500))._same(t2))
 
         t2 = hl.read_table(f, _intervals=[
@@ -1530,7 +1530,6 @@ def test_map_partitions_errors():
     with pytest.raises(ValueError, match='must preserve key fields'):
         ht._map_partitions(lambda rows: rows.map(lambda r: r.drop('idx')))
 
-@fails_local_backend()
 def test_map_partitions_indexed():
     tmp_file = new_temp_file()
     hl.utils.range_table(100, 8).write(tmp_file)

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

@@ -225,79 +225,12 @@ object HailContext {
     bounds: Option[Interval],
     metrics: InputMetrics = null
   ): Iterator[Long] =
-    if (bounds.isEmpty) {
-      idxr.close()
-      HailContext.readRowsPartition(makeDec)(ctx.r, in, metrics)
-    } else {
-      new Iterator[Long] {
-        private val region = ctx.region
-        private val idx = idxr.queryByInterval(bounds.get).buffered
-
-        private val trackedIn = new ByteTrackingInputStream(in)
-        private val field = offsetField.map { f =>
-          idxr.annotationType.asInstanceOf[TStruct].fieldIdx(f)
-        }
-        private val dec =
-          try {
-            if (idx.hasNext) {
-              val dec = makeDec(trackedIn)
-              val i = idx.head
-              val off = field.map { j =>
-                i.annotation.asInstanceOf[Row].getAs[Long](j)
-              }.getOrElse(i.recordOffset)
-              dec.seek(off)
-              dec
-            } else {
-              in.close()
-              null
-            }
-          } catch {
-            case e: Exception =>
-              idxr.close()
-              in.close()
-              throw e
-          }
-
-        private var cont: Byte = if (dec != null) dec.readByte() else 0
-        if (cont == 0) {
-          idxr.close()
-          if (dec != null) dec.close()
-        }
-
-        def hasNext: Boolean = cont != 0 && idx.hasNext
-
-        def next(): Long = {
-          if (!hasNext)
-            throw new NoSuchElementException("next on empty iterator")
-
-          try {
-            idx.next()
-            val res = dec.readRegionValue(region)
-            cont = dec.readByte()
-            if (metrics != null) {
-              ExposedMetrics.incrementRecord(metrics)
-              ExposedMetrics.incrementBytes(metrics, trackedIn.bytesReadAndClear())
-            }
-
-            if (cont == 0) {
-              dec.close()
-              idxr.close()
-            }
-
-            res
-          } catch {
-            case e: Exception =>
-              dec.close()
-              idxr.close()
-              throw e
-          }
-        }
-
-        override def finalize(): Unit = {
-          idxr.close()
-          if (dec != null) dec.close()
-        }
-      }
+    bounds match {
+      case Some(b) =>
+        new IndexReadIterator(makeDec, ctx.r, in, idxr, offsetField.orNull, b, metrics)
+      case None =>
+        idxr.close()
+        HailContext.readRowsPartition(makeDec)(ctx.r, in, metrics)
     }
 
   def readSplitRowsPartition(

hail/src/main/scala/is/hail/asm4s/Code.scala

@@ -188,6 +188,10 @@ object Code {
   )(implicit a1ct: ClassTag[A1], a2ct: ClassTag[A2], a3ct: ClassTag[A3], a4ct: ClassTag[A4], a5ct: ClassTag[A5], tct: ClassTag[T], tti: TypeInfo[T]): Code[T] =
     newInstance[T](Array[Class[_]](a1ct.runtimeClass, a2ct.runtimeClass, a3ct.runtimeClass, a4ct.runtimeClass, a5ct.runtimeClass), Array[Code[_]](a1, a2, a3, a4, a5))
 
+  def newInstance7[T <: AnyRef, A1, A2, A3, A4, A5, A6, A7](a1: Code[A1], a2: Code[A2], a3: Code[A3], a4: Code[A4], a5: Code[A5], a6: Code[A6], a7: Code[A7]
+  )(implicit a1ct: ClassTag[A1], a2ct: ClassTag[A2], a3ct: ClassTag[A3], a4ct: ClassTag[A4], a5ct: ClassTag[A5], a6ct: ClassTag[A6], a7ct: ClassTag[A7], tct: ClassTag[T], tti: TypeInfo[T]): Code[T] =
+    newInstance[T](Array[Class[_]](a1ct.runtimeClass, a2ct.runtimeClass, a3ct.runtimeClass, a4ct.runtimeClass, a5ct.runtimeClass, a6ct.runtimeClass, a7ct.runtimeClass), Array[Code[_]](a1, a2, a3, a4, a5, a6, a7))
+
   def newArray[T](size: Code[Int])(implicit tti: TypeInfo[T]): Code[Array[T]] =
     Code(size, lir.newArray(tti))
 

hail/src/main/scala/is/hail/expr/ir/Compile.scala

@@ -197,14 +197,14 @@ object CompileIterator {
     InferPType(ir, Env.empty[PType])
     val returnType = ir.pType.asInstanceOf[PStream].elementType.asInstanceOf[PStruct].setRequired(true)
 
-    val optStream = EmitStream.emit(emitter, ir, stepF, outerRegion, Env.empty, None)
+    val optStream = EmitStream.emit(ctx, emitter, ir, stepF, outerRegion, Env.empty, None)
 
     val elementAddress = stepF.genFieldThisRef[Long]("elementAddr")
 
     val didSetup = stepF.genFieldThisRef[Boolean]("didSetup")
     stepF.cb.emitInit(didSetup := false)
 
-    implicit val ecc: EmitStreamContext = EmitStreamContext(stepF)
+    implicit val ecc: EmitStreamContext = EmitStreamContext(stepF, ctx)
 
     val pullLabel = CodeLabel()
     val eosField = stepF.genFieldThisRef[Boolean]("eos")

hail/src/main/scala/is/hail/expr/ir/Emit.scala

@@ -456,7 +456,7 @@ class Emit[C](
       this.emit(ir, mb, region, env, container, loopEnv)
 
     def emitStream(ir: IR, outerRegion: ParentStagedRegion, mb: EmitMethodBuilder[C] = mb): EmitCode =
-      EmitStream.emit(this, ir, mb, outerRegion, env, container)
+      EmitStream.emit(ctx, this, ir, mb, outerRegion, env, container)
 
     def emitVoid(ir: IR, cb: EmitCodeBuilder = cb, mb: EmitMethodBuilder[C] = mb, region: StagedRegion = region, env: E = env, container: Option[AggContainer] = container, loopEnv: Option[Env[LoopRef]] = loopEnv): Unit =
       this.emitVoid(cb, ir, mb, region, env, container, loopEnv)
@@ -512,7 +512,7 @@ class Emit[C](
           {},
           { s =>
             cb += eltRegion.allocateRegion(Region.REGULAR)
-            cb += s.asStream.stream.getStream(eltRegion).forEach(mb, forBody)
+            cb += s.asStream.stream.getStream(eltRegion).forEach(ctx, mb, forBody)
             cb += eltRegion.free()
           })
 
@@ -638,7 +638,7 @@ class Emit[C](
       this.emitI(ir, cb, region, env, container, loopEnv)
 
     def emitStream(ir: IR, outerRegion: ParentStagedRegion): IEmitCode =
-      EmitStream.emit(this, ir, mb, outerRegion, env, container).toI(cb)
+      EmitStream.emit(ctx, this, ir, mb, outerRegion, env, container).toI(cb)
 
     def emitVoid(ir: IR, env: E = env, container: Option[AggContainer] = container, loopEnv: Option[Env[LoopRef]] = loopEnv): Unit =
       this.emitVoid(cb, ir: IR, mb, region, env, container, loopEnv)
@@ -1423,7 +1423,7 @@ class Emit[C](
           cb += tmpRegion.allocateRegion(Region.REGULAR)
           cb.assign(xAcc, emitI(zero, eltRegion).map(cb)(_.castTo(mb, eltRegion.code, accType)))
 
-          stream.asStream.stream.getStream(eltRegion).forEachI(cb, { elt =>
+          stream.asStream.stream.getStream(eltRegion).forEachI(ctx, cb, { elt =>
             // pre- and post-condition: 'xAcc' contains current accumulator,
             // whose heap memory is contained in 'eltRegion'. 'tmpRegion' is
             // empty.
@@ -1466,7 +1466,7 @@ class Emit[C](
           (accVars, acc).zipped.foreach { case (xAcc, (_, x)) =>
             cb.assign(xAcc, emitI(x, eltRegion).map(cb)(_.castTo(mb, eltRegion.code, xAcc.pt)))
           }
-          stream.asStream.stream.getStream(eltRegion).forEachI(cb, { elt =>
+          stream.asStream.stream.getStream(eltRegion).forEachI(ctx, cb, { elt =>
             // pre- and post-condition: 'accVars' contain current accumulators,
             // all of whose heap memory is contained in 'eltRegion'. 'tmpRegion'
             // is empty.
@@ -1544,7 +1544,7 @@ class Emit[C](
           .get(cb, "rows stream was missing in shuffle write")
           .asStream.stream.getStream(eltRegion)
         cb += eltRegion.allocateRegion(Region.REGULAR)
-        cb += rows.forEach(mb, { row: EmitCode =>
+        cb += rows.forEach(ctx, mb, { row: EmitCode =>
           Code(
             row.setup,
             row.m.mux(
@@ -1655,7 +1655,7 @@ class Emit[C](
     }
 
     def emitStream(ir: IR, outerRegion: ParentStagedRegion): EmitCode =
-      EmitStream.emit(this, ir, mb, outerRegion, env, container)
+      EmitStream.emit(ctx, this, ir, mb, outerRegion, env, container)
 
     def emitNDArrayColumnMajorStrides(ir: IR): EmitCode = {
       EmitCode.fromI(mb) { cb =>
@@ -1831,7 +1831,7 @@ class Emit[C](
         val optStream = emitStream(array, outerRegion)
         optStream.map { stream =>
           PCode(pt, Code(
-            EmitStream.write(mb, stream.asStream, vab, outerRegion),
+            EmitStream.write(ctx, mb, stream.asStream, vab, outerRegion),
             sort,
             distinct,
             sorter.toRegion()))
@@ -1844,7 +1844,7 @@ class Emit[C](
       case ToArray(a) =>
         val outerRegion = region.asParent(coerce[PStream](a.pType).separateRegions, "ToArray")
         emitStream(a, outerRegion).map { stream =>
-          EmitStream.toArray(mb, coerce[PArray](pt), stream.asStream, outerRegion)
+          EmitStream.toArray(ctx, mb, coerce[PArray](pt), stream.asStream, outerRegion)
         }
 
       case x@LowerBoundOnOrderedCollection(orderedCollection, elem, onKey) =>
@@ -1927,7 +1927,7 @@ class Emit[C](
         val outerRegion = region.asParent(atyp.separateRegions, "GroupByKey")
         emitStream(collection, outerRegion).map { stream =>
           PCode(pt, Code(
-            EmitStream.write(mb, stream.asStream, eab, outerRegion),
+            EmitStream.write(ctx, mb, stream.asStream, eab, outerRegion),
             sorter.sort(sortF),
             sorter.pruneMissing,
             eab.size.ceq(0).mux(
@@ -1998,7 +1998,7 @@ class Emit[C](
                   count := 0,
                   setup,
                   eltRegion.allocateRegion(Region.REGULAR),
-                  stream(eltRegion).forEach(mb, _ => Code(count := count + 1, eltRegion.clear())),
+                  stream(eltRegion).forEach(ctx, mb, _ => Code(count := count + 1, eltRegion.clear())),
                   eltRegion.free(),
                   count.get
                 )
@@ -2322,7 +2322,7 @@ class Emit[C](
               setup,
               ctxab.invoke[Int, Unit]("ensureCapacity", len.getOrElse(16)),
               eltRegion.allocateRegion(Region.REGULAR),
-              stream(eltRegion).map(etToTuple(_, ctxType)).forEach(mb, { offset =>
+              stream(eltRegion).map(etToTuple(_, ctxType)).forEach(ctx, mb, { offset =>
                 Code(
                   baos.invoke[Unit]("reset"),
                   Code.memoize(offset, "cda_add_contexts_addr") { offset =>
@@ -2374,7 +2374,7 @@ class Emit[C](
           decodeResult))
         }
 
-        COption.toEmitCode(optRes, mb)
+        COption.toEmitCode(ctx, optRes, mb)
 
       case x@TailLoop(name, args, body) =>
         val label = CodeLabel()
@@ -2412,14 +2412,14 @@ class Emit[C](
           // dead code
           const(true), pt.defaultValue)
 
-      case x@WritePartition(stream, ctx, writer) =>
-        val ctxCode = emit(ctx)
+      case x@WritePartition(stream, pctx, writer) =>
+        val ctxCode = emit(pctx)
         val streamType = coerce[PStream](stream.pType)
         val eltType = coerce[PStruct](streamType.elementType)
         val outerRegion = region.asParent(streamType.separateRegions, "WritePartition")
-        COption.toEmitCode(
+        COption.toEmitCode(ctx,
           COption.fromEmitCode(emitStream(stream, outerRegion)).flatMap { s =>
-            COption.fromEmitCode(writer.consumeStream(ctxCode, eltType, mb, outerRegion, s.asStream.stream))
+            COption.fromEmitCode(writer.consumeStream(ctx, ctxCode, eltType, mb, outerRegion, s.asStream.stream))
           }, mb)
 
       case x =>