[query] force aggregate_cols to be local

[patrick-schultz]

CHANGELOG: MatrixTable.aggregate_cols no longer forces a distributed computation. This should be what you want in the majority of cases. In case you know the aggregation is very slow and should be parallelized, use mt.cols().aggregate instead.

Most of the time, aggregate_cols will be much faster performing the aggregation locally. Currently, we generate a TableAggregate over a TableParallelize of the columns. We shouldn't try to optimize that to a local computation during compilation; TableParallelize should express the intent that the computation is expensive and really should be parallelized. This should be considered part of the semantics the compiler must preserve.

This PR changes aggregate_cols to explicitly generate a local computation using StreamAgg (which was only exposed in Python relatively recently, which is why we haven't made this change sooner). Longer term, aggregating columns should probably get its own IR node, especially once we start partitioning along columns.

[danking]

Looks like you need the globals too. Two concerns:

1. Does localize entries not produce a table aggregate?
2. Does this preserve the column ordering that we all agreed on? Is that tested?

[danking]

see comment

[patrick-schultz]

1. localize_entries is a no-op at runtime, it just exposes the table representation to python
2. Good point. We agreed on ordering by column key, right?

[danking]

I can't find the conversation we all had about this, but I strongly disagree with (2). It has to be the case that this always evaluates to true. This would be profoundly confusing if not.

mt = mt.annotate_globals(
    phenos = hl.agg.collect(mt.pheno)
)
mt = mt.add_col_index('col_idx')
mt = mt.annotate_cols(
    same = mt.phenos[mt.col_idx] == mt.pheno
)

I hate it but I'm willing to accept that

phenos = mt.cols().pheno.collect()

Returns them out of order, even though I don't like that.

EDIT: hit enter too fast

[danking]

Heh. We can't annotate_globals.

[danking]

OK, alright. I think you're right. I'm just reliving how frustrated I am by this situation. Can you modify aggregate_cols to include the same warning from cols()?

In [6]: import hail as hl
   ...: mt = hl.utils.range_matrix_table(3, 3)
   ...: mt = mt.choose_cols([2, 1, 0])
   ...: mt = mt.checkpoint('/tmp/foo.mt', overwrite=True)
   ...: mt.col_idx.show()
   ...: print(mt.aggregate_cols(hl.agg.collect(mt.col_idx)))
   ...: mt = mt.key_cols_by()
   ...: print(mt.aggregate_cols(hl.agg.collect(mt.col_idx)))
2023-08-10 14:43:14.286 Hail: INFO: wrote matrix table with 3 rows and 3 columns in 3 partitions to /tmp/foo.mt
+---------+
| col_idx |
+---------+
|   int32 |
+---------+
|       2 |
|       1 |
|       0 |
+---------+
2023-08-10 14:43:16.338 Hail: INFO: Coerced sorted dataset
[0, 1, 2]
[2, 1, 0]

[danking]

I can't find good tests either, can you add some tests in the spirit of my shared ipython session?

done

[patrick-schultz]

Had to chase down a latent bug in StreamAgg, but I think everything should be sorted now (pun intended).

[danking]

Can you update your PR message with a CHANGELOG line indicating that we fixed the performance regression?

[danking]

oops, sorry, assert in Python doesn't use parentheses. My bad

done