Possible feature: Improved indexing for GroupedDataFrames

[pdeffebach]

Here is some code I wrote for indexing into a GroupedDataFrame in cases where you want all the values corresponding to one key, or some combination thereof.

julia> using DataFrames;

julia> df = DataFrame(x = [1, 1, 2, 2], y = [200, 100, 200, 100], z = rand(4));

julia> gd = groupby(df, [:x, :y]);

julia> struct Groups{T}
           nt::T
       end;

julia> function Groups(; kwargs...)
           Groups((; kwargs...,))
       end;

julia> function Base.getindex(gd::GroupedDataFrame, x::Groups)
           ks = keys(gd)
           new_indices = Int[]
           nt = x.nt
           k1 = first(ks)
           nms = propertynames(k1)
           vals = map(nms) do nm
               if nm in keys(nt)
                   nt[nm]
               else
                   Colon()
               end
           end
           nt_cols_added = NamedTuple{(nms...,)}((vals...,)) 
           for k in ks
               to_push = true
               for nm in nms
                   if k[nm] != nt_cols_added[nm] && !(nt_cols_added[nm] isa Colon) && !(k[nm] in nt_cols_added[nm])
                       to_push = false
                   end
               end
               if to_push
                   inds = parentindices(gd[k])[1]
                   append!(new_indices, inds)
               end
           end
           view(parent(gd), new_indices, :)
       end;

julia> gd[Groups(x = 1, y = 200)]
1×3 SubDataFrame
 Row │ x      y      z        
     │ Int64  Int64  Float64  
─────┼────────────────────────
   1 │     1    200  0.423354

julia> gd[Groups(x = 1, y = :)]
2×3 SubDataFrame
 Row │ x      y      z        
     │ Int64  Int64  Float64  
─────┼────────────────────────
   1 │     1    200  0.423354
   2 │     1    100  0.542757

julia> gd[Groups(x = 1, y = [100, 200])]
2×3 SubDataFrame
 Row │ x      y      z        
     │ Int64  Int64  Float64  
─────┼────────────────────────
   1 │     1    200  0.423354
   2 │     1    100  0.542757

julia> gd[Groups(x = 1)]
2×3 SubDataFrame
 Row │ x      y      z        
     │ Int64  Int64  Float64  
─────┼────────────────────────
   1 │     1    200  0.423354
   2 │     1    100  0.542757