# Help writing a \`collect\` for iterators

**URL:** <https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494>\
**Category:** General Usage\
**Tags:** question, iterators\
**Created:** [July 25, 2024, 8:36pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494 "2024-07-25T20:36:06Z")\
**Posts on this page:** 12\
**Page:** 1

<div class="post-metadata">

**Author:** ![juliohm](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juliohm/32/215266_2.png) [@juliohm](https://discourse.julialang.org/u/juliohm)\
**Post date:** [July 25, 2024, 8:36pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/1 "2024-07-25T20:36:06Z")

</div>

Here is a task that I am having a hard time to accomplish with Julia Base:

> Given an iterator `iter` and a set of linear indices `inds`, write a function `collectat(iter, inds)` that collects the iterator at the indices without materializing the full list of items.

I asked this question on Zulip some time ago and @Mason helped with a solution that depends on Transducers.jl:

```julia
using Transducers: Filter, Map, TakeWhile, tcollect, ⨟

function collectat(iter, inds)
  if isempty(inds)
    eltype(iter)[]
  else
    selectat(inds) = enumerate ⨟ TakeWhile(x -> first(x) ≤ last(inds)) ⨟ Filter(y -> first(y) ∈ inds) ⨟ Map(last)
    iter |> selectat(inds) |> tcollect
  end
end

```

Here is an example of usage:

```julia
julia> iter = (rand() for _ in 1:200_000_000)
Base.Generator{UnitRange{Int64}, var"#8#9"}(var"#8#9"(), 1:200000000)

julia> inds = [1,1_000_000,1_999_999]
3-element Vector{Int64}:
       1
 1000000
 1999999

julia> @allocated collectat(iter, inds)
4608

```

Is it possible to write something that works in most cases without Transducers.jl?

---

<div class="post-metadata">

**Author:** ![aplavin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aplavin/32/222056_2.png) [@aplavin](https://discourse.julialang.org/u/aplavin)\
**Post date:** [July 25, 2024, 9:21pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/2 "2024-07-25T21:21:01Z")

</div>

Is the Transducers solution different from this one using just `Iterators`?

```julia
collectat_itr(itr, idxs) = @p let
    enumerate(itr)
    Iterators.takewhile(first(_) ≤ last(idxs))
    Iterators.filter(first(_) ∈ idxs)
    map(last)
end

```

(copied from that zulip thread)

---

<div class="post-metadata">

**Author:** ![rafael.guerra](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rafael.guerra/32/216610_2.png) [@rafael.guerra](https://discourse.julialang.org/u/rafael.guerra)\
**Post date:** [July 25, 2024, 9:41pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/3 "2024-07-25T21:41:02Z")

</div>

@aplavin, how would you write that function without DataPipes.jl?

Like this?

```julia
function collectat_itr2(itr, idxs)
    e = enumerate(itr)
    it = Iterators.takewhile(x->(first(x) ≤ last(idxs)), e)
    f = Iterators.filter(x->(first(x) ∈ idxs), it)
    return map(last, f)
end

```

---

<div class="post-metadata">

**Author:** ![aplavin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aplavin/32/222056_2.png) [@aplavin](https://discourse.julialang.org/u/aplavin)\
**Post date:** [July 25, 2024, 10:05pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/4 "2024-07-25T22:05:23Z")

</div>

Yeah, exactly – the translation back and forth is really straightforward (:  
Note that some parens in your example aren’t necessary, eg `x->(first(x) ≤ last(idxs))` could be just `x->first(x) ≤ last(idxs)`.

---

<div class="post-metadata">

**Author:** ![juliohm](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juliohm/32/215266_2.png) [@juliohm](https://discourse.julialang.org/u/juliohm)\
**Post date:** [July 25, 2024, 10:20pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/5 "2024-07-25T22:20:24Z")

</div>

That is perfect! Thank you all!

---

<div class="post-metadata">

**Author:** ![sadish-d](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sadish-d/32/48058_2.png) [@sadish-d](https://discourse.julialang.org/u/sadish-d)\
**Post date:** [July 26, 2024, 2:19am UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/6 "2024-07-26T02:19:27Z")

</div>

I think the solution returns a different number of elements depending on the order of the indices. If that’s not what you want I guess you’d do something like `collectat_itr3`? It returns the correct number of elements but ignores the order of indices (preserving the order of elements in the iterator).

Edit: I guess when the OP said linear indices, they meant ordered indices, in which case `collectat_itr2` makes sense. I’ll still leave my comment here to draw attention to it.

```julia
struct I
    n::Int
    I(n::Int) = n >= 0 ? new(n) : error("Expected integer >=0. Got: $n")
end
Iterators.iterate(i::I, state::Int = i.n) = state == 0 ? nothing : (state, state - 1)

function collectat_itr2(itr, idxs)
    e = enumerate(itr)
    it = Iterators.takewhile(x->(first(x) ≤ last(idxs)), e)
    f = Iterators.filter(x->(first(x) ∈ idxs), it)
    return map(last, f)
end

function collectat_itr3(itr, idxs)
    e = enumerate(itr)
    f = Iterators.filter(x->(first(x) ∈ idxs), e)
    return map(last, f)
end

@assert collectat_itr2(I(9), [1, 2]) == [9, 8]
@assert collectat_itr2(I(9), [2, 1]) == [9]

@assert collectat_itr3(I(9), [1, 2]) == [9, 8]
@assert collectat_itr3(I(9), [2, 1]) == [9, 8]

```

---

<div class="post-metadata">

**Author:** ![bertschi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bertschi/32/33462_2.png) [@bertschi](https://discourse.julialang.org/u/bertschi)\
**Post date:** [July 26, 2024, 6:29am UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/7 "2024-07-26T06:29:08Z")

</div>

Good point, for unordered indices, I would probably do something like this:

```julia
function collectat_itr4(itr, idxs)
    idxs = Set(idxs) # Want to check unordered membership
    itr |> it ->
      Iterators.take(it, maximum(idxs)) |> # just take what we might need
      enumerate |> it ->
      Iterators.filter(x->(first(x) ∈ idxs), it) |> it ->
      map(last, it)
end

```

Using `take` or `takewhile` also has the nice feature that it works on infinite iterators then, i.e., only taking the finite indices that are requested.

---

<div class="post-metadata">

**Author:** ![juliohm](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juliohm/32/215266_2.png) [@juliohm](https://discourse.julialang.org/u/juliohm)\
**Post date:** [July 26, 2024, 9:40am UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/8 "2024-07-26T09:40:15Z")

</div>

Thank you for raising the issue @sadish-d. By linear indices I meant `LinearIndices`. I will update the answer to make sure that all elements are returned.

---

<div class="post-metadata">

**Author:** ![juliohm](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juliohm/32/215266_2.png) [@juliohm](https://discourse.julialang.org/u/juliohm)\
**Post date:** [July 26, 2024, 9:53am UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/9 "2024-07-26T09:53:51Z")

</div>

Final solution:

```julia
function collectat(iter, inds)
  if isempty(inds)
    eltype(iter)[]
  else
    m = maximum(inds)
    e = Iterators.enumerate(iter)
    w = Iterators.takewhile(x -> (first(x) ≤ m), e)
    f = Iterators.filter(x -> (first(x) ∈ inds), w)
    map(last, f)
  end
end

```

---

<div class="post-metadata">

**Author:** ![sadish-d](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sadish-d/32/48058_2.png) [@sadish-d](https://discourse.julialang.org/u/sadish-d)\
**Post date:** [July 26, 2024, 6:19pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/10 "2024-07-26T18:19:17Z")

</div>

This still preserves the order of `iter` and ignores the order of indices in `inds`. It does not iterate over `inds` one by one, it iterates over `iter`.

---

<div class="post-metadata">

**Author:** ![juliohm](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juliohm/32/215266_2.png) [@juliohm](https://discourse.julialang.org/u/juliohm)\
**Post date:** [July 26, 2024, 6:36pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/11 "2024-07-26T18:36:16Z")

</div>

You are correct @sadish-d , I removed the word “ordered” in my previous comment before the solution to avoid confusion.

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [July 26, 2024, 8:38pm UTC](https://discourse.julialang.org/t/help-writing-a-collect-for-iterators/117494/12 "2024-07-26T20:38:58Z")

</div>

In case the ordered indicators are important (and with slight improvement for non-ordered case):

```julia
function collectat2(iter, inds)
    isempty(inds) && return eltype(iter)[]
    wassorted = issorted(inds)
    if !wassorted
        perm = sortperm(inds)
        iperm = invperm(perm)
    end
    L, T, j = length(inds), eltype(iter), 1
    M = wassorted ? last(inds) : inds[last(perm)]
    res = Vector{T}(undef, L)
    for (i,v) in enumerate(Iterators.take(iter, M))
        if i == inds[wassorted ? j : perm[j]]
            res[wassorted ? j : iperm[j]] = v
            j == L && break
            j += 1
        end
    end
    return res
end

```

and some timings:

```julia
julia> @btime collectat(1:10,[2,3,4,5]);
  104.165 ns (4 allocations: 304 bytes)

julia> @btime collectat2(1:10,[2,3,4,5]);
  34.675 ns (2 allocations: 192 bytes)

```