JuliaLang / julia

The Julia Programming Language
https://julialang.org/
MIT License
45.56k stars 5.47k forks source link

Multi index `Base.setindex` #54508

Open rayegun opened 4 months ago

rayegun commented 4 months ago

As far as I can tell there is no multi-index Base.setindex such as setindex((1,2,3,4,5), (10, 20), 1:2) or setindex((1,2,3,4,5), (10, 40), [1, 4])

I'm not sure what the most performant method for this is, perhaps something like:

return ntuple(j -> if else(j in i, v[j], x[j]), Val{N}())
arhik commented 4 months ago

If you are implying for setindex rather than inplace setindex! version. No.

You can do it with Accesssors.jl.

julia> using Accessors

julia> a = (1, 2, 3, 4, 5)
(1, 2, 3, 4, 5)

julia> @set a[[2, 5]] = (-1, -1)
(1, -1, 3, 4, -1)

julia> using Chairmarks

julia> @b @set a[[2, 5]] = (-1, -1)
41.047 ns (2.02 allocs: 128.760 bytes)
nsajko commented 4 months ago

Base.setindex itself isn't public API. Although perhaps it should be?

Guess I'll make a package.

arhik commented 4 months ago

I quickly tried something like this below, but it is not going to be performant version.


g(c) = j -> (j in i) && (c += 1) > 0 ? v[c] : x[j]
ntuple(g(0), Val{N}()) 

julia> @b setindex(x, v, i)
391.071 ns (4.11 allocs: 131.657 bytes)
g(c) = j -> (j in i) && (c += 1) > 0 ? v[c] : x[j]
ntuple(g(0),  length(x)) 

Second is relatively more performant but still 10x slower than accessors.jl version above.

arhik commented 4 months ago

Also this scanning version is better

function setindex(x::Tuple, v::Tuple, i::Tuple)
    N = length(x)
    is = ntuple(j ->  ifelse(j in i, 1, 0), Val{N}())
    vs = ntuple(Val{N}()) do j
        if is[j] === 1
            return v[sum(is[1:j])] 
        else
            return x[j]
        end
    end
    return vs
end

julia> @b setindex(x, v, i)
70.008 ns (3.02 allocs: 112.650 bytes)
arhik commented 4 months ago

simply iterating through works too

function setindex(x::Tuple, v::Tuple, i::Tuple)
    out = x
    for (val, idx) in zip(v, i)
        out = Base.setindex(out, val, idx)
    end
    return out
end

julia> @b setindex(x, v, i)
22.901 ns (1.01 allocs: 48.403 bytes)

This is better than Accesssors.jl version.

Onesweep algorithm is possible I think. That can give best solution. I couldn't come up with one yet.

arhik commented 4 months ago

Since single setindex in base is taking roughly 16ns, then

julia> @b Base.setindex(x, 1, 4)
15.798 ns (1.01 allocs: 48.424 bytes)

This is best so far.

simply iterating through works too

function setindex(x::Tuple, v::Tuple, i::Tuple)
  out = x
  for (val, idx) in zip(v, i)
      out = Base.setindex(out, val, idx)
  end
  return out
end

julia> @b setindex(x, v, i)
22.901 ns (1.01 allocs: 48.403 bytes)

This is better than Accesssors.jl version.

Onesweep algorithm is possible I think. That can give best solution. I couldn't come up with one yet.