Группы выборки в Pandas

Скажем, я хочу сделать стратифицированный образец из фреймворка данных в Pandas, чтобы получить 5% строк для каждого значения данного столбца. Как я могу это сделать?

Например, в приведенной ниже таблице данных я хотел бы попробовать 5% строк, связанных с каждым значением столбца Z. Есть ли способ отобразить группы из загружаемой в память фрейма данных?

> df 

   X   Y  Z
   1 123  a
   2  89  b
   1 234  a
   4 893  a
   6 234  b
   2 893  b
   3 200  c
   5 583  c
   2 583  c
   6 100  c

В общем, что, если я этот фреймворк данных на диске в огромном файле (например, 8 ГБ файла csv). Есть ли способ сделать эту выборку, не загружая весь блок данных в памяти?

Ответ 1

Как насчет загрузки только столбца "Z" в память с помощью опции "usecols". Скажем, файл sample.csv. Это должно использовать гораздо меньше памяти, если у вас есть куча столбцов. Тогда, полагая, что это вписывается в память, я думаю, что это сработает для вас.

stratfraction = 0.05
#Load only the Z column
df = pd.read_csv('sample.csv', usecols = ['Z'])
#Generate the counts per value of Z
df['Obs']  = 1
gp = df.groupby('Z')
#Get number of samples per group 
df2 = np.ceil(gp.count()*stratfraction)
#Generate the indices of the request sample (first entrie)
stratsample = []
for i, key in enumerate(gp.groups):
    FirstFracEntries = gp.groups[key][0:int(df2['Obs'][i])]
    stratsample.extend(FirstFracEntries) 
#Generate a list of rows to skip since read_csv doesn't have a rows to keep option
stratsample.sort
RowsToSkip = set(df.index.values).difference(stratsample)
#Load only the requested rows (no idea how well this works for a really giant list though)         
df3 = df = pd.read_csv('sample.csv', skiprows  = RowsToSkip)