[query] Move LoweredTableReaderCoercer into ExecuteContext

ehigham · ehigham · commit 9615ece71b85 · 2024-10-16T17:30:10.000-04:00
diff --git a/hail/src/main/scala/is/hail/backend/Backend.scala b/hail/src/main/scala/is/hail/backend/Backend.scala
@@ -88,8 +88,6 @@ abstract class Backend extends Closeable {
   def asSpark(implicit E: Enclosing): SparkBackend =
     fatal(s"${getClass.getSimpleName}: ${E.value} requires SparkBackend")
 
-  def shouldCacheQueryInfo: Boolean = true
-
   def lowerDistributedSort(
     ctx: ExecuteContext,
     stage: TableStage,
diff --git a/hail/src/main/scala/is/hail/backend/ExecuteContext.scala b/hail/src/main/scala/is/hail/backend/ExecuteContext.scala
@@ -5,6 +5,7 @@ import is.hail.annotations.{Region, RegionPool}
 import is.hail.asm4s.HailClassLoader
 import is.hail.backend.local.LocalTaskContext
 import is.hail.expr.ir.{BaseIR, CodeCacheKey, CompiledFunction}
+import is.hail.expr.ir.LoweredTableReader.LoweredTableReaderCoercer
 import is.hail.expr.ir.lowering.IrMetadata
 import is.hail.io.fs.FS
 import is.hail.linalg.BlockMatrix
@@ -74,6 +75,7 @@ object ExecuteContext {
     blockMatrixCache: mutable.Map[String, BlockMatrix],
     codeCache: mutable.Map[CodeCacheKey, CompiledFunction[_]],
     irCache: mutable.Map[Int, BaseIR],
+    coercerCache: mutable.Map[Any, LoweredTableReaderCoercer],
   )(
     f: ExecuteContext => T
   ): T = {
@@ -95,6 +97,7 @@ object ExecuteContext {
           blockMatrixCache,
           codeCache,
           irCache,
+          coercerCache,
         ))(f(_))
       }
     }
@@ -127,6 +130,7 @@ class ExecuteContext(
   val BlockMatrixCache: mutable.Map[String, BlockMatrix],
   val CodeCache: mutable.Map[CodeCacheKey, CompiledFunction[_]],
   val IrCache: mutable.Map[Int, BaseIR],
+  val CoercerCache: mutable.Map[Any, LoweredTableReaderCoercer],
 ) extends Closeable {
 
   val rngNonce: Long =
@@ -199,6 +203,7 @@ class ExecuteContext(
     blockMatrixCache: mutable.Map[String, BlockMatrix] = this.BlockMatrixCache,
     codeCache: mutable.Map[CodeCacheKey, CompiledFunction[_]] = this.CodeCache,
     irCache: mutable.Map[Int, BaseIR] = this.IrCache,
+    coercerCache: mutable.Map[Any, LoweredTableReaderCoercer] = this.CoercerCache,
   )(
     f: ExecuteContext => A
   ): A =
@@ -218,5 +223,6 @@ class ExecuteContext(
       blockMatrixCache,
       codeCache,
       irCache,
+      coercerCache,
     ))(f)
 }
diff --git a/hail/src/main/scala/is/hail/backend/caching/package.scala b/hail/src/main/scala/is/hail/backend/caching/package.scala
@@ -0,0 +1,16 @@
+package is.hail.backend
+
+import scala.collection.mutable
+
+package object caching {
+  private[this] object NoCachingInstance extends mutable.AbstractMap[Any, Any] {
+    override def +=(kv: (Any, Any)): NoCachingInstance.this.type = this
+    override def -=(key: Any): NoCachingInstance.this.type = this
+    override def get(key: Any): Option[Any] = None
+    override def iterator: Iterator[(Any, Any)] = Iterator.empty
+    override def getOrElseUpdate(key: Any, op: => Any): Any = op
+  }
+
+  def NoCaching[K, V]: mutable.Map[K, V] =
+    NoCachingInstance.asInstanceOf[mutable.Map[K, V]]
+}
diff --git a/hail/src/main/scala/is/hail/backend/local/LocalBackend.scala b/hail/src/main/scala/is/hail/backend/local/LocalBackend.scala
@@ -7,6 +7,7 @@ import is.hail.backend._
 import is.hail.backend.py4j.Py4JBackendExtensions
 import is.hail.expr.Validate
 import is.hail.expr.ir._
+import is.hail.expr.ir.LoweredTableReader.LoweredTableReaderCoercer
 import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.compile.Compile
 import is.hail.expr.ir.lowering._
@@ -81,6 +82,7 @@ class LocalBackend(
   private[this] val theHailClassLoader = new HailClassLoader(getClass.getClassLoader)
   private[this] val codeCache = new Cache[CodeCacheKey, CompiledFunction[_]](50)
   private[this] val persistedIR: mutable.Map[Int, BaseIR] = mutable.Map()
+  private[this] val coercerCache = new Cache[Any, LoweredTableReaderCoercer](32)
 
   // flags can be set after construction from python
   def fs: FS = RouterFS.buildRoutes(CloudStorageFSConfig.fromFlagsAndEnv(None, flags))
@@ -106,6 +108,7 @@ class LocalBackend(
         ImmutableMap.empty,
         codeCache,
         persistedIR,
+        coercerCache,
       )(f)
     }
 
diff --git a/hail/src/main/scala/is/hail/backend/py4j/Py4JBackendExtensions.scala b/hail/src/main/scala/is/hail/backend/py4j/Py4JBackendExtensions.scala
@@ -3,7 +3,10 @@ package is.hail.backend.py4j
 import is.hail.HailFeatureFlags
 import is.hail.backend.{Backend, ExecuteContext, NonOwningTempFileManager, TempFileManager}
 import is.hail.expr.{JSONAnnotationImpex, SparkAnnotationImpex}
-import is.hail.expr.ir.{BaseIR, BlockMatrixIR, EncodedLiteral, GetFieldByIdx, IRParser, Interpret, MatrixIR, MatrixNativeReader, MatrixRead, NativeReaderOptions, TableLiteral, TableValue}
+import is.hail.expr.ir.{
+  BaseIR, BlockMatrixIR, EncodedLiteral, GetFieldByIdx, IRParser, Interpret, MatrixIR,
+  MatrixNativeReader, MatrixRead, NativeReaderOptions, TableLiteral, TableValue,
+}
 import is.hail.expr.ir.IRParser.parseType
 import is.hail.expr.ir.functions.IRFunctionRegistry
 import is.hail.io.reference.{IndexedFastaSequenceFile, LiftOver}
diff --git a/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala b/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala
@@ -4,6 +4,7 @@ import is.hail.{CancellingExecutorService, HailContext, HailFeatureFlags}
 import is.hail.annotations._
 import is.hail.asm4s._
 import is.hail.backend._
+import is.hail.backend.caching.NoCaching
 import is.hail.backend.service.ServiceBackend.MaxAvailableGcsConnections
 import is.hail.expr.Validate
 import is.hail.expr.ir.{
@@ -63,8 +64,6 @@ class ServiceBackend(
   private[this] var stageCount = 0
   private[this] val executor = Executors.newFixedThreadPool(MaxAvailableGcsConnections)
 
-  override def shouldCacheQueryInfo: Boolean = false
-
   def defaultParallelism: Int = 4
 
   def broadcast[T: ClassTag](_value: T): BroadcastValue[T] = {
@@ -316,9 +315,10 @@ class ServiceBackend(
         ),
         new IrMetadata(),
         references,
-        ImmutableMap.empty,
-        mutable.Map.empty,
-        ImmutableMap.empty,
+        NoCaching,
+        NoCaching,
+        NoCaching,
+        NoCaching,
       )(f)
     }
 }
diff --git a/hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala b/hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala
@@ -4,10 +4,11 @@ import is.hail.{HailContext, HailFeatureFlags}
 import is.hail.annotations._
 import is.hail.asm4s._
 import is.hail.backend._
-import is.hail.backend.caching.BlockMatrixCache
+import is.hail.backend.caching.{BlockMatrixCache, NoCaching}
 import is.hail.backend.py4j.Py4JBackendExtensions
 import is.hail.expr.Validate
 import is.hail.expr.ir._
+import is.hail.expr.ir.LoweredTableReader.LoweredTableReaderCoercer
 import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.compile.Compile
 import is.hail.expr.ir.lowering._
@@ -343,6 +344,7 @@ class SparkBackend(
   private[this] val bmCache = new BlockMatrixCache()
   private[this] val codeCache = new Cache[CodeCacheKey, CompiledFunction[_]](50)
   private[this] val persistedIr = mutable.Map.empty[Int, BaseIR]
+  private[this] val coercerCache = new Cache[Any, LoweredTableReaderCoercer](32)
 
   def createExecuteContextForTests(
     timer: ExecutionTimer,
@@ -365,8 +367,9 @@ class SparkBackend(
       new IrMetadata(),
       references,
       ImmutableMap.empty,
-      mutable.Map.empty,
+      NoCaching,
       ImmutableMap.empty,
+      NoCaching,
     )
 
   override def withExecuteContext[T](f: ExecuteContext => T)(implicit E: Enclosing): (T, Timings) =
@@ -393,6 +396,7 @@ class SparkBackend(
         bmCache,
         codeCache,
         persistedIr,
+        coercerCache,
       )(f)
     }
 
diff --git a/hail/src/main/scala/is/hail/expr/ir/GenericTableValue.scala b/hail/src/main/scala/is/hail/expr/ir/GenericTableValue.scala
@@ -3,6 +3,7 @@ package is.hail.expr.ir
 import is.hail.annotations.Region
 import is.hail.asm4s._
 import is.hail.backend.ExecuteContext
+import is.hail.expr.ir.LoweredTableReader.LoweredTableReaderCoercer
 import is.hail.expr.ir.functions.UtilFunctions
 import is.hail.expr.ir.lowering.{TableStage, TableStageDependency}
 import is.hail.expr.ir.streams.StreamProducer
@@ -143,16 +144,6 @@ class PartitionIteratorLongReader(
     )
 }
 
-abstract class LoweredTableReaderCoercer {
-  def coerce(
-    ctx: ExecuteContext,
-    globals: IR,
-    contextType: Type,
-    contexts: IndexedSeq[Any],
-    body: IR => IR,
-  ): TableStage
-}
-
 class GenericTableValue(
   val fullTableType: TableType,
   val uidFieldName: String,
@@ -168,12 +159,11 @@ class GenericTableValue(
   assert(contextType.hasField("partitionIndex"))
   assert(contextType.fieldType("partitionIndex") == TInt32)
 
-  private var ltrCoercer: LoweredTableReaderCoercer = _
-
   private def getLTVCoercer(ctx: ExecuteContext, context: String, cacheKey: Any)
-    : LoweredTableReaderCoercer = {
-    if (ltrCoercer == null) {
-      ltrCoercer = LoweredTableReader.makeCoercer(
+    : LoweredTableReaderCoercer =
+    ctx.CoercerCache.getOrElseUpdate(
+      (1, contextType, fullTableType.key, cacheKey),
+      LoweredTableReader.makeCoercer(
         ctx,
         fullTableType.key,
         1,
@@ -184,11 +174,8 @@ class GenericTableValue(
         bodyPType,
         body,
         context,
-        cacheKey,
-      )
-    }
-    ltrCoercer
-  }
+      ),
+    )
 
   def toTableStage(ctx: ExecuteContext, requestedType: TableType, context: String, cacheKey: Any)
     : TableStage = {
@@ -217,11 +204,13 @@ class GenericTableValue(
       val contextsIR = ToStream(Literal(TArray(contextType), contexts))
       TableStage(globalsIR, p, TableStageDependency.none, contextsIR, requestedBody)
     } else {
-      getLTVCoercer(ctx, context, cacheKey).coerce(
+      getLTVCoercer(ctx, context, cacheKey)(
         ctx,
         globalsIR,
-        contextType, contexts,
-        requestedBody)
+        contextType,
+        contexts,
+        requestedBody,
+      )
     }
   }
 }
diff --git a/hail/src/main/scala/is/hail/expr/ir/TableIR.scala b/hail/src/main/scala/is/hail/expr/ir/TableIR.scala
