Julia: Sample N groups from a grouped dataframe

Question

I have a dataframe of time series records of item product sales that I am using to plot, but there are many and I want a random sample of N.

Here is a simplified example of the data with three items, and I want to sample two of them randomly:

df = DataFrame(time = [0, 1, 0, 1, 0, 1]
    , amt = [19.00, 11.00, 35.50, 32.50, 5.99, 5.99]
    , item = ["B001", "B001", "B020", "B020", "BX00", "BX00"])

6×3 DataFrame
│ Row │ time  │ amt     │ item   │
│     │ Int64 │ Float64 │ String │
├─────┼───────┼─────────┼────────┤
│ 1   │ 0     │ 19.0    │ B001   │
│ 2   │ 1     │ 11.0    │ B001   │
│ 3   │ 0     │ 35.5    │ B020   │
│ 4   │ 1     │ 32.5    │ B020   │
│ 5   │ 0     │ 5.99    │ BX00   │
│ 6   │ 1     │ 5.99    │ BX00   │

I have found a solution after some study, but it doesn't seem like the simple way to express this.

# this attaches a random number to each group, sorts it, and then ranks each group:

using StatsBase

@pipe df |> groupby(_, :item) |>
     combine(_, :time, :amt, :item, :item => (x -> rand()) => :rando) |>
     sort(_, :rando) |>
     transform(_, :rando => denserank => :rnk_rnd)

6×5 DataFrame
│ Row │ item   │ time  │ amt     │ rando    │ rnk_rnd │
│     │ String │ Int64 │ Float64 │ Float64  │ Int64   │
├─────┼────────┼───────┼─────────┼──────────┼─────────┤
│ 1   │ B001   │ 0     │ 19.0    │ 0.449577 │ 1       │
│ 2   │ B001   │ 1     │ 11.0    │ 0.449577 │ 1       │
│ 3   │ BX00   │ 0     │ 5.99    │ 0.482569 │ 2       │
│ 4   │ BX00   │ 1     │ 5.99    │ 0.482569 │ 2       │
│ 5   │ B020   │ 0     │ 35.5    │ 0.612401 │ 3       │
│ 6   │ B020   │ 1     │ 32.5    │ 0.612401 │ 3       │


# I only need the original columns, and I'll filter for the first N=2 items from the re-constituted dataframe

@pipe ans |> filter(:rnk_rnd => <=(2), _)  |>
     select(_, :item, :time, :amt)

4×3 DataFrame
│ Row │ item   │ time  │ amt     │
│     │ String │ Int64 │ Float64 │
├─────┼────────┼───────┼─────────┤
│ 1   │ BX00   │ 0     │ 5.99    │
│ 2   │ BX00   │ 1     │ 5.99    │
│ 3   │ B001   │ 0     │ 19.0    │
│ 4   │ B001   │ 1     │ 11.0    │

# this is exactly what I'm looking for

Is there no other more compact way to take a random sample of groups from a grouped dataframe?

user_1269567 · Answer

An alternative is to use sample from StatsBase.jl:

@pipe df |>
      groupby(_, :item) |>
      _[sample(1:length(_), 2, replace=false)] |>
      DataFrame

if you accept a random fraction q from your DataFrame (not a fixed number) then it is even easier:

@pipe df |>
      groupby(_, :item) |>
      combine(sdf -> rand() < q ? sdf : DataFrame(), _)

user_2223505 · Answer

I picked up a more compact expression from an issue in DataFrames.jl

@pipe df |> 
    groupby(_, :item) |>
    _[shuffle(1:end)] |>
    combine(_[1:2], :)

Which results in my same randomly selected groups, back as a dataframe:

4×3 DataFrame
│ Row │ item   │ time  │ amt     │
│     │ String │ Int64 │ Float64 │
├─────┼────────┼───────┼─────────┤
│ 1   │ BX00   │ 0     │ 5.99    │
│ 2   │ BX00   │ 1     │ 5.99    │
│ 3   │ B020   │ 0     │ 35.5    │
│ 4   │ B020   │ 1     │ 32.5    │

I think there will eventually be a shuffle function for grouped df, if we all upvote the issue!

Julia: Sample N groups from a grouped dataframe

2 Answers