[query] eliminate optimization that can blow RAM

[danking]

CHANGELOG: On some pipelines, since at least 0.2.58 (commit 23813af), Hail could use essentially unbounded amounts of memory. This change removes "optimization" rules that accidentally caused that.

Closes #13606

[patrick-schultz]

These optimizations are too important to just disable. I think a better quickfix is to do

    case ToStream(ToArray(s), false) if s.typ.isInstanceOf[TStream] => s

    case ToStream(Let(name, value, ToArray(x)), false) if x.typ.isInstanceOf[TStream] =>
      Let(name, value, x)

The right fix is probably to add a requiresMemoryManagementPerElement flag to StreamMap---which could be generally useful when the producer didn't care about memory management, but the map body allocates a lot and wants to free after each row---and then use that to make these smarter:

    case ToStream(ToArray(s), false) if s.typ.isInstanceOf[TStream] => s
    case ToStream(ToArray(s), true) if s.typ.isInstanceOf[TStream] =>
      StreamMap(s, uid, Ref(uid), requiresMemoryManagementPerElement = true)

    case ToStream(Let(name, value, ToArray(x)), false) if x.typ.isInstanceOf[TStream] =>
      Let(name, value, x)
    case ToStream(Let(name, value, ToArray(x)), false) if x.typ.isInstanceOf[TStream] =>
      Let(name, value, StreamMap(x, uid, Ref(uid), requiresMemoryManagementPerElement = true))

And I still don't fully understand why these were deoptimizing the gnomad pipeline.

[danking]

FWIW, I finally found a simpler reproducer. It really takes some doing to convince the simplifier to apply this rule.

This operation should use a constant ~1GiB of RAM (in reality, in a non-broken pipeline it uses closer to 8GiB, but, still, a constant amount of RAM), but in reality memory use grows with each row processed

import hail as hl
ht = hl.utils.range_table(1)
ht = ht.key_by()
ht = ht.select(rows = hl.range(10))
ht = ht.explode('rows')
ht = ht.annotate(garbage=hl.range(1024 ** 3))
ht.write('/tmp/foo.ht', overwrite=True)

The simplifier cannot simplify the pipeline if the key is still present so this pipeline is sufficient to restore normal memory usage:

import hail as hl
ht = hl.utils.range_table(1)
ht = ht.select(rows = hl.range(10))
ht = ht.explode('rows')
ht = ht.annotate(garbage=hl.range(1024 ** 3))
ht.write('/tmp/foo.ht', overwrite=True)

The "bad" WritePartition body IR looks like this:

(StreamFlatMap __iruid_447
  (StreamRange -1 True
    (GetField start (Ref __iruid_446))
    (GetField end (Ref __iruid_446))
    (I32 1))
  (StreamMap __iruid_448
    (StreamRange 1 False (I32 0) (I32 10) (I32 1))
    (InsertFields
      (Literal Struct{} <literal value>)
      ("rows" "garbage")
      (rows (Ref __iruid_448))
      (garbage
        (ToArray
          (StreamRange 2 False
            (I32 0)
            (I32 1073741824)
            (I32 1)))))))

The "good" IR looks like this:

(StreamFlatMap __iruid_480
  (StreamRange -1 True
    (GetField start (Ref __iruid_479))
    (GetField end (Ref __iruid_479))
    (I32 1))
  (Let __iruid_481
    (MakeStruct
      (idx (Ref __iruid_480))
      (rows
        (ToArray
          (StreamRange 1 False (I32 0) (I32 10) (I32 1)))))
    (StreamMap __iruid_482
      (ToStream True (GetField rows (Ref __iruid_481)))
      (InsertFields
        (Ref __iruid_481)
        ("idx" "rows" "garbage")
        (rows (Ref __iruid_482))
        (garbage
          (ToArray
            (StreamRange 2 False
              (I32 0)
              (I32 1073741824)
              (I32 1))))))))

Notice, in particular, that the StreamMap inside the StreamFlatMap uses memory management because the originating ToStream uses memory management.

[danking]

Issue for StreamMap memory management: #13623

done.

[danking]

@patrick-schultz bump for tomorrow

[patrick-schultz]

That's a very clear example, thanks!