juliajulia-dataframe

Stratified Sampling of a DataFrame


Given a dataframe with columns "a", "b", and "value", I'd like to sample N rows from each pair of ("a", "b"). In python pandas, this is easy to do with the following syntax:

import pandas as pd
df.groupby(["a", "b"]).sample(n=10)

In Julia, I found a way to achieve something similar using:

using DataFrames, StatsBase

combine(groupby(df, [:a, :b]),
names(df) .=> sample .=> names(df)
)

However, I don't know how to extend this to n>1. I tried

combine(groupby(df, [:a, :b]),
names(df) .=> x -> sample(x, n) .=> names(df)
)

but this returned the error (for n=3):

DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 3 and 7")

One method (with slightly different syntax) that I found was:

combine(groupby(df, [:a, :b]), x -> x[sample(1:nrow(x), n), :])

but I'm interested in knowing if there are better alternatives


Solution

  • Maybe as an additional comment. If you have an id column in your data frame (holding a row number) then:

    df[combine(groupby(df, [:a, :b]), :id => (x -> rand(x, n)) => :id).id, :]
    

    will be a bit faster (but not by much).

    Here is an example:

    using DataFrames
    n = 10
    df = DataFrame(a=rand(1:1000, 10^8), b=rand(1:1000, 10^8), id=1:10^8)
    combine(groupby(df, [:a, :b]), x -> x[rand(1:nrow(x), n), :]); # around 16.5 seconds on my laptop
    df[combine(groupby(df, [:a, :b]), :id => (x -> rand(x, n)) => :id).id, :]; # around 14 seconds on my laptop