カテゴリごとにRow数を調整（削除）する

Question

やりたい事

下のようなカテゴリ属性をもつ、データフレームがあるとして、(例のためpandasデータフレームで作成しています)

import pandas as pd
df= pd.DataFrame({
              '場所':['東京','大阪','京都','神奈川','東京','東京','京都','神奈川'],
              '他属性1':['a','a','a','a','a','a','a','a'],
              '他属性2':['a','a','a','a','a','a','a','a'],
              })

場所 他属性1 他属性2
0 東京 a a
1 大阪 a a
2 京都 a a
3 神奈川 a a
4 東京 a a
5 東京 a a
6 京都 a a
7 神奈川 a a

これらの件数を下のデータフレームを参考に各カテゴリごとに件数調整し抽出したいです。　

df = pd.DataFrame({
              '場所':['東京','大阪','京都','神奈川'],
              '抽出件数':[2,1,1,1]
              })

場所 抽出件数
0 東京 2
1 大阪 1
2 京都 1
3 神奈川 1

＜目標イメージ＞

場所 他属性1 他属性2
0 東京 a a
1 大阪 a a
3 神奈川 a a
4 東京 a a
6 京都 a a

順番は問わず（各カテゴリごとにシャフルしたい）

課題

pysparkを使っているのですが、対象データにおけるカテゴリ数が多く、以下のような一つずつするのは避けたいです。

df_agg = df.filter(F.col('場所')=='東京')
df_agg = df.filter(F.col('場所')=='大阪')
df_agg = df.filter(F.col('場所')=='神奈川')

Window関数？などを使えばそのような処理ができるのでしょうか。

Comments

[bsdfan]

質問としては、pysparkでやる方法を聞きたいのでしょうか？
pandasでもいいのでしょうか？

[sasakin]

修正依頼ありがとうございます。
できればpysparkの方がありがたいですが、pandasでキレイにできる方法などあれば（map関数？）そちら参考にしてpysparkで実装させていただきます！

Answer 1

シャッフルのために乱数の列を追加して、window内でソートして行番号と取得したい数と比較して抜き出しています。

import pandas as pd
import pyspark
from pyspark.sql import SparkSession
from pyspark.sql import Window
from pyspark.sql.functions import col, rand, row_number

spark = SparkSession.builder.getOrCreate()

df = pd.DataFrame({
              '場所':['東京','大阪','京都','神奈川','東京','東京','京都','神奈川'],
              '他属性1':['a','b','c','d','e','f','g','h'],
              '他属性2':['a','a','a','a','a','a','a','a'],
              })

df_num = pd.DataFrame({
              '場所':['東京','大阪','京都','神奈川'],
              '抽出件数':[2,1,1,1]
              })

sdf = spark.createDataFrame(df)
sdf_num = spark.createDataFrame(df_num)

result = (sdf
    .join(sdf_num, ['場所'], 'left')
    .withColumn('x', rand())
    .withColumn('rank', row_number().over(Window.partitionBy(col('場所')).orderBy(col('x'))))
    .where(col('rank') <= col('抽出件数'))
    .drop('抽出件数', 'x', 'rank'))

result.show()
#+------+-------+-------+
#|  場所|他属性1|他属性2|
#+------+-------+-------+
#|  大阪|      b|      a|
#|神奈川|      h|      a|
#|  東京|      f|      a|
#|  東京|      e|      a|
#|  京都|      c|      a|
#+------+-------+-------+

Answer 2

以下のコードで実現できました。
多分関数でもっとシンプルに書ける気がしますが、悪くはない気がします。
※シャフルは入っていません

import pandas as pd
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from functools import reduce
spark = SparkSession.builder.getOrCreate()
df= pd.DataFrame({
              '場所':['東京','大阪','京都','神奈川','東京','東京','京都','神奈川'],
              '他属性1':['a','a','a','a','a','a','a','a'],
              '他属性2':['a','a','a','a','a','a','a','a'],
              })
main_df = spark.createDataFrame(df)
main_df.show()

df = pd.DataFrame({
              '場所':['東京','大阪','京都','神奈川'],
              '抽出件数':[2,1,1,1]
              })
count_df = spark.createDataFrame(df)
count_df.show()

filter_df_list = [main_df.filter(F.col('場所') == row.場所).limit(row.抽出件数) for row in count_df.collect()]
reduce(lambda x,y: x.union(y), filter_df_list).show()