Define length(::Flatten) for some collections of arrays

[mcabbott]

As suggested at #23431, this defines length (and IteratorSize trait) for a few more Iterators.flatten objects: those flattening an arrays of array, or a tuples of arrays.

Computing the total length involves adding the lengths of the constituent arrays. At present length(::Flatten) never does this. For these cases it seems to be very cheap.

Knowing the length speeds up collect ∘ Iterators.flatten, although there is still room for improvement:

julia> vecs = [rand(100) for _ in 1:100];

julia> @btime length(Iterators.flatten($vecs))
  min 54.133 ns, mean 54.270 ns (0 allocations)  # this PR
10000

julia> @btime collect(Iterators.flatten($vecs));
  min 68.875 μs, mean 81.634 μs (9 allocations, 326.55 KiB) # master
  min 12.167 μs, mean 17.848 μs (2 allocations, 78.17 KiB)  # this PR

julia> @btime reduce(vcat, $vecs);
  min 2.403 μs, mean 8.332 μs (3 allocations, 79.05 KiB)  # ideal?
  
julia> tvecs = Tuple(rand(1000) for _ in 1:10);

julia> @btime collect(Iterators.flatten($tvecs));
  min 69.167 μs, mean 81.648 μs (9 allocations, 326.55 KiB)  # master
  min 12.833 μs, mean 22.149 μs (2 allocations, 78.17 KiB)   # this PR

julia> @btime vcat($tvecs...);
  min 1.780 μs, mean 11.348 μs (2 allocations, 78.17 KiB)  # ideal?

Some other cases with already known length remain quite slow:

julia> tups = [Tuple(rand(10)) for _ in 1:1000];

julia> Base.haslength(tups)  # unchanged, uses eltype of array
true

julia> v = @btime collect(Iterators.flatten($tups));
  min 117.583 μs, mean 131.905 μs (2 allocations, 78.17 KiB)

julia> v == vec(@btime stack($tups))
  min 8.597 μs, mean 18.170 μs (2 allocations, 78.17 KiB)
true

[LilithHafner]

For flattened tuples of AbstractArrays this seems quite sensible.

How is the performance for long vectors of short vectors?

[LilithHafner]

I don't see how HasEltype is relevant here, would this work too?

Suggested change

	_flatten_iteratorsize(sz, ::HasEltype, ::Type{<:Tuple{Vararg{AbstractArray}}}) = HasLength()
	_flatten_iteratorsize(sz, ::HasEltype, ::Type{<:AbstractArray{<:AbstractArray}}) = HasLength()
	_flatten_iteratorsize(sz, _, ::Type{<:Tuple{Vararg{AbstractArray}}}) = HasLength()
	_flatten_iteratorsize(sz, _, ::Type{<:AbstractArray{<:AbstractArray}}) = HasLength()

[mcabbott]

For a large enough array of tiny arrays, length takes longer. A possible torture test, in which collect still benefits greatly:

julia> let vs = [rand(2) for _ in 1:10^6]
         @btime sum(length, $vs)
         @btime sum(first, $vs)  # a cheap operation visiting all arrays
         @btime collect(Iterators.flatten($vs))
       end;
  2.020 ms (0 allocations: 0 bytes)  # length
  2.259 ms (0 allocations: 0 bytes)
  17.394 ms (15 allocations: 18.90 MiB)  # flatten, master

  6.333 ms (2 allocations: 15.26 MiB)    # flatten, this PR

StaticArrays provides a length based on type, in src/flatten.jl. In its present state this PR dispatches earlier & breaks that:

julia> using StaticArrays

julia> let vs = [@SVector(rand(2)) for _ in 1:10^6]
         @btime sum(length, $vs)
         @btime sum(first, $vs)
         @btime collect(Iterators.flatten($vs))
       end;
  5.535 μs (0 allocations: 0 bytes)  # special length(flatten(vs)) takes 2.5ns
  674.875 μs (0 allocations: 0 bytes)
  2.688 ms (2 allocations: 15.26 MiB)  # flatten, master

  3.119 ms (2 allocations: 15.26 MiB)  # flatten, this PR

[LilithHafner]

To fix the interaction with StaticArrays it would work to define the changes here at the flatten_length method instead of the length method, or to make a PR to StaticArrays that changes he definition there to the length method. flatten_length is internal, so it is acceptable to simply break StaticArrays's dispatch, but it is polite to make PR that fixes it too.

Either way, weshould probably define length and IteratorSize in the same way (i.e. using both flatten_length and flatten_iteratorsize or neither)

[fredcallaway]

Bump on this. From a user perspective, it seems really weird that length(::Flatten) doesn't work when flattening any iterators that have known length. I think it shouldn't be limited to arrays and tuples though, e.g. UnitRange and CartesianIndices should work as well. From my perspective, it's about semantics, not just efficiency for things like map.

[nsajko]

I think this is not in general type stable when f.it is empty. Is there even a way to make it type stable?

[nsajko]

Because length(x) isa Int is not necessarily true.

[LilithHafner]

I'm willing to accept non-type-stable length(::Flatten) for inputs that have lengths that are neither Int nor smaller integers that promote to Int when added to 0.

length almost always returns Int.

[nsajko]

julia> length(big(1):big(2)) isa Int
false

UnitRange{BigInt} does not seem that exceptional, for what that is worth. Both UnitRange and BigInt are exported from Base and commonly used.

[nsajko]

This method relies for correctness on the assumption that any subtype of AbstractArray is a stateless iterator. I am not completely sure if that assumption is justified. A stateful AbstractArray would certainly be uncommon and weird, but perhaps still valid?

[nsajko]

xref:

• HasLength for even more Flatten iterators #23431 (comment)

• Introduce iterator trait: IsStateful #43388