feedback from @patrick-schultz

hail-is · Oct 16, 2024 · 2f05ecd · 2f05ecd
1 parent 1087175
commit 2f05ecd
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/hail/src/main/scala/is/hail/backend/ExecuteContext.scala b/hail/src/main/scala/is/hail/backend/ExecuteContext.scala
@@ -36,7 +36,7 @@ class OwningTempFileManager(fs: FS) extends TempFileManager {
   }
 }
 
-case class NonOwningTempFileManager(owner: TempFileManager) extends TempFileManager {
+class NonOwningTempFileManager private(owner: TempFileManager) extends TempFileManager {
   override def newTmpPath(tmpdir: String, prefix: String, extension: String): String =
     owner.newTmpPath(tmpdir, prefix, extension)
 

diff --git a/hail/src/main/scala/is/hail/expr/ir/CompileAndEvaluate.scala b/hail/src/main/scala/is/hail/expr/ir/CompileAndEvaluate.scala
@@ -57,7 +57,12 @@ object CompileAndEvaluate {
             optimize = optimize,
           )
 
-          Left(ctx.scopedExecution((hcl, fs, htc, r) => f(hcl, fs, htc, r)(r)))
+          val unit: Unit = ctx.scopedExecution { (hcl, fs, htc, r) =>
+            val execute = f(hcl, fs, htc, r)
+            ctx.time(execute(r))
+          }
+
+          Left(unit)
 
         case _ =>
           val (Some(PTypeReferenceSingleCodeType(resType: PTuple)), f) =
@@ -71,7 +76,12 @@ object CompileAndEvaluate {
               optimize = optimize,
             )
 
-          Right((resType, ctx.scopedExecution((hcl, fs, htc, r) => f(hcl, fs, htc, r)(r))))
+          val res = ctx.scopedExecution { (hcl, fs, htc, r) =>
+            val execute = f(hcl, fs, htc, r)
+            ctx.time(execute(r))
+          }
+
+          Right((resType, res))
       }
     }
 }