Factor out field access of `AbstractArray` and/or iterator interface

General Usage
,
1

Take

function foo(a)
    for x in a
        x.field # do something with x.field
    end
end

function foo2(a)
    for x in a
        x # do something with x directly
    end
end

The foo2 version accesses elements directly, and the foo accesses a field of the elements of an AbstractArray or iterator.

I would like to just define foo2, and factor out the field access in a way that is generic and performant.
In other words, I want to write foo(x) instead as foo2(tx(x2)) for some tx. I think the tx5 solution below is reasonably generic, but I’m unhappy with the performance.

Perhaps someone has a suggestion for how to do it better!

(https://gist.github.com/goretkin/3300946be7c698b6e2104852afcf98ed , attached below)

using Test
using BenchmarkTools
using DataStructures: OrderedDict

using MappedArrays: mappedarray

# use the iterator interface
function foo(a)
    r = nothing
    for x in a
        y = x.field
        if r === nothing
            r = y
        else
            r += y
        end
    end
    return r
end

# use the `AbstractArray` interface
function bar(a)
    T = fieldtype(eltype(typeof(a)), :field)
    return a[3].field + one(T)
end

function foo2(a)
    r = nothing
    for y in a
        if r === nothing
            r = y
        else
            r += y
        end
    end
    return r
end

function bar2(a)
    T = eltype(typeof(a))
    return a[3] + one(T)
end

a = [(field=i,) for i = 1:5]
b = ((field=i,) for i in Iterators.TakeWhile(<=(5), Iterators.Count(1,1)))  # has `Any` `eltype`
c = Iterators.TakeWhile(x->true, a) # has an `eltype`

@test foo(a) == 15
@test foo(b) == 15
@test foo(c) == 15
@test bar(a) == 4


# materialize an intermediate array
tx1(A, f) = getfield.(A, Ref(f))

@test foo2(tx1(a, :field)) == foo(a)
@test foo2(tx1(b, :field)) == foo(b)
@test foo2(tx1(c, :field)) == foo(c)
@test bar2(tx1(a, :field)) == bar(a)

# don't materialize an intermediate array
tx2(A, f) = (getfield(x, f) for x in A)

@test foo2(tx2(a, :field)) == foo(a)
@test foo2(tx2(b, :field)) == foo(b)
@test foo2(tx2(c, :field)) == foo(c)
@test_throws Any bar2(tx2(a, :field)) == bar(a) # but `tx2` doesn't preserve the `AbstractArray` interface

# don't materialize an intermediate array
tx3(A, f) = mappedarray(Base.Fix2(getfield, f), A)

@test foo2(tx3(a, :field)) == foo(a)
@test_throws Any foo2(tx3(b, :field)) == foo(b) # but `tx3` requires the `AbstractArray` interface
@test_throws Any foo2(tx3(c, :field)) == foo(c) # ditto
@test bar2(tx3(a, :field)) == bar(a)

lazy_map(f, A::AbstractArray) = mappedarray(f, A)
lazy_map(f, A) = (f(x) for x in A)

tx4(A, f) = lazy_map(Base.Fix2(getfield, f), A)

@test foo2(tx4(a, :field)) == foo(a)
@test foo2(tx4(b, :field)) == foo(b)
@test foo2(tx4(c, :field)) == foo(c)
@test bar2(tx4(a, :field)) == bar(a)
@show eltype(tx4(c, :field))            # but `tx4` doesn't have right `eltype` since generators do not infer `eltype`

# force the `eltype` of the generator to be G
struct TypedGenerator{T, G}
    g::G
end
Base.eltype(::Type{TypedGenerator{T, G}}) where {T, G} = T
Base.axes(g::TypedGenerator) = axes(g.g)
Base.collect(g::TypedGenerator) = collect(g.g)
Base.iterate(g::TypedGenerator, args...) = iterate(g.g, args...)
Base.length(g::TypedGenerator) = length(g.g)
Base.ndims(g::TypedGenerator) = ndims(g.g)
# Base.nextind(g::TypedGenerator, args...) = nextind(g.g, args...)    # AbstractTrees defines this
Base.size(g::TypedGenerator) = size(g.g)
Base.IteratorSize(::Type{TypedGenerator{T, G}}) where {T ,G} = Base.IteratorSize(G)

TypedGenerator{T}(g) where {T} = TypedGenerator{T, typeof(g)}(g)

tg = TypedGenerator{Int64}((x for x in Iterators.TakeWhile(<(5), Iterators.Count(1, 1))))
tg2 = TypedGenerator{Int64}((x for x in [1,2,3]))


fieldtype_any(T::Type{Any}, f) = Any
fieldtype_any(T::Type{<:Any}, f) = fieldtype(T, f)

lazy_map2(f, A::AbstractArray) = mappedarray(f, A)
lazy_map2(f::Base.Fix2{typeof(getfield)}, A) = TypedGenerator{fieldtype_any(eltype(A), f.x)}(f(x) for x in A)
lazy_map2(f::Base.Fix2{typeof(getfield)}, A::AbstractArray) = mappedarray(f, A) # resolve ambiguity

tx5(A, f) = lazy_map2(Base.Fix2(getfield, f), A)

@test foo2(tx5(a, :field)) == foo(a)
@test foo2(tx5(b, :field)) == foo(b)
@test foo2(tx5(c, :field)) == foo(c)
@test bar2(tx5(a, :field)) == bar(a)
@test eltype(tx5(c, :field)) == Int64            # `tx5` has right `eltype`


a1000 = [(field=i,) for i = 1:1000]
b1000 = ((field=i,) for i in Iterators.TakeWhile(<=(1000), Iterators.Count(1,1)))  # has `Any` `eltype`
c1000 = Iterators.TakeWhile(x->true, a1000) # has an `eltype`

B = OrderedDict{Any, Any}()

for (abc, abc_s) = zip((a1000, b1000, c1000), (:a, :b, :c))
    println("\n\n-------------------------------------------------")
    @show abc_s

    b = B[(abc_s, :unfactored)] = @benchmark foo($abc)

    println("unfactored")
    show(stdout, "text/plain", b); println(stdout)

    b = B[(abc_s, :factored)] = @benchmark foo2(tx5($abc, :field))

    println("\n\nfactored")
    show(stdout, "text/plain", b); println(stdout)
end

Gives timings which shows that the performance on Array is 5x slower:

julia> B
OrderedDict{Any,Any} with 6 entries:
  (:a, :unfactored) => Trial(95.608 ns)
  (:a, :factored)   => Trial(503.808 ns)
  (:b, :unfactored) => Trial(397.259 ns)
  (:b, :factored)   => Trial(430.373 ns)
  (:c, :unfactored) => Trial(549.579 ns)
  (:c, :factored)   => Trial(599.636 ns)

Those timings were after replacing all @propagate_inbounds with @inline @propagate_inbounds in MappedArrays.jl. That changes the timings, but I didn’t notice much of a difference after running a few times.

2

I am not sure if I understand the specs (the code doesn’t really do anything), but what about

foreach(x -> x.field, a) # foo
foreach(identity, a)     # foo2

?

3

(the code doesn’t really do anything)

The gist/code below has a foo/foo2 that does something with the elements.

4

A good idiom for this is to take a transformation function as the first argument (parameterized by a type parameter F to be sure that the compiler specializes on the specific function passed).

function foo(f::F, a) where {F<:Function}
    for x in a
        f(x) # do something with f(x)
    end
end
foo(a) = foo(identity, a)

Then you can use e.g. foo(x -> x.field, a) if you need field access, and more generally you can do any kind of transformation.

3 Likes
5

Thanks for mentioning this. I figured this might be the best way currently, which is why sum, and all, and so on can take a function as their first argument. But I was hoping for an alternative.

6

The 5X performance gap of tx5 can be eliminated: `iterate` fallback causes 5X performance drop · Issue #45 · JuliaArrays/MappedArrays.jl · GitHub

1 Like