custom_pyspark_groupby.py

from pyspark.sql.functions import col

def groupby(df, columns):
  sdf = df.to_spark()
  _groups = sdf.select(*columns).distinct().collect()
  _groups = [group.asDict() for group in _groups]
  groups = []
  for group in _groups:
    tmp = []
    for column in columns:
      tmp.append(group[column])
     groups.append(tmp)
   
  for group in groups:
    mask = (col(columns[0]) == group[0])
    for index, column in enunerate(columns):
      if index == 0:
        continue
      mask &= (col(column) == group[index])
    yield (group, sdf.filter(mask).to_pandas_on_spark())