[query] Avoid py4j for python-backend interactions

hail/python/hail/backend/local_backend.py

@@ -67,8 +67,7 @@ def __init__(self, tmpdir, log, quiet, append, branching_factor,
         jhc = hail_package.HailContext.apply(
             jbackend,
             branching_factor,
-            optimizer_iterations,
-            True
+            optimizer_iterations
         )
 
         super(LocalBackend, self).__init__(self._gateway.jvm, jbackend, jhc)
@@ -102,11 +101,8 @@ def register_ir_function(self,
             jbody)
 
     def stop(self):
-        self._backend_server.stop()
-        self._jhc.stop()
-        self._jhc = None
+        super().stop()
         self._gateway.shutdown()
-        self._registered_ir_function_names = set()
         uninstall_exception_handler()
 
     @property

hail/python/hail/backend/py4j_backend.py

@@ -287,3 +287,10 @@ def _to_java_value_ir(self, ir):
 
     def _to_java_blockmatrix_ir(self, ir):
         return self._to_java_ir(ir, self._parse_blockmatrix_ir)
+
+    def stop(self):
+        self._backend_server.stop()
+        self._jhc.stop()
+        self._jhc = None
+        self._registered_ir_function_names = set()
+        uninstall_exception_handler()

hail/python/hail/backend/spark_backend.py

@@ -13,7 +13,7 @@
 from hailtop.aiotools.router_fs import RouterAsyncFS
 from hailtop.aiotools.validators import validate_file
 
-from .py4j_backend import Py4JBackend, uninstall_exception_handler
+from .py4j_backend import Py4JBackend
 from .backend import local_jar_information
 
 
@@ -114,7 +114,7 @@ def __init__(self, idempotent, sc, spark_conf, app_name, master,
                 jsc, app_name, master, local, log, True, append, skip_logging_configuration, min_block_size, tmpdir, local_tmpdir,
                 gcs_requester_pays_project, gcs_requester_pays_buckets)
             jhc = hail_package.HailContext.apply(
-                jbackend, branching_factor, optimizer_iterations, True)
+                jbackend, branching_factor, optimizer_iterations)
 
         self._jsc = jbackend.sc()
         if sc:
@@ -146,13 +146,10 @@ def validate_file(self, uri: str) -> None:
         validate_file(uri, self._router_async_fs)
 
     def stop(self):
+        super().stop()
         self._jbackend.close()
-        self._jhc.stop()
-        self._jhc = None
         self.sc.stop()
         self.sc = None
-        self._registered_ir_function_names = set()
-        uninstall_exception_handler()
 
     @property
     def fs(self):

hail/python/hail/ir/ir.py

@@ -3718,10 +3718,10 @@ def __del__(self):
 
 
 class JavaIR(IR):
-    def __init__(self, hail_type, ir_id, ref: Optional[JavaIRSharedReference] = None):
+    def __init__(self, hail_type: HailType, ir_id: int, ref: Optional[JavaIRSharedReference] = None):
         super(JavaIR, self).__init__()
-        self._type: HailType = hail_type
-        self._id: int = ir_id
+        self._type = hail_type
+        self._id = ir_id
         self._ref = ref or JavaIRSharedReference(ir_id)
 
     def copy(self):

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

@@ -107,8 +107,7 @@ object HailContext {
 
   def apply(backend: Backend,
     branchingFactor: Int = 50,
-    optimizerIterations: Int = 3,
-    addDefaultReferences: Boolean = true): HailContext = synchronized {
+    optimizerIterations: Int = 3): HailContext = synchronized {
     require(theContext == null)
     checkJavaVersion()
 
@@ -124,11 +123,6 @@ object HailContext {
 
     info(s"Running Hail version ${ theContext.version }")
 
-    // needs to be after `theContext` is set, since this creates broadcasts
-    if (addDefaultReferences) {
-      backend.addDefaultReferences()
-    }
-
     theContext
   }
 

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

@@ -56,6 +56,7 @@ object LocalBackend {
       gcsRequesterPaysProject,
       gcsRequesterPaysBuckets
     )
+    theLocalBackend.addDefaultReferences()
     theLocalBackend
   }
 

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

@@ -336,6 +336,7 @@ class ServiceBackend(
         elementType.virtualType,
         BufferSpec.parseOrDefault(bufferSpecString)
       )
+      assert(pt.isFieldDefined(off, 0))
       codec.encode(ctx, elementType, pt.loadField(off, 0))
     }
   }
@@ -449,7 +450,7 @@ object ServiceBackendAPI {
       HailContext.get.backend = backend
       log.info("Default references added to already initialized HailContexet.")
     } else {
-      HailContext(backend, 50, 3, false)
+      HailContext(backend, 50, 3)
       log.info("HailContexet initialized.")
     }
 

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

@@ -247,6 +247,7 @@ object SparkBackend {
     sc1.uiWebUrl.foreach(ui => info(s"SparkUI: $ui"))
 
     theSparkBackend = new SparkBackend(tmpdir, localTmpdir, sc1, gcsRequesterPaysProject, gcsRequesterPaysBuckets)
+    theSparkBackend.addDefaultReferences()
     theSparkBackend
   }