[query] Lift backend state into (Batch|Py4J)QueryDriver (#14698)

This change extracts common "driver" logic and state out of the various
backend classes into two new classes:
- `BatchQueryDriver`
- `Py4JQueryDriver`

These drivers handle requests from the python clients via their
respective rpc apis. The `BackendRpc` trait implements the core logic of
reading a command, dispatching to the backend for execution, and
responding to the user. Each driver implements the specifics of
- how requests are handled
  - over http in the case of py4j
  - via object storage for batch 
- how execute contexts are created
- how state is managed between requests

With this, driver state and lifecycle becomes decoupled from how queries
are executed (ie which "backend" implementation). It would be trivial to
use a `ServiceBackend` with the `Py4JQueryDriver` to test generated code
on batch without building and uploading a query jar.

Author: ehigham

hail/Makefile

@@ -348,7 +348,7 @@ upload-artifacts: $(WHEEL)
 # target must be run at least once a day. To trigger this target to re-run,
 # > rm upload-remote-test-resources
 upload-remote-test-resources: install-editable
-upload-remote-test-resources: $(shell git ls-files src/test/resources)
+upload-remote-test-resources: $(shell git ls-files hail/test/resources)
 upload-remote-test-resources: $(shell git ls-files python/hail/docs/data)
 	# # If hailtop.aiotools.copy gives you trouble:
 	# gcloud storage cp -r src/test/resources/\* $(CLOUD_HAIL_TEST_RESOURCES_DIR)
@@ -357,7 +357,7 @@ upload-remote-test-resources: $(shell git ls-files python/hail/docs/data)
             -vvv \
             'null' \
             "[\
-              {\"from\":\"src/test/resources\",\"to\":\"$(CLOUD_HAIL_TEST_RESOURCES_DIR)\"}, \
+              {\"from\":\"hail/test/resources\",\"to\":\"$(CLOUD_HAIL_TEST_RESOURCES_DIR)\"}, \
               {\"from\":\"python/hail/docs/data\",\"to\":\"$(CLOUD_HAIL_DOCTEST_DATA_DIR)\"} \
              ]" \
             --timeout 600

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

@@ -19,6 +19,7 @@ import org.apache.spark.executor.InputMetrics
 import org.apache.spark.rdd.RDD
 import org.json4s.Extraction
 import org.json4s.jackson.JsonMethods
+import sourcecode.Enclosing
 
 case class FilePartition(index: Int, file: String) extends Partition
 
@@ -41,7 +42,7 @@ object HailContext {
 
   def backend: Backend = get.backend
 
-  def sparkBackend(op: String): SparkBackend = get.sparkBackend(op)
+  def sparkBackend(implicit E: Enclosing): SparkBackend = get.backend.asSpark
 
   def configureLogging(logFile: String, quiet: Boolean, append: Boolean): Unit = {
     org.apache.log4j.helpers.LogLog.setInternalDebugging(true)
@@ -153,7 +154,7 @@ object HailContext {
 
     val fsBc = fs.broadcast
 
-    new RDD[T](SparkBackend.sparkContext("readPartition"), Nil) {
+    new RDD[T](SparkBackend.sparkContext, Nil) {
       def getPartitions: Array[Partition] =
         Array.tabulate(nPartitions)(i => FilePartition(i, partFiles(i)))
 
@@ -176,8 +177,6 @@ class HailContext private (
 ) {
   def stop(): Unit = HailContext.stop()
 
-  def sparkBackend(op: String): SparkBackend = backend.asSpark(op)
-
   var checkRVDKeys: Boolean = false
 
   def version: String = is.hail.HAIL_PRETTY_VERSION
@@ -189,7 +188,7 @@ class HailContext private (
     maxLines: Int,
   ): Map[String, Array[WithContext[String]]] = {
     val regexp = regex.r
-    SparkBackend.sparkContext("fileAndLineCounts").textFilesLines(fs.globAll(files).map(_.getPath))
+    SparkBackend.sparkContext.textFilesLines(fs.globAll(files).map(_.getPath))
       .filter(line => regexp.findFirstIn(line.value).isDefined)
       .take(maxLines)
       .groupBy(_.source.file)

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

@@ -1,28 +1,20 @@
 package is.hail.backend
 
-import is.hail.asm4s._
-import is.hail.backend.Backend.jsonToBytes
+import is.hail.asm4s.HailClassLoader
 import is.hail.backend.spark.SparkBackend
-import is.hail.expr.ir.{IR, IRParser, LoweringAnalyses, SortField, TableIR, TableReader}
+import is.hail.expr.ir.{IR, LoweringAnalyses, SortField, TableIR, TableReader}
 import is.hail.expr.ir.lowering.{TableStage, TableStageDependency}
 import is.hail.io.{BufferSpec, TypedCodecSpec}
-import is.hail.io.fs._
-import is.hail.io.plink.LoadPlink
-import is.hail.io.vcf.LoadVCF
-import is.hail.types._
+import is.hail.io.fs.FS
+import is.hail.types.RTable
 import is.hail.types.encoded.EType
 import is.hail.types.physical.PTuple
-import is.hail.types.virtual.TFloat64
-import is.hail.utils._
-import is.hail.variant.ReferenceGenome
+import is.hail.utils.fatal
 
 import scala.reflect.ClassTag
 
-import java.io._
-import java.nio.charset.StandardCharsets
+import java.io.{Closeable, OutputStream}
 
-import org.json4s._
-import org.json4s.jackson.JsonMethods
 import sourcecode.Enclosing
 
 object Backend {
@@ -38,23 +30,19 @@ object Backend {
     ctx: ExecuteContext,
     t: PTuple,
     off: Long,
-    bufferSpecString: String,
+    bufferSpec: BufferSpec,
     os: OutputStream,
   ): Unit = {
-    val bs = BufferSpec.parseOrDefault(bufferSpecString)
     assert(t.size == 1)
     val elementType = t.fields(0).typ
     val codec = TypedCodecSpec(
       EType.fromPythonTypeEncoding(elementType.virtualType),
       elementType.virtualType,
-      bs,
+      bufferSpec,
     )
     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 }
@@ -82,8 +70,8 @@ abstract class Backend extends Closeable {
     f: (Array[Byte], HailTaskContext, HailClassLoader, FS) => Array[Byte]
   ): (Option[Throwable], IndexedSeq[(Array[Byte], Int)])
 
-  def asSpark(op: String): SparkBackend =
-    fatal(s"${getClass.getSimpleName}: $op requires SparkBackend")
+  def asSpark(implicit E: Enclosing): SparkBackend =
+    fatal(s"${getClass.getSimpleName}: ${E.value} requires SparkBackend")
 
   def lowerDistributedSort(
     ctx: ExecuteContext,
@@ -116,70 +104,7 @@ abstract class Backend extends Closeable {
   def tableToTableStage(ctx: ExecuteContext, inputIR: TableIR, analyses: LoweringAnalyses)
     : TableStage
 
-  def withExecuteContext[T](f: ExecuteContext => T)(implicit E: Enclosing): T
-
-  final def valueType(s: String): Array[Byte] =
-    withExecuteContext { ctx =>
-      jsonToBytes {
-        IRParser.parse_value_ir(ctx, s).typ.toJSON
-      }
-    }
-
-  final def tableType(s: String): Array[Byte] =
-    withExecuteContext { ctx =>
-      jsonToBytes {
-        IRParser.parse_table_ir(ctx, s).typ.toJSON
-      }
-    }
-
-  final def matrixTableType(s: String): Array[Byte] =
-    withExecuteContext { ctx =>
-      jsonToBytes {
-        IRParser.parse_matrix_ir(ctx, s).typ.toJSON
-      }
-    }
-
-  final def blockMatrixType(s: String): Array[Byte] =
-    withExecuteContext { ctx =>
-      jsonToBytes {
-        IRParser.parse_blockmatrix_ir(ctx, s).typ.toJSON
-      }
-    }
-
-  def loadReferencesFromDataset(path: String): Array[Byte]
-
-  def fromFASTAFile(
-    name: String,
-    fastaFile: String,
-    indexFile: String,
-    xContigs: Array[String],
-    yContigs: Array[String],
-    mtContigs: Array[String],
-    parInput: Array[String],
-  ): Array[Byte] =
-    withExecuteContext { ctx =>
-      jsonToBytes {
-        ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
-          xContigs, yContigs, mtContigs, parInput).toJSON
-      }
-    }
-
-  def parseVCFMetadata(path: String): Array[Byte] =
-    withExecuteContext { ctx =>
-      jsonToBytes {
-        Extraction.decompose {
-          LoadVCF.parseHeaderMetadata(ctx.fs, Set.empty, TFloat64, path)
-        }(defaultJSONFormats)
-      }
-    }
-
-  def importFam(path: String, isQuantPheno: Boolean, delimiter: String, missingValue: String)
-    : Array[Byte] =
-    withExecuteContext { ctx =>
-      LoadPlink.importFamJSON(ctx.fs, path, isQuantPheno, delimiter, missingValue).getBytes(
-        StandardCharsets.UTF_8
-      )
-    }
-
   def execute(ctx: ExecuteContext, ir: IR): Either[Unit, (PTuple, Long)]
+
+  def backendContext(ctx: ExecuteContext): BackendContext
 }