Skip to content

Commit b19073d

Browse files
ehighamehigham
committed
[query] Expose references via ExecuteContext 
1 parent 2d4edc1 commit b19073d

File tree

19 files changed

+315
-350
lines changed

19 files changed

+315
-350
lines changed

‎hail/src/main/scala/is/hail/HailFeatureFlags.scala‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,9 @@ class HailFeatureFlags private (
6868
flags.update(flag, value)
6969
}
7070

71+
def +(feature: (String, String)): HailFeatureFlags =
72+
new HailFeatureFlags(flags + (feature._1 -> feature._2))
73+
7174
def get(flag: String): String = flags(flag)
7275

7376
def lookup(flag: String): Option[String] =

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

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,6 @@ abstract class Backend extends Closeable {
105105
def lookupOrCompileCachedFunction[T](k: CodeCacheKey)(f: => CompiledFunction[T])
106106
: CompiledFunction[T]
107107

108-
def references: mutable.Map[String, ReferenceGenome]
109-
110108
def lowerDistributedSort(
111109
ctx: ExecuteContext,
112110
stage: TableStage,
@@ -181,10 +179,8 @@ abstract class Backend extends Closeable {
181179
): Array[Byte] =
182180
withExecuteContext { ctx =>
183181
jsonToBytes {
184-
Extraction.decompose {
185-
ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
186-
xContigs, yContigs, mtContigs, parInput).toJSON
187-
}(defaultJSONFormats)
182+
ReferenceGenome.fromFASTAFile(ctx, name, fastaFile, indexFile,
183+
xContigs, yContigs, mtContigs, parInput).toJSON
188184
}
189185
}
190186

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ object ExecuteContext {
6363
tmpdir: String,
6464
localTmpdir: String,
6565
backend: Backend,
66+
references: Map[String, ReferenceGenome],
6667
fs: FS,
6768
timer: ExecutionTimer,
6869
tempFileManager: TempFileManager,
@@ -79,6 +80,7 @@ object ExecuteContext {
7980
tmpdir,
8081
localTmpdir,
8182
backend,
83+
references,
8284
fs,
8385
region,
8486
timer,
@@ -107,6 +109,7 @@ class ExecuteContext(
107109
val tmpdir: String,
108110
val localTmpdir: String,
109111
val backend: Backend,
112+
val references: Map[String, ReferenceGenome],
110113
val fs: FS,
111114
val r: Region,
112115
val timer: ExecutionTimer,
@@ -128,7 +131,7 @@ class ExecuteContext(
128131
)
129132
}
130133

131-
def stateManager = HailStateManager(backend.references.toMap)
134+
val stateManager = HailStateManager(references)
132135

133136
val tempFileManager: TempFileManager =
134137
if (_tempFileManager != null) _tempFileManager else new OwningTempFileManager(fs)
@@ -154,7 +157,7 @@ class ExecuteContext(
154157

155158
def getFlag(name: String): String = flags.get(name)
156159

157-
def getReference(name: String): ReferenceGenome = backend.references(name)
160+
def getReference(name: String): ReferenceGenome = references(name)
158161

159162
def shouldWriteIRFiles(): Boolean = getFlag("write_ir_files") != null
160163

@@ -174,6 +177,7 @@ class ExecuteContext(
174177
tmpdir: String = this.tmpdir,
175178
localTmpdir: String = this.localTmpdir,
176179
backend: Backend = this.backend,
180+
references: Map[String, ReferenceGenome] = this.references,
177181
fs: FS = this.fs,
178182
r: Region = this.r,
179183
timer: ExecutionTimer = this.timer,
@@ -189,6 +193,7 @@ class ExecuteContext(
189193
tmpdir,
190194
localTmpdir,
191195
backend,
196+
references,
192197
fs,
193198
r,
194199
timer,

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ class LocalBackend(
8686
override val references: mutable.Map[String, ReferenceGenome],
8787
) extends Backend with BackendWithCodeCache with Py4JBackendExtensions {
8888

89+
override def backend: Backend = this
8990
override val flags: HailFeatureFlags = HailFeatureFlags.fromEnv()
9091
override def longLifeTempFileManager: TempFileManager = null
9192

@@ -101,6 +102,7 @@ class LocalBackend(
101102
tmpdir,
102103
tmpdir,
103104
this,
105+
references.toMap,
104106
fs,
105107
timer,
106108
null,

‎hail/src/main/scala/is/hail/backend/py4j/Py4JBackendExtensions.scala‎

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import is.hail.expr.ir.{
1010
}
1111
import is.hail.expr.ir.IRParser.parseType
1212
import is.hail.expr.ir.functions.IRFunctionRegistry
13+
import is.hail.io.reference.{IndexedFastaSequenceFile, LiftOver}
1314
import is.hail.linalg.RowMatrix
1415
import is.hail.types.physical.PStruct
1516
import is.hail.types.virtual.{TArray, TInterval}
@@ -30,7 +31,9 @@ import org.json4s.Formats
3031
import org.json4s.jackson.{JsonMethods, Serialization}
3132
import sourcecode.Enclosing
3233

33-
trait Py4JBackendExtensions { this: Backend =>
34+
trait Py4JBackendExtensions {
35+
def backend: Backend
36+
def references: mutable.Map[String, ReferenceGenome]
3437
def persistedIR: mutable.Map[Int, BaseIR]
3538
def flags: HailFeatureFlags
3639
def longLifeTempFileManager: TempFileManager
@@ -61,7 +64,9 @@ trait Py4JBackendExtensions { this: Backend =>
6164
persistedIR.remove(id)
6265

6366
def pyAddSequence(name: String, fastaFile: String, indexFile: String): Unit =
64-
withExecuteContext(ctx => references(name).addSequence(ctx, fastaFile, indexFile))
67+
backend.withExecuteContext { ctx =>
68+
references(name).addSequence(IndexedFastaSequenceFile(ctx.fs, fastaFile, indexFile))
69+
}
6570

6671
def pyRemoveSequence(name: String): Unit =
6772
references(name).removeSequence()
@@ -76,7 +81,7 @@ trait Py4JBackendExtensions { this: Backend =>
7681
partitionSize: java.lang.Integer,
7782
entries: String,
7883
): Unit =
79-
withExecuteContext { ctx =>
84+
backend.withExecuteContext { ctx =>
8085
val rm = RowMatrix.readBlockMatrix(ctx.fs, pathIn, partitionSize)
8186
entries match {
8287
case "full" =>
@@ -114,7 +119,7 @@ trait Py4JBackendExtensions { this: Backend =>
114119
returnType: String,
115120
bodyStr: String,
116121
): Unit = {
117-
withExecuteContext { ctx =>
122+
backend.withExecuteContext { ctx =>
118123
IRFunctionRegistry.registerIR(
119124
ctx,
120125
name,
@@ -128,10 +133,10 @@ trait Py4JBackendExtensions { this: Backend =>
128133
}
129134

130135
def pyExecuteLiteral(irStr: String): Int =
131-
withExecuteContext { ctx =>
136+
backend.withExecuteContext { ctx =>
132137
val ir = IRParser.parse_value_ir(irStr, IRParserEnvironment(ctx, persistedIR.toMap))
133138
assert(ir.typ.isRealizable)
134-
execute(ctx, ir) match {
139+
backend.execute(ctx, ir) match {
135140
case Left(_) => throw new HailException("Can't create literal")
136141
case Right((pt, addr)) =>
137142
val field = GetFieldByIdx(EncodedLiteral.fromPTypeAndAddress(pt, addr, ctx), 0)
@@ -160,13 +165,13 @@ trait Py4JBackendExtensions { this: Backend =>
160165
}
161166

162167
def pyToDF(s: String): DataFrame =
163-
withExecuteContext { ctx =>
168+
backend.withExecuteContext { ctx =>
164169
val tir = IRParser.parse_table_ir(s, IRParserEnvironment(ctx, irMap = persistedIR.toMap))
165170
Interpret(tir, ctx).toDF()
166171
}
167172

168173
def pyReadMultipleMatrixTables(jsonQuery: String): util.List[MatrixIR] =
169-
withExecuteContext { ctx =>
174+
backend.withExecuteContext { ctx =>
170175
log.info("pyReadMultipleMatrixTables: got query")
171176
val kvs = JsonMethods.parse(jsonQuery) match {
172177
case json4s.JObject(values) => values.toMap
@@ -195,19 +200,24 @@ trait Py4JBackendExtensions { this: Backend =>
195200
addReference(ReferenceGenome.fromJSON(jsonConfig))
196201

197202
def pyRemoveReference(name: String): Unit =
198-
references.remove(name)
203+
removeReference(name)
199204

200205
def pyAddLiftover(name: String, chainFile: String, destRGName: String): Unit =
201-
withExecuteContext(ctx => references(name).addLiftover(ctx, chainFile, destRGName))
206+
backend.withExecuteContext { ctx =>
207+
references(name).addLiftover(references(destRGName), LiftOver(ctx.fs, chainFile))
208+
}
202209

203210
def pyRemoveLiftover(name: String, destRGName: String): Unit =
204211
references(name).removeLiftover(destRGName)
205212

206213
private[this] def addReference(rg: ReferenceGenome): Unit =
207214
ReferenceGenome.addFatalOnCollision(references, FastSeq(rg))
208215

216+
private[this] def removeReference(name: String): Unit =
217+
references -= name
218+
209219
def parse_value_ir(s: String, refMap: java.util.Map[String, String]): IR =
210-
withExecuteContext { ctx =>
220+
backend.withExecuteContext { ctx =>
211221
IRParser.parse_value_ir(
212222
s,
213223
IRParserEnvironment(ctx, irMap = persistedIR.toMap),
@@ -233,7 +243,7 @@ trait Py4JBackendExtensions { this: Backend =>
233243
}
234244

235245
def loadReferencesFromDataset(path: String): Array[Byte] =
236-
withExecuteContext { ctx =>
246+
backend.withExecuteContext { ctx =>
237247
val rgs = ReferenceGenome.fromHailDataset(ctx.fs, path)
238248
ReferenceGenome.addFatalOnCollision(references, rgs)
239249

@@ -247,7 +257,7 @@ trait Py4JBackendExtensions { this: Backend =>
247257
f: ExecuteContext => T
248258
)(implicit E: Enclosing
249259
): T =
250-
withExecuteContext { ctx =>
260+
backend.withExecuteContext { ctx =>
251261
val tempFileManager = longLifeTempFileManager
252262
if (selfContainedExecution && tempFileManager != null) f(ctx)
253263
else ctx.local(tempFileManager = NonOwningTempFileManager(tempFileManager))(f)

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

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import is.hail.expr.ir.analyses.SemanticHash
1313
import is.hail.expr.ir.functions.IRFunctionRegistry
1414
import is.hail.expr.ir.lowering._
1515
import is.hail.io.fs._
16+
import is.hail.io.reference.{IndexedFastaSequenceFile, LiftOver}
1617
import is.hail.linalg.BlockMatrix
1718
import is.hail.services.{BatchClient, JobGroupRequest, _}
1819
import is.hail.services.JobGroupStates.{Cancelled, Failure, Running, Success}
@@ -93,11 +94,20 @@ object ServiceBackend {
9394
rpcConfig.custom_references.map(ReferenceGenome.fromJSON),
9495
)
9596

96-
val backend = new ServiceBackend(
97+
rpcConfig.liftovers.foreach { case (sourceGenome, liftoversForSource) =>
98+
liftoversForSource.foreach { case (destGenome, chainFile) =>
99+
references(sourceGenome).addLiftover(references(destGenome), LiftOver(fs, chainFile))
100+
}
101+
}
102+
rpcConfig.sequences.foreach { case (rg, seq) =>
103+
references(rg).addSequence(IndexedFastaSequenceFile(fs, seq.fasta, seq.index))
104+
}
105+
106+
new ServiceBackend(
97107
JarUrl(jarLocation),
98108
name,
99109
theHailClassLoader,
100-
references,
110+
references.toMap,
101111
batchClient,
102112
batchConfig,
103113
flags,
@@ -106,27 +116,14 @@ object ServiceBackend {
106116
backendContext,
107117
scratchDir,
108118
)
109-
110-
backend.withExecuteContext { ctx =>
111-
rpcConfig.liftovers.foreach { case (sourceGenome, liftoversForSource) =>
112-
liftoversForSource.foreach { case (destGenome, chainFile) =>
113-
references(sourceGenome).addLiftover(ctx, chainFile, destGenome)
114-
}
115-
}
116-
rpcConfig.sequences.foreach { case (rg, seq) =>
117-
references(rg).addSequence(ctx, seq.fasta, seq.index)
118-
}
119-
}
120-
121-
backend
122119
}
123120
}
124121

125122
class ServiceBackend(
126123
val jarSpec: JarSpec,
127124
var name: String,
128125
val theHailClassLoader: HailClassLoader,
129-
override val references: mutable.Map[String, ReferenceGenome],
126+
val references: Map[String, ReferenceGenome],
130127
val batchClient: BatchClient,
131128
val batchConfig: BatchConfig,
132129
val flags: HailFeatureFlags,
@@ -397,6 +394,7 @@ class ServiceBackend(
397394
tmpdir,
398395
"file:///tmp",
399396
this,
397+
references,
400398
fs,
401399
timer,
402400
null,

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,7 @@ class SparkBackend(
345345
new HadoopFS(new SerializableHadoopConfiguration(conf))
346346
}
347347

348+
override def backend: Backend = this
348349
override val flags: HailFeatureFlags = HailFeatureFlags.fromEnv()
349350

350351
override val longLifeTempFileManager: TempFileManager =
@@ -374,6 +375,7 @@ class SparkBackend(
374375
tmpdir,
375376
localTmpdir,
376377
this,
378+
references.toMap,
377379
fs,
378380
region,
379381
timer,
@@ -393,6 +395,7 @@ class SparkBackend(
393395
tmpdir,
394396
localTmpdir,
395397
this,
398+
references.toMap,
396399
fs,
397400
timer,
398401
null,

‎hail/src/main/scala/is/hail/io/reference/LiftOver.scala‎

Lines changed: 0 additions & 69 deletions
This file was deleted.

0 commit comments

Comments
 (0)