Is there a non-allocating version of findall?

General Usage Performance
,
1
MWE 
using StaticArrays, BenchmarkTools

# Data vector
dat = collect(1:10);

# Always five "trues" but at any indices
flag = [trues(3); falses(5); trues(2)];

# Index with vector of bools
@btime ($dat[$flag]);
@btime (@views $dat[$flag]);
@btime SVector{5}($dat[$flag]);
@btime SVector{5}(@views $dat[$flag]);


# Index with vector of Ints
idx = @btime findall($flag);
@btime ($dat[$idx]);
@btime (@views $dat[$idx]);
@btime SVector{5}($dat[$idx]);
@btime SVector{5}(@views $dat[$idx]);

# Comprehensions
@btime ($dat[i] for i in $idx);
@btime SVector{5}($dat[i] for i in $idx);

I have two vectors like below and want to subset the first to get only the indices flagged as true in the second, without any heap allocations:

# Data vector
dat = collect(1:10);

# Always five "trues" but at any indices
flag = [trues(3); falses(5); trues(2)];

Index with vector of Bools:

julia> @btime ($dat[$flag]);
  51.061 ns (1 allocation: 96 bytes)

julia> @btime (@views $dat[$flag]);
  80.513 ns (1 allocation: 96 bytes)

julia> @btime SVector{5}($dat[$flag]);
  54.985 ns (1 allocation: 96 bytes)

julia> @btime SVector{5}(@views $dat[$flag]);
 88.470 ns (1 allocation: 96 bytes)

Index with vector of Ints:

julia> idx = @btime findall($flag);
  43.921 ns (1 allocation: 96 bytes)

julia> @btime ($dat[$idx]);
  50.275 ns (1 allocation: 96 bytes)

julia> @btime (@views $dat[$idx]);
  9.700 ns (0 allocations: 0 bytes)

julia> @btime SVector{5}($dat[$idx]);
  53.179 ns (1 allocation: 96 bytes)

julia> @btime SVector{5}(@views $dat[$idx]);
  11.041 ns (0 allocations: 0 bytes)

Comprehensions:

julia> @btime ($dat[i] for i in $idx);
  3.132 ns (0 allocations: 0 bytes)

julia> @btime SVector{5}($dat[i] for i in $idx);
  7.366 ns (0 allocations: 0 bytes)

Some of the latter methods work, but only move the allocation to the line idx = findall($flag).

It seems I am looking for a non-allocating version of findall. Does this exist?

2

As you have shown, collecting a generator into an SVector does not allocate, i.e., this works for finding the indices as well:

julia> find5(flag) = SVector{5}(i for (i, b) in enumerate(flag) if b);

julia> @btime SVector{5}(@view $dat[find5($flag)]);
  40.874 ns (0 allocations: 0 bytes)

# Seems like the indices being a static vector also gives a static result
julia> @btime $dat[find5($flag)];
  34.445 ns (0 allocations: 0 bytes)
1 Like