Skip to content

Support indexing with generators #37648

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
+36 −0
Open

Support indexing with generators #37648

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0c7856b
Support indexing with generators
omus Sep 18, 2020
f40839b
Limit indexing to AbstractArray
omus Sep 28, 2020
bad925e
Support all iterators again
omus Oct 9, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions NEWS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ New library features
* New constructor `NamedTuple(iterator)` that constructs a named tuple from a key-value pair iterator.
* A new `reinterpret(reshape, T, a::AbstractArray{S})` reinterprets `a` to have eltype `T` while potentially
inserting or consuming the first dimension depending on the ratio of `sizeof(T)` and `sizeof(S)`.
* Generators support `getindex`, `firstindex`, and `lastindex` ([#37648]).

Standard library changes
------------------------
Expand Down
22 changes: 22 additions & 0 deletions base/generator.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,28 @@ size(g::Generator) = size(g.iter)
axes(g::Generator) = axes(g.iter)
ndims(g::Generator) = ndims(g.iter)

function getindex(g::Generator{<:AbstractArray}, I...)
I′ = to_indices(g.iter, I)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

to_indices doesn't work on non-array indexables like NamedTuple and Dict. As I commented earlier, I don't think we should support vectorized indexing.

You can use LazyArrays.jl or MappedArrays.jl to get a full array API. If we were to add vectorized indexing for lazy mapping in Base, Broadcasted is probably a better interface to do it since it's already very array-like.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As I commented earlier, I don't think we should support vectorized indexing.

The tricky part of your statement is how would we limit getindex to just use scalar indexing?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just define getindex(::Generator, ::Integer...) and nothing else. I also think this is the right thing to do.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how would we limit getindex to just use scalar indexing?

Ah, good point. How about introducing

function to_scalar_indices(A, I)
    J = to_indices(A, I)
    J isa Tuple{Vararg{Integer}} || error("expected scalar indices. got: ", I)
    return J
end

scalar_getindex(A::AbstractArray, I...) = getindex(A, to_scalar_indices(A, I)...)
scalar_getindex(A, I...) = getindex(A, I...)

and use it for Generator?

Just define getindex(::Generator, ::Integer...)

Isn't it too restrictive and a half-way solution if it were to restrict the API? For example, it'd mean x[1] for x = Iterators.map(string, (a=1, b=2)) works but x[:a] doesn't.

Also, I guess we need a different implementation for AbstractDict or error out? Applying f after getindex is not correct for AbstractDict.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just define getindex(::Generator, ::Integer...)

Mostly that works but we get into some strange corner cases with AbstractDict:

julia> function Base.getindex(g::Base.Generator, I::Integer...)
           g.f(g.iter[I...])
       end

julia> g = (x * y for (x, y) in Dict(3 => 4))
Base.Generator{Dict{Int64,Int64},var"#1#2"}(var"#1#2"(),
Dict(3 => 4))

julia> collect(g)[1]
12

julia> g[1]
ERROR: KeyError: key 1 not found
Stacktrace:
 [1] getindex at ./dict.jl:467 [inlined]
 [2] getindex(::Base.Generator{Dict{Int64,Int64},var"#1#2"}, ::Int64) at ./REPL[1]:2
 [3] top-level scope at REPL[3]:1

This was the main reason I ended up restricting the logic to AbstractArray

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we need a different implementation for AbstractDict or error out

That was also my conclusion. A first pass at the AbstractDict version looks like:

function Base.getindex(g::Base.Generator{<:AbstractDict}, I...)
    subset = collect(pairs(g.iter))[I...]
    if subset isa Pair
        g.f(subset)
    else
        map(g.f, subset)
    end
end

Copy link
Contributor

@rapus95 rapus95 Dec 17, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As in my #37648 (comment) to this whole PR, why do you want to support AbstractDict in general?
IMO Generator should behave like a lazy map. And for map we have:

julia> map(x->println(x), Dict(:a=>2, :c=>"h"))
ERROR: map is not defined on dictionaries
Stacktrace:
 [1] error(::String) at .\error.jl:33
 [2] map(::Function, ::Dict{Symbol,Any}) at .\abstractarray.jl:2190
 [3] top-level scope at REPL[10]:1

So I'd definitively error out, optimally with the same error.

subset = g.iter[I′...]
if isempty(index_shape(I′...))
g.f(subset)
else
map(g.f, subset)
end
end

function getindex(g::Generator, I...)
A = collect(g.iter)
subset = A[I...]
if typeof(subset) == eltype(A)
g.f(subset)
else
map(g.f, subset)
end
end

firstindex(g::Generator) = firstindex(g.iter)
lastindex(g::Generator) = lastindex(g.iter)

## iterator traits

Expand Down
13 changes: 13 additions & 0 deletions test/functional.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,3 +219,16 @@ end
let (:)(a,b) = (i for i in Base.:(:)(1,10) if i%2==0)
@test Int8[ i for i = 1:2 ] == [2,4,6,8,10]
end

@testset "indexing" begin
g = (i % 2 == 0 ? error("invalid") : i for i in 1:3)

@test g[3] == 3
@test g[1:2:3] == [1, 3]
@test firstindex(g) == 1
@test lastindex(g) == 3

@test (tuple(x) for x in ["a"])[1] == ("a",)
@test (tuple(x) for x in [[0]])[1] == ([0],)
@test (x * y for (x, y) in Dict(3 => 4))[1] == 12
end