[query] Expose references via ExecuteContext

Author: ehigham

hail/src/main/scala/is/hail/HailFeatureFlags.scala

@@ -68,6 +68,9 @@ class HailFeatureFlags private (
     flags.update(flag, value)
   }
 
+  def +(feature: (String, String)): HailFeatureFlags =
+    new HailFeatureFlags(flags + (feature._1 -> feature._2))
+
   def get(flag: String): String = flags(flag)
 
   def lookup(flag: String): Option[String] =

hail/src/main/scala/is/hail/backend/Backend.scala

@@ -105,8 +105,6 @@ abstract class Backend extends Closeable {
   def lookupOrCompileCachedFunction[T](k: CodeCacheKey)(f: => CompiledFunction[T])
     : CompiledFunction[T]
 
-  def references: mutable.Map[String, ReferenceGenome]
-
   def lowerDistributedSort(
     ctx: ExecuteContext,
     stage: TableStage,
@@ -181,10 +179,8 @@ abstract class Backend extends Closeable {
   ): Array[Byte] =
     withExecuteContext { ctx =>
       jsonToBytes {
-        Extraction.decompose {
-          ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
-            xContigs, yContigs, mtContigs, parInput).toJSON
-        }(defaultJSONFormats)
+        ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
+          xContigs, yContigs, mtContigs, parInput).toJSON
       }
     }
 

hail/src/main/scala/is/hail/backend/ExecuteContext.scala

@@ -63,6 +63,7 @@ object ExecuteContext {
     tmpdir: String,
     localTmpdir: String,
     backend: Backend,
+    references: Map[String, ReferenceGenome],
     fs: FS,
     timer: ExecutionTimer,
     tempFileManager: TempFileManager,
@@ -79,6 +80,7 @@ object ExecuteContext {
           tmpdir,
           localTmpdir,
           backend,
+          references,
           fs,
           region,
           timer,
@@ -107,6 +109,7 @@ class ExecuteContext(
   val tmpdir: String,
   val localTmpdir: String,
   val backend: Backend,
+  val references: Map[String, ReferenceGenome],
   val fs: FS,
   val r: Region,
   val timer: ExecutionTimer,
@@ -128,7 +131,7 @@ class ExecuteContext(
         )
     }
 
-  def stateManager = HailStateManager(backend.references.toMap)
+  val stateManager = HailStateManager(references)
 
   val tempFileManager: TempFileManager =
     if (_tempFileManager != null) _tempFileManager else new OwningTempFileManager(fs)
@@ -154,8 +157,6 @@ class ExecuteContext(
 
   def getFlag(name: String): String = flags.get(name)
 
-  def getReference(name: String): ReferenceGenome = backend.references(name)
-
   def shouldWriteIRFiles(): Boolean = getFlag("write_ir_files") != null
 
   def shouldNotLogIR(): Boolean = flags.get("no_ir_logging") != null
@@ -174,6 +175,7 @@ class ExecuteContext(
     tmpdir: String = this.tmpdir,
     localTmpdir: String = this.localTmpdir,
     backend: Backend = this.backend,
+    references: Map[String, ReferenceGenome] = this.references,
     fs: FS = this.fs,
     r: Region = this.r,
     timer: ExecutionTimer = this.timer,
@@ -189,6 +191,7 @@ class ExecuteContext(
       tmpdir,
       localTmpdir,
       backend,
+      references,
       fs,
       r,
       timer,

hail/src/main/scala/is/hail/backend/local/LocalBackend.scala

@@ -86,6 +86,7 @@ class LocalBackend(
   override val references: mutable.Map[String, ReferenceGenome],
 ) extends Backend with BackendWithCodeCache with Py4JBackendExtensions {
 
+  override def backend: Backend = this
   override val flags: HailFeatureFlags = HailFeatureFlags.fromEnv()
   override def longLifeTempFileManager: TempFileManager = null
 
@@ -101,6 +102,7 @@ class LocalBackend(
         tmpdir,
         tmpdir,
         this,
+        references.toMap,
         fs,
         timer,
         null,

hail/src/main/scala/is/hail/backend/py4j/Py4JBackendExtensions.scala

@@ -10,12 +10,14 @@ import is.hail.expr.ir.{
 }
 import is.hail.expr.ir.IRParser.parseType
 import is.hail.expr.ir.functions.IRFunctionRegistry
+import is.hail.io.reference.{IndexedFastaSequenceFile, LiftOver}
 import is.hail.linalg.RowMatrix
 import is.hail.types.physical.PStruct
 import is.hail.types.virtual.{TArray, TInterval}
 import is.hail.utils.{defaultJSONFormats, log, toRichIterable, FastSeq, HailException, Interval}
 import is.hail.variant.ReferenceGenome
 
+import scala.collection.mutable
 import scala.jdk.CollectionConverters.{
   asScalaBufferConverter, mapAsScalaMapConverter, seqAsJavaListConverter,
 }
@@ -29,7 +31,10 @@ import org.json4s.Formats
 import org.json4s.jackson.{JsonMethods, Serialization}
 import sourcecode.Enclosing
 
-trait Py4JBackendExtensions { this: Backend =>
+trait Py4JBackendExtensions {
+  def backend: Backend
+  def references: mutable.Map[String, ReferenceGenome]
+  def persistedIR: mutable.Map[Int, BaseIR]
   def flags: HailFeatureFlags
   def longLifeTempFileManager: TempFileManager
 
@@ -59,7 +64,9 @@ trait Py4JBackendExtensions { this: Backend =>
     persistedIR.remove(id)
 
   def pyAddSequence(name: String, fastaFile: String, indexFile: String): Unit =
-    withExecuteContext(ctx => references(name).addSequence(ctx, fastaFile, indexFile))
+    backend.withExecuteContext { ctx =>
+      references(name).addSequence(IndexedFastaSequenceFile(ctx.fs, fastaFile, indexFile))
+    }
 
   def pyRemoveSequence(name: String): Unit =
     references(name).removeSequence()
@@ -74,7 +81,7 @@ trait Py4JBackendExtensions { this: Backend =>
     partitionSize: java.lang.Integer,
     entries: String,
   ): Unit =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       val rm = RowMatrix.readBlockMatrix(ctx.fs, pathIn, partitionSize)
       entries match {
         case "full" =>
@@ -112,7 +119,7 @@ trait Py4JBackendExtensions { this: Backend =>
     returnType: String,
     bodyStr: String,
   ): Unit = {
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       IRFunctionRegistry.registerIR(
         ctx,
         name,
@@ -126,10 +133,10 @@ trait Py4JBackendExtensions { this: Backend =>
   }
 
   def pyExecuteLiteral(irStr: String): Int =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       val ir = IRParser.parse_value_ir(irStr, IRParserEnvironment(ctx, persistedIR.toMap))
       assert(ir.typ.isRealizable)
-      execute(ctx, ir) match {
+      backend.execute(ctx, ir) match {
         case Left(_) => throw new HailException("Can't create literal")
         case Right((pt, addr)) =>
           val field = GetFieldByIdx(EncodedLiteral.fromPTypeAndAddress(pt, addr, ctx), 0)
@@ -158,13 +165,13 @@ trait Py4JBackendExtensions { this: Backend =>
   }
 
   def pyToDF(s: String): DataFrame =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       val tir = IRParser.parse_table_ir(s, IRParserEnvironment(ctx, irMap = persistedIR.toMap))
       Interpret(tir, ctx).toDF()
     }
 
   def pyReadMultipleMatrixTables(jsonQuery: String): util.List[MatrixIR] =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       log.info("pyReadMultipleMatrixTables: got query")
       val kvs = JsonMethods.parse(jsonQuery) match {
         case json4s.JObject(values) => values.toMap
@@ -193,19 +200,24 @@ trait Py4JBackendExtensions { this: Backend =>
     addReference(ReferenceGenome.fromJSON(jsonConfig))
 
   def pyRemoveReference(name: String): Unit =
-    references.remove(name)
+    removeReference(name)
 
   def pyAddLiftover(name: String, chainFile: String, destRGName: String): Unit =
-    withExecuteContext(ctx => references(name).addLiftover(ctx, chainFile, destRGName))
+    backend.withExecuteContext { ctx =>
+      references(name).addLiftover(references(destRGName), LiftOver(ctx.fs, chainFile))
+    }
 
   def pyRemoveLiftover(name: String, destRGName: String): Unit =
     references(name).removeLiftover(destRGName)
 
   private[this] def addReference(rg: ReferenceGenome): Unit =
     ReferenceGenome.addFatalOnCollision(references, FastSeq(rg))
 
+  private[this] def removeReference(name: String): Unit =
+    references -= name
+
   def parse_value_ir(s: String, refMap: java.util.Map[String, String]): IR =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       IRParser.parse_value_ir(
         s,
         IRParserEnvironment(ctx, irMap = persistedIR.toMap),
@@ -231,7 +243,7 @@ trait Py4JBackendExtensions { this: Backend =>
     }
 
   def loadReferencesFromDataset(path: String): Array[Byte] =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       val rgs = ReferenceGenome.fromHailDataset(ctx.fs, path)
       ReferenceGenome.addFatalOnCollision(references, rgs)
 
@@ -245,7 +257,7 @@ trait Py4JBackendExtensions { this: Backend =>
     f: ExecuteContext => T
   )(implicit E: Enclosing
   ): T =
-    withExecuteContext { ctx =>
+    backend.withExecuteContext { ctx =>
       val tempFileManager = longLifeTempFileManager
       if (selfContainedExecution && tempFileManager != null) f(ctx)
       else ctx.local(tempFileManager = NonOwningTempFileManager(tempFileManager))(f)

hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala

@@ -13,6 +13,7 @@ import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.functions.IRFunctionRegistry
 import is.hail.expr.ir.lowering._
 import is.hail.io.fs._
+import is.hail.io.reference.{IndexedFastaSequenceFile, LiftOver}
 import is.hail.linalg.BlockMatrix
 import is.hail.services.{BatchClient, JobGroupRequest, _}
 import is.hail.services.JobGroupStates.{Cancelled, Failure, Running, Success}
@@ -93,7 +94,16 @@ object ServiceBackend {
       rpcConfig.custom_references.map(ReferenceGenome.fromJSON),
     )
 
-    val backend = new ServiceBackend(
+    rpcConfig.liftovers.foreach { case (sourceGenome, liftoversForSource) =>
+      liftoversForSource.foreach { case (destGenome, chainFile) =>
+        references(sourceGenome).addLiftover(references(destGenome), LiftOver(fs, chainFile))
+      }
+    }
+    rpcConfig.sequences.foreach { case (rg, seq) =>
+      references(rg).addSequence(IndexedFastaSequenceFile(fs, seq.fasta, seq.index))
+    }
+
+    new ServiceBackend(
       JarUrl(jarLocation),
       name,
       theHailClassLoader,
@@ -106,27 +116,14 @@ object ServiceBackend {
       backendContext,
       scratchDir,
     )
-
-    backend.withExecuteContext { ctx =>
-      rpcConfig.liftovers.foreach { case (sourceGenome, liftoversForSource) =>
-        liftoversForSource.foreach { case (destGenome, chainFile) =>
-          references(sourceGenome).addLiftover(ctx, chainFile, destGenome)
-        }
-      }
-      rpcConfig.sequences.foreach { case (rg, seq) =>
-        references(rg).addSequence(ctx, seq.fasta, seq.index)
-      }
-    }
-
-    backend
   }
 }
 
 class ServiceBackend(
   val jarSpec: JarSpec,
   var name: String,
   val theHailClassLoader: HailClassLoader,
-  override val references: mutable.Map[String, ReferenceGenome],
+  val references: mutable.Map[String, ReferenceGenome],
   val batchClient: BatchClient,
   val batchConfig: BatchConfig,
   val flags: HailFeatureFlags,
@@ -397,6 +394,7 @@ class ServiceBackend(
         tmpdir,
         "file:///tmp",
         this,
+        references.toMap,
         fs,
         timer,
         null,

hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala

@@ -345,6 +345,7 @@ class SparkBackend(
     new HadoopFS(new SerializableHadoopConfiguration(conf))
   }
 
+  override def backend: Backend = this
   override val flags: HailFeatureFlags = HailFeatureFlags.fromEnv()
 
   override val longLifeTempFileManager: TempFileManager =
@@ -374,6 +375,7 @@ class SparkBackend(
       tmpdir,
       localTmpdir,
       this,
+      references.toMap,
       fs,
       region,
       timer,
@@ -393,6 +395,7 @@ class SparkBackend(
         tmpdir,
         localTmpdir,
         this,
+        references.toMap,
         fs,
         timer,
         null,

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

@@ -71,7 +71,7 @@ class EmitModuleBuilder(val ctx: ExecuteContext, val modb: ModuleBuilder) {
   }
 
   def referenceGenomes(): IndexedSeq[ReferenceGenome] =
-    rgContainers.keys.map(ctx.getReference(_)).toIndexedSeq.sortBy(_.name)
+    rgContainers.keys.map(ctx.references(_)).toIndexedSeq.sortBy(_.name)
 
   def referenceGenomeFields(): IndexedSeq[StaticField[ReferenceGenome]] =
     rgContainers.toFastSeq.sortBy(_._1).map(_._2)

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

@@ -646,7 +646,7 @@ class ExtractIntervalFilters(ctx: ExecuteContext, keyType: TStruct) {
           BoolValue.fromComparison(l, op).restrict(keySet)
         case Contig(rgStr) =>
           // locus contig equality comparison
-          val b = getIntervalFromContig(l.asInstanceOf[String], ctx.getReference(rgStr)) match {
+          val b = getIntervalFromContig(l.asInstanceOf[String], ctx.references(rgStr)) match {
             case Some(i) =>
               val b = BoolValue(
                 KeySet(i),
@@ -670,7 +670,7 @@ class ExtractIntervalFilters(ctx: ExecuteContext, keyType: TStruct) {
         case Position(rgStr) =>
           // locus position comparison
           val posBoolValue = BoolValue.fromComparison(l, op)
-          val rg = ctx.getReference(rgStr)
+          val rg = ctx.references(rgStr)
           val b = BoolValue(
             KeySet(liftPosIntervalsToLocus(posBoolValue.trueBound, rg, ctx)),
             KeySet(liftPosIntervalsToLocus(posBoolValue.falseBound, rg, ctx)),

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

@@ -210,7 +210,7 @@ case class MatrixValue(
       ReferenceGenome.exportReferences(
         fs,
         refPath,
-        ReferenceGenome.getReferences(t).map(ctx.getReference(_)),
+        ReferenceGenome.getReferences(t).map(ctx.references(_)),
       )
     }