[query] Enable use of aggregators on array expressions #12988
Conversation
hail/python/hail/ir/ir.py
Outdated
| tuple, uid, elt = unpack_uid(a.typ) | ||
| body = AggLet(value_name, elt, body) | ||
| body = with_split_rng_state(body, uid) | ||
| value_name = tuple |
There was a problem hiding this comment.
Isn't it possible for the body to use randomness in a value context? E.g. see TableAggregate.__init__.
There was a problem hiding this comment.
Check if how I just updated this is sane.
There was a problem hiding this comment.
Thanks Chris. That's almost right, but we still need to bind the elt (body expects value_name to be an element without uid). And lets keep using with_split_rng_state to keep that logic in one place:
tup, uid, elt = unpack_uid(a.typ)
body = AggLet(value_name, elt, body)
if body.uses_value_randomness:
body = Let('__rng_state', RNGStateLiteral(), body)
if body.uses_agg_randomness(is_scan=False):
body = with_split_rng_state(body, is_scan=False)
value_name = tup
There was a problem hiding this comment.
thank you, should be fixed now.
| @@ -238,6 +238,36 @@ def context(self): | |||
| return 'agg' | |||
|
|
|||
|
|
|||
| def aggregate_local_array(array, f): | |||
There was a problem hiding this comment.
We had discussed adding a CollectionExpression.aggregate method, similar to fold. Did you change your mind?
There was a problem hiding this comment.
nope, forgot over the weekend amid the presentation rush. Thanks for the reminder.
hail/python/hail/ir/ir.py
Outdated
| a = a.handle_randomness(body.uses_randomness) | ||
| if body.uses_randomness: | ||
| a = a.handle_randomness(body.uses_agg_randomness) | ||
| if body.uses_agg_randomness: | ||
| tup, uid, elt = unpack_uid(a.typ) | ||
| body = AggLet(value_name, elt, body) | ||
| if body.uses_value_randomness: | ||
| body = Let('__rng_state', RNGStateLiteral(), body) | ||
| if body.uses_agg_randomness(is_scan=False): | ||
| body = with_split_rng_state(body, is_scan=False) | ||
| body = AggLet(value_name, elt, body, is_scan=False) | ||
| body = with_split_rng_state(body, uid, is_scan=False) |
There was a problem hiding this comment.
The test caught some problems with this. Most substantially, we shouldn't be making a new root RNGStateLiteral for any value randomness. That means if this occurs inside some other looping context, value randomness will be identical on each outer iteration. Instead, we should just use the rng state bound in the outer context, which means we don't need to do anything here.
There was a problem hiding this comment.
Thanks for explaining. Clears up why this was necessary.
chrisvittal
changed the title
This method is useful for accessing functionality that exists in the aggregator library but not the basic expression library, for instance,
call_stats.