[query] Use valid globals reference in MWZJ and TABK (hail-is#14246)

CHANGELOG: Fix a bug, introduced in 0.2.114, in which
`Table.multi_way_zip_join` and `Table.aggregate_by_key` could throw
"NoSuchElementException: Ref with name `__iruid_...`" when one or more
of the tables had a number of partitions substantially different from
the desired number of output partitions.

Fixes hail-is#14245.

In both MultiWayZipJoin and TableAggregateByKey, we repartition the
child but neglect to use the new globals `Ref` from the repartitioned
child. As long as `repartitionNoShuffle` does not create a new
TableStage with new globals, this is fine, but that is not, in general,
true. It seems that recently, in lowered backends, when the repartition
cost is deemed "high" we generate a fresh TableStage with a fresh
globals ref.

hail/python/test/hail/table/test_table.py

@@ -668,6 +668,23 @@ def test_multi_way_zip_join_key_downcast2(self):
  ht = hl.Table.multi_way_zip_join(vcfs, 'data', 'new_globals')
  assert exp_count == ht._force_count()
 
+ def test_multi_way_zip_join_highly_unbalanced_partitions__issue_14245(self):
+ def import_vcf(file: str, partitions: int):
+ return (
+ hl.import_vcf(file, force_bgz=True, reference_genome='GRCh38', min_partitions=partitions)
+ .rows()
+ .select()
+ )
+
+ hl.Table.multi_way_zip_join(
+ [
+ import_vcf(resource('gvcfs/HG00096.g.vcf.gz'), 100),
+ import_vcf(resource('gvcfs/HG00268.g.vcf.gz'), 1),
+ ],
+ 'data',
+ 'new_globals',
+ ).write(new_temp_file(extension='ht'))
+
  def test_index_maintains_count(self):
  t1 = hl.Table.parallelize(
  [{'a': 'foo', 'b': 1}, {'a': 'bar', 'b': 2}, {'a': 'bar', 'b': 2}],

hail/src/main/scala/is/hail/expr/ir/lowering/LowerTableIR.scala

@@ -1197,39 +1197,39 @@ object LowerTableIR {
 
  case TableAggregateByKey(child, expr) =>
  val loweredChild = lower(child)
-
- loweredChild.repartitionNoShuffle(
+ val repartitioned = loweredChild.repartitionNoShuffle(
  ctx,
  loweredChild.partitioner.coarsen(child.typ.key.length).strictify(),
  )
- .mapPartition(Some(child.typ.key)) { partition =>
- Let(
- FastSeq("global" -> loweredChild.globals),
- mapIR(StreamGroupByKey(partition, child.typ.key, missingEqual = true)) { groupRef =>
- StreamAgg(
- groupRef,
- "row",
- bindIRs(
- ArrayRef(
- ApplyAggOp(
- FastSeq(I32(1)),
- FastSeq(SelectFields(Ref("row", child.typ.rowType), child.typ.key)),
- AggSignature(Take(), FastSeq(TInt32), FastSeq(child.typ.keyType)),
- ),
- I32(0),
- ), // FIXME: would prefer a First() agg op
- expr,
- ) { case Seq(key, value) =>
- MakeStruct(child.typ.key.map(k =>
- (k, GetField(key, k))
- ) ++ expr.typ.asInstanceOf[TStruct].fieldNames.map { f =>
- (f, GetField(value, f))
- })
- },
- )
- },
- )
- }
+
+ repartitioned.mapPartition(Some(child.typ.key)) { partition =>
+ Let(
+ FastSeq("global" -> repartitioned.globals),
+ mapIR(StreamGroupByKey(partition, child.typ.key, missingEqual = true)) { groupRef =>
+ StreamAgg(
+ groupRef,
+ "row",
+ bindIRs(
+ ArrayRef(
+ ApplyAggOp(
+ FastSeq(I32(1)),
+ FastSeq(SelectFields(Ref("row", child.typ.rowType), child.typ.key)),
+ AggSignature(Take(), FastSeq(TInt32), FastSeq(child.typ.keyType)),
+ ),
+ I32(0),
+ ), // FIXME: would prefer a First() agg op
+ expr,
+ ) { case Seq(key, value) =>
+ MakeStruct(child.typ.key.map(k =>
+ (k, GetField(key, k))
+ ) ++ expr.typ.asInstanceOf[TStruct].fieldNames.map { f =>
+ (f, GetField(value, f))
+ })
+ },
+ )
+ },
+ )
+ }
 
  case TableDistinct(child) =>
  val loweredChild = lower(child)
@@ -2155,7 +2155,10 @@ object LowerTableIR {
  )
  val repartitioned = lowered.map(_.repartitionNoShuffle(ctx, newPartitioner))
  val newGlobals = MakeStruct(FastSeq(
- globalName -> MakeArray(lowered.map(_.globals), TArray(lowered.head.globalType))
+ globalName -> MakeArray(
+ repartitioned.map(_.globals),
+ TArray(repartitioned.head.globalType),
+ )
  ))
  val globalsRef = Ref(genUID(), newGlobals.typ)