Is there the equivalent to find all indexes for an Array Expression?

Hail Query & hailctl
1

I was looking for something that will return me the indices (indexes?) for ALL members of an array, based on some arbitrary function…

A bit like how index returns the FIRST occurence of a value:

s3 = hl.literal(['Alice','Bob','Charlie'])
hl.eval(s3.index(lambda x: x.contains('l')))
0

The equivalent i would assume something like this

hl.eval(s3.all_index(lambda x: x.contains('l')))
[0, 2]

Am i missing something?

I went ahead and implemented it myself, since it was trivial once I found the code for index, but somehow I feel I am missing something obvious that is already there…

def all_index(self, x):
    """Returns the all indexes of `x`, or missing.

    Parameters
    ----------
    x : :class:`.Expression` or :obj:`typing.Callable`
        Value to find, or function from element to Boolean expression.

    Returns
    -------
    :class:`.ArrayExpression`

    Examples
    --------
    >>> hl.eval(names.all_index('Bob'))
    1

    >>> hl.eval(names.all_index('Beth'))
    None

    >>> hl.eval(names.all_index(lambda x: x.contains('l')))
    [0, 3]

    >>> hl.eval(names.all_index(lambda x: x.endswith('h')))
    None
    """
    if callable(x):
        def f(elt, x):
            return x(elt)
    else:
        def f(elt, x):
            return elt == x
    return hl.bind(lambda a: hl.range(0, a.length()).filter(lambda i: f(a[i], x)), self)

and then:

hl.expr.expressions.typed_expressions.ArrayExpression.all_index = all_index

2

There’s no single function to do this, but quite easy to compose:

hl.enumerate(a)\
  .filter(lambda idx_and_element: filter_func(idx_and_element[1]))\
  .map(lambda idx_and_element: element[1])
3

Ah…yeah, I was looking for a way to enumerate, since I figured something like that woudl be needed althoug I must have overlook it (may have something to do with the website help being a little broken, yesterday :slight_smile:

How would you rate my implementation I provided wrs speed and scalability etc., compared to yours? Feels pretty similar although all_index takes a function as a argument, so slightly more general, but maybe more costly?

4

Should be identical, right now. At some point we’ll improve code generation of fused array operations such that my version will be a little faster in some cases (when a doesn’t need to be realized as an intermediate, just iterated over once), and then some point after that, we’ll improve our optimizer so your code generates the same thing as mine.

1 Like