What's the best way to optimize findmax with a function?

maxlength(l) = findmax(map(length, l))
When I look at the compiled code, I see a reference to collect_similar, so it is materializing the map result. I don’t want that inefficiency.

I know I can write the equivalent using foldl, and get the call to length inside the loop for foldl. But that is not very expressive. I see that many functions over iterables, such as maximum, take a function argument for this purpose. So it would be okay with me if there were a findmax method that took a function as first argument. If that’s the Julia way of dealing with this sort of issue, then I’m happy to submit a PR to implement it.

But the language designer in me cringes at the proliferation of such methods. I would prefer that the code I wrote above, would optimize the function call into the loop, so the result of the map does not get materialized. Then I would consider deprecating methods like the maximum and sum methods that take a function argument, since maximum(f, i) === maximum(map(f, i))

I will understand if having methods that take a function argument are preferred, since they work well with do.

If there is agreement that optimizing map in this way is desirable, I might consider implementing it. Many years ago I was a compiler writer. I have not looked at the Julia compiler. I would need someone to point me in the right direction. I would try to generalize this optimization, but I’d need guidance with that.

Anyway – I want to be able to write either the code at the top, or findmax(length, l). Enough that I’m willing to do the work, if you tell me which way is preferable for Julia.

There is a PR for such a findmax that seems to be almost done: https://github.com/JuliaLang/julia/pull/35316. But another thing to consider is that map semantically materializes, while Iterators.map does not. So you can probably get the behavior you want just by findmax(Iterators.map(length, l)). Note that the Iterators submodule of Base provides lazy versions for other functions in Base too (and the package IterTools.jl has some more).

One likely cannot completely remove the materialization as a compiler pass, because sometimes different algorithms are used. For example, given a vector v, sum(map(x -> x^2, v)) will materialize the vector and then apply sum, which will use pairwise summation (since its input is a vector). But sum(Iterators.map(x -> x^2, v)) (or sum(x -> x^2, v)) will apply a naive accumulation, which might have more floating point error. (I did not actually check these statements so please correct me if I’m wrong). So the compiler won’t convert the first into the second, because you might be relying on pairwise summation. A hypothetical super smart compiler might be able to remove details like the actual memory allocation for the vector, etc, but semantically it has to use the vector method/algorithm, since that’s what you told it to do. I’m not sure to what extent this applies to findmax but I hope it helps explain why Julia doesn’t always do these things automatically.

edit: by the way, the findmax code from the PR (by @ColinCaine) is very simple, so it’s easy to just copy it into your own projects until the PR makes it through to a Julia release:

findmax(f, domain) = mapfoldl(x -> (f(x), x), _rf_findmax, domain)
_rf_findmax((fm, m), (fx, x)) = isless(fm, fx) ? (fx, x) : (fm, m)

Thanks for the quick reply, and the detailed explanation! I tried Iterators.map, and it still materializes the map result. (Calls collect_to_with_first!)

I’ll look at the PR you mentioned, and use the code from there.

1 Like

Ah, good to know! Seems like it shouldn’t, maybe there’s something that can be fixed/improved there (or maybe my understanding is off).

what about

maxlength2(l) = begin
  ml = maximum(length, l)
  ml, findfirst(x->length(x) == ml, l)

and I wonder why there isn’t a findmax(fn, l) implementation

Thanks for the suggestion! That doesn’t have a map result to materialize, but it is a little inefficient in requiring a second scan of the input.

@ericphanson pointed out that there is a PR to add findmax(fn, l) to Julia, and gave the relevant code from that PR. That code does work. The only thing I could complain about is that there is still a memory allocation (call to jl_gc_pool_alloc). Probably making the state for the foldl. In some future optimization, maybe that state can be consed on the stack. But it’s plenty good for now.

Thanks again, @ericphanson!

julia> map == Iterators.map

Ah, it doesn’t exist! I must’ve been thinking of IterTools.imap then.

Did a bunch of benchmarks and the very original way seems the fastest.

using Random

l = [randstring(x) for x in rand(UInt8, 1000)]

function maxlen_reduc((last_len, last_val), val)
    new_len = length(val)
    (new_len > last_len) && return new_len, val
    return last_len, val

maxlength2(l) = begin
  ml = maximum(length, l)
  ml, findfirst(x->length(x) == ml, l)

maxlength(l) = findmax(map(length, l))

using BenchmarkTools
@benchmark reduce(maxlen_reduc, l; init=(length(l[1]), l[1]))

@benchmark foldl(maxlen_reduc, l; init=(length(l[1]), l[1]))

@benchmark maxlength(l)

@benchmark maxlength2(l)

import Base: findmax
findmax(f, domain) = mapfoldl(x -> (f(x), x), _rf_findmax, domain)
_rf_findmax((fm, m), (fx, x)) = isless(fm, fx) ? (fx, x) : (fm, m)

@benchmark findmax(length, l)

using MappedArrays
@benchmark findmax(mappedarray(length, l))

function manualmaxlen(l)
    maxl = l[1]
    len = length(maxl)
    for m in l
        lm = length(m)
        if lm > len
            len = lm
            maxl = m
    maxl, len

@benchmark manualmaxlen(l)