[query] Extract Backend Methods called from Python into Py4JBackendExtensions

Author: ehigham

hail/python/hail/backend/py4j_backend.py

@@ -237,17 +237,17 @@ def _rpc(self, action, payload) -> Tuple[bytes, Optional[dict]]:
 
     def persist_expression(self, expr):
         t = expr.dtype
-        return construct_expr(JavaIR(t, self._jbackend.executeLiteral(self._render_ir(expr._ir))), t)
+        return construct_expr(JavaIR(t, self._jbackend.pyExecuteLiteral(self._render_ir(expr._ir))), t)
 
     def _is_registered_ir_function_name(self, name: str) -> bool:
         return name in self._registered_ir_function_names
 
     def set_flags(self, **flags: Mapping[str, str]):
-        available = self._jbackend.availableFlags()
+        available = self._jbackend.pyAvailableFlags()
         invalid = []
         for flag, value in flags.items():
             if flag in available:
-                self._jbackend.setFlag(flag, value)
+                self._jbackend.pySetFlag(flag, value)
             else:
                 invalid.append(flag)
         if len(invalid) != 0:
@@ -256,7 +256,7 @@ def set_flags(self, **flags: Mapping[str, str]):
             )
 
     def get_flags(self, *flags) -> Mapping[str, str]:
-        return {flag: self._jbackend.getFlag(flag) for flag in flags}
+        return {flag: self._jbackend.pyGetFlag(flag) for flag in flags}
 
     def _add_reference_to_scala_backend(self, rg):
         self._jbackend.pyAddReference(orjson.dumps(rg._config).decode('utf-8'))

hail/python/hail/ir/ir.py

@@ -3880,7 +3880,7 @@ def __del__(self):
         if Env._hc:
             backend = Env.backend()
             assert isinstance(backend, Py4JBackend)
-            backend._jbackend.removeJavaIR(self._id)
+            backend._jbackend.pyRemoveJavaIR(self._id)
 
 
 class JavaIR(IR):

hail/python/hail/ir/table_ir.py

@@ -1215,4 +1215,4 @@ def __del__(self):
         if Env._hc:
             backend = Env.backend()
             assert isinstance(backend, Py4JBackend)
-            backend._jbackend.removeJavaIR(self._id)
+            backend._jbackend.pyRemoveJavaIR(self._id)

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

@@ -1,12 +1,12 @@
 package is.hail.backend
 
 import is.hail.asm4s._
+import is.hail.backend.Backend.jsonToBytes
 import is.hail.backend.spark.SparkBackend
 import is.hail.expr.ir.{
   BaseIR, CodeCacheKey, CompiledFunction, IR, IRParser, IRParserEnvironment, LoweringAnalyses,
   SortField, TableIR, TableReader,
 }
-import is.hail.expr.ir.functions.IRFunctionRegistry
 import is.hail.expr.ir.lowering.{TableStage, TableStageDependency}
 import is.hail.io.{BufferSpec, TypedCodecSpec}
 import is.hail.io.fs._
@@ -20,7 +20,6 @@ import is.hail.types.virtual.{BlockMatrixType, TFloat64}
 import is.hail.utils._
 import is.hail.variant.ReferenceGenome
 
-import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.reflect.ClassTag
 
@@ -29,7 +28,7 @@ import java.nio.charset.StandardCharsets
 
 import com.fasterxml.jackson.core.StreamReadConstraints
 import org.json4s._
-import org.json4s.jackson.{JsonMethods, Serialization}
+import org.json4s.jackson.JsonMethods
 import sourcecode.Enclosing
 
 object Backend {
@@ -41,13 +40,6 @@ object Backend {
     s"hail_query_$id"
   }
 
-  private var irID: Int = 0
-
-  def nextIRID(): Int = {
-    irID += 1
-    irID
-  }
-
   def encodeToOutputStream(
     ctx: ExecuteContext,
     t: PTuple,
@@ -66,6 +58,9 @@ object Backend {
     assert(t.isFieldDefined(off, 0))
     codec.encode(ctx, elementType, t.loadField(off, 0), os)
   }
+
+  def jsonToBytes(f: => JValue): Array[Byte] =
+    JsonMethods.compact(f).getBytes(StandardCharsets.UTF_8)
 }
 
 abstract class BroadcastValue[T] { def value: T }
@@ -89,14 +84,6 @@ abstract class Backend extends Closeable {
 
   val persistedIR: mutable.Map[Int, BaseIR] = mutable.Map()
 
-  protected[this] def addJavaIR(ir: BaseIR): Int = {
-    val id = Backend.nextIRID()
-    persistedIR += (id -> ir)
-    id
-  }
-
-  def removeJavaIR(id: Int): Unit = persistedIR.remove(id)
-
   def defaultParallelism: Int
 
   def canExecuteParallelTasksOnDriver: Boolean = true
@@ -131,30 +118,7 @@ abstract class Backend extends Closeable {
   def lookupOrCompileCachedFunction[T](k: CodeCacheKey)(f: => CompiledFunction[T])
     : CompiledFunction[T]
 
-  var references: Map[String, ReferenceGenome] = Map.empty
-
-  def addDefaultReferences(): Unit =
-    references = ReferenceGenome.builtinReferences()
-
-  def addReference(rg: ReferenceGenome): Unit = {
-    references.get(rg.name) match {
-      case Some(rg2) =>
-        if (rg != rg2) {
-          fatal(
-            s"Cannot add reference genome '${rg.name}', a different reference with that name already exists. Choose a reference name NOT in the following list:\n  " +
-              s"@1",
-            references.keys.truncatable("\n  "),
-          )
-        }
-      case None =>
-        references += (rg.name -> rg)
-    }
-  }
-
-  def hasReference(name: String) = references.contains(name)
-
-  def removeReference(name: String): Unit =
-    references -= name
+  def references: mutable.Map[String, ReferenceGenome]
 
   def lowerDistributedSort(
     ctx: ExecuteContext,
@@ -189,9 +153,6 @@ abstract class Backend extends Closeable {
 
   def withExecuteContext[T](f: ExecuteContext => T)(implicit E: Enclosing): T
 
-  private[this] def jsonToBytes(f: => JValue): Array[Byte] =
-    JsonMethods.compact(f).getBytes(StandardCharsets.UTF_8)
-
   final def valueType(s: String): Array[Byte] =
     jsonToBytes {
       withExecuteContext { ctx =>
@@ -220,15 +181,7 @@ abstract class Backend extends Closeable {
       }
     }
 
-  def loadReferencesFromDataset(path: String): Array[Byte] = {
-    withExecuteContext { ctx =>
-      val rgs = ReferenceGenome.fromHailDataset(ctx.fs, path)
-      rgs.foreach(addReference)
-
-      implicit val formats: Formats = defaultJSONFormats
-      Serialization.write(rgs.map(_.toJSON).toFastSeq).getBytes(StandardCharsets.UTF_8)
-    }
-  }
+  def loadReferencesFromDataset(path: String): Array[Byte]
 
   def fromFASTAFile(
     name: String,
@@ -240,18 +193,22 @@ abstract class Backend extends Closeable {
     parInput: Array[String],
   ): Array[Byte] =
     withExecuteContext { ctx =>
-      val rg = ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
-        xContigs, yContigs, mtContigs, parInput)
-      rg.toJSONString.getBytes(StandardCharsets.UTF_8)
+      jsonToBytes {
+        Extraction.decompose {
+          ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
+            xContigs, yContigs, mtContigs, parInput).toJSON
+        }(defaultJSONFormats)
+      }
     }
 
-  def parseVCFMetadata(path: String): Array[Byte] = jsonToBytes {
+  def parseVCFMetadata(path: String): Array[Byte] =
     withExecuteContext { ctx =>
-      val metadata = LoadVCF.parseHeaderMetadata(ctx.fs, Set.empty, TFloat64, path)
-      implicit val formats = defaultJSONFormats
-      Extraction.decompose(metadata)
+      jsonToBytes {
+        Extraction.decompose {
+          LoadVCF.parseHeaderMetadata(ctx.fs, Set.empty, TFloat64, path)
+        }(defaultJSONFormats)
+      }
     }
-  }
 
   def importFam(path: String, isQuantPheno: Boolean, delimiter: String, missingValue: String)
     : Array[Byte] =
@@ -261,27 +218,6 @@ abstract class Backend extends Closeable {
       )
     }
 
-  def pyRegisterIR(
-    name: String,
-    typeParamStrs: java.util.ArrayList[String],
-    argNameStrs: java.util.ArrayList[String],
-    argTypeStrs: java.util.ArrayList[String],
-    returnType: String,
-    bodyStr: String,
-  ): Unit = {
-    withExecuteContext { ctx =>
-      IRFunctionRegistry.registerIR(
-        ctx,
-        name,
-        typeParamStrs.asScala.toArray,
-        argNameStrs.asScala.toArray,
-        argTypeStrs.asScala.toArray,
-        returnType,
-        bodyStr,
-      )
-    }
-  }
-
   def execute(ctx: ExecuteContext, ir: IR): Either[Unit, (PTuple, Long)]
 }
 

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

@@ -112,6 +112,7 @@ class BackendHttpHandler(backend: Backend) extends HttpHandler {
         }
         return
       }
+
       val response: Array[Byte] = exchange.getRequestURI.getPath match {
         case "/value/type" => backend.valueType(body.extract[IRTypePayload].ir)
         case "/table/type" => backend.tableType(body.extract[IRTypePayload].ir)

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

@@ -128,7 +128,7 @@ class ExecuteContext(
         )
     }
 
-  val stateManager = HailStateManager(backend.references)
+  def stateManager = HailStateManager(backend.references.toMap)
 
   val tempFileManager: TempFileManager =
     if (_tempFileManager != null) _tempFileManager else new OwningTempFileManager(fs)

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

@@ -4,8 +4,9 @@ import is.hail.{CancellingExecutorService, HailContext, HailFeatureFlags}
 import is.hail.annotations.Region
 import is.hail.asm4s._
 import is.hail.backend._
+import is.hail.backend.py4j.Py4JBackendExtensions
 import is.hail.expr.Validate
-import is.hail.expr.ir.{IRParser, _}
+import is.hail.expr.ir._
 import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.lowering._
 import is.hail.io.fs._
@@ -17,7 +18,7 @@ import is.hail.types.virtual.{BlockMatrixType, TVoid}
 import is.hail.utils._
 import is.hail.variant.ReferenceGenome
 
-import scala.collection.JavaConverters._
+import scala.collection.mutable
 import scala.reflect.ClassTag
 
 import java.io.PrintWriter
@@ -47,8 +48,11 @@ object LocalBackend {
     if (!skipLoggingConfiguration)
       HailContext.configureLogging(logFile, quiet, append)
 
-    theLocalBackend = new LocalBackend(tmpdir)
-    theLocalBackend.addDefaultReferences()
+    theLocalBackend = new LocalBackend(
+      tmpdir,
+      mutable.Map(ReferenceGenome.builtinReferences().toSeq: _*),
+    )
+
     theLocalBackend
   }
 
@@ -64,18 +68,15 @@ object LocalBackend {
   }
 }
 
-class LocalBackend(val tmpdir: String) extends Backend with BackendWithCodeCache {
-
-  private[this] val flags = HailFeatureFlags.fromEnv()
-  private[this] val theHailClassLoader = new HailClassLoader(getClass().getClassLoader())
+class LocalBackend(
+  val tmpdir: String,
+  override val references: mutable.Map[String, ReferenceGenome],
+) extends Backend with BackendWithCodeCache with Py4JBackendExtensions {
 
-  def getFlag(name: String): String = flags.get(name)
+  override val flags: HailFeatureFlags = HailFeatureFlags.fromEnv()
+  override def longLifeTempFileManager: TempFileManager = null
 
-  def setFlag(name: String, value: String) = flags.set(name, value)
-
-  // called from python
-  val availableFlags: java.util.ArrayList[String] =
-    flags.available
+  private[this] val theHailClassLoader = new HailClassLoader(getClass().getClassLoader())
 
   // flags can be set after construction from python
   def fs: FS = RouterFS.buildRoutes(CloudStorageFSConfig.fromFlagsAndEnv(None, flags))
@@ -190,81 +191,6 @@ class LocalBackend(val tmpdir: String) extends Backend with BackendWithCodeCache
       res
     }
 
-  def executeLiteral(irStr: String): Int =
-    withExecuteContext { ctx =>
-      val ir = IRParser.parse_value_ir(irStr, IRParserEnvironment(ctx, persistedIR.toMap))
-      assert(ir.typ.isRealizable)
-      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)
-          addJavaIR(field)
-      }
-    }
-
-  def pyAddReference(jsonConfig: String): Unit = addReference(ReferenceGenome.fromJSON(jsonConfig))
-  def pyRemoveReference(name: String): Unit = removeReference(name)
-
-  def pyAddLiftover(name: String, chainFile: String, destRGName: String): Unit =
-    withExecuteContext(ctx => references(name).addLiftover(ctx, chainFile, destRGName))
-
-  def pyRemoveLiftover(name: String, destRGName: String) =
-    references(name).removeLiftover(destRGName)
-
-  def pyFromFASTAFile(
-    name: String,
-    fastaFile: String,
-    indexFile: String,
-    xContigs: java.util.List[String],
-    yContigs: java.util.List[String],
-    mtContigs: java.util.List[String],
-    parInput: java.util.List[String],
-  ): String =
-    withExecuteContext { ctx =>
-      val rg = ReferenceGenome.fromFASTAFile(
-        ctx,
-        name,
-        fastaFile,
-        indexFile,
-        xContigs.asScala.toArray,
-        yContigs.asScala.toArray,
-        mtContigs.asScala.toArray,
-        parInput.asScala.toArray,
-      )
-      rg.toJSONString
-    }
-
-  def pyAddSequence(name: String, fastaFile: String, indexFile: String): Unit =
-    withExecuteContext(ctx => references(name).addSequence(ctx, fastaFile, indexFile))
-
-  def pyRemoveSequence(name: String) = references(name).removeSequence()
-
-  def parse_value_ir(s: String, refMap: java.util.Map[String, String]): IR =
-    withExecuteContext { ctx =>
-      IRParser.parse_value_ir(
-        s,
-        IRParserEnvironment(ctx, persistedIR.toMap),
-        BindingEnv.eval(refMap.asScala.toMap.map { case (n, t) =>
-          Name(n) -> IRParser.parseType(t)
-        }.toSeq: _*),
-      )
-    }
-
-  def parse_table_ir(s: String): TableIR =
-    withExecuteContext { ctx =>
-      IRParser.parse_table_ir(s, IRParserEnvironment(ctx, irMap = persistedIR.toMap))
-    }
-
-  def parse_matrix_ir(s: String): MatrixIR =
-    withExecuteContext { ctx =>
-      IRParser.parse_matrix_ir(s, IRParserEnvironment(ctx, irMap = persistedIR.toMap))
-    }
-
-  def parse_blockmatrix_ir(s: String): BlockMatrixIR =
-    withExecuteContext { ctx =>
-      IRParser.parse_blockmatrix_ir(s, IRParserEnvironment(ctx, irMap = persistedIR.toMap))
-    }
-
   override def lowerDistributedSort(
     ctx: ExecuteContext,
     stage: TableStage,