Scatter tasks around predefined batches inside a WDL workflow

Question

Suppose I have two arrays Array[File] X and Array[String] Y both of length N, with Y being a sort of R-like factor that defines batch membership for X, for a total of M batches with M < N. A valid example would be:

Array[File] X = ["file1", "file2", "file3", "file4", "file5"]
Array[String] Y  =  ["A",     "A",     "B",     "A",     "B"]

In which case N=5 and M=2. I want to scatter M tasks so that each task takes only the files corresponding to the same batch. In the previous example, one task would take as input:

Array[File] X =["file1", "file2", "file4"]
String Y = "A"

And the other task would take as input:

Array[File] X =["file3", "file5"]
String Y = "B"

Is there an easy way to do this within a WDL workflow in the general case? I understand that there are less elegant ways to do this as its own task but, as this seems like it could be a general purpose scatter/gather problem, an easy solution might be of general interest to the community.

Answer 1

Here's an alternative to using a custom task, albeit still rather unappetizing:

workflow groupby {
  input {
    File? nil
  }

  Array[File] X = ["file1", "file2", "file3", "file4", "file5"]
  Array[String] Y  =  ["A",     "A",     "B",     "A",     "B"]

  scatter (y in ["A", "B"]) {
    scatter (p in zip(X,Y)) {
      File? x_or_nil = if (p.right == y) then p.left else nil
    }
    Pair[String,Array[File]] group = (y, select_all(x_or_nil))
  }

  output {
    Array[Pair[String,Array[File]]] ans = group
  }
}

(It'll be slightly cleaner with the upcoming None literal instead of dummy nil)

FWIW, here's a lengthy discussion about the idea of allowing developers to define custom library functions which just run a small task that manipulates data structures (hinting that it needn't necessarily offload to a worker node, with the huge latency that typically entails).

Answer 2

It turns out the development version has a function collect_by_key that is particularly helpful for scattering around predefined batches:

Array[String] X  =  ["A",     "A",     "B",     "A",     "B"]
Array[File] Y = ["file1", "file2", "file3", "file4", "file5"]
scatter (pair in as_pairs(collect_by_key(zip(X,Y)))) {
  call test {input: X = pair.left, Y = pair.right}
}

This will call task t first with input:

String X = "A"
Array[File] Y =["file1", "file2", "file4"]

And then with input:

String X = "B"
Array[File] Y =["file3", "file5"]

This is a readable solution to this problem allowed by collect_by_key. This function is not very well advertised (I had completely missed it) but kudos to who introduced it in the language. The only drawback is that collect_by_key is not available in version 1.0.