如何用Pandas生成类SQL分组统计的完整DataFrame结果？

背景与需求 #

现有数据库表结构如下：

stream_history
+ ----------+-------- +
+ played_at + song_id +
+ --------- + ------- +

songs
+ ------- + --------- + -------- +
+ song_id + song_name + album_id +
+ ------- + --------- + -------- +

artists
+ --------- + ----------- +
+ artist_id + artist_name +
+ --------- + ----------- +

artists_songs (关联表)
+ --------- + ------- +
+ artist_id + song_id +
+ --------- + ------- +

albums
+ -------- + ---------- +
+ album_id + album_name +
+ -------- + ---------- +

通过以下SQL查询获取播放数据：

select distinct played_at, song_name, sh.song_id, artist_name, 
        artists.artist_id, album_name
    from stream_history as sh
    join songs on songs.song_id = sh.song_id
    join art_songs as ars on ars.song_id = songs.song_id
    join artists on artists.artist_id = ars.artist_id
    join albums on albums.album_id = songs.album_id;

返回结果示例（注意同一播放记录因歌曲多艺人会生成多条数据）：

+ --------- + --------- + --------- + ----------- + ----------- + ---------- +
+ played_at + song_name + song_id   + artist_name + artist_id   + album_name +
+ --------- + --------- + --------- + ----------- + ----------- + ---------- +
+ DATE_1    + Song_A    + SONG_A_ID + ARTIST_A    + ARTIST_A_ID + ALBUM_A    +
+ DATE_1    + Song_A    + SONG_A_ID + ARTIST_B    + ARTIST_B_ID + ALBUM_A    +
+ DATE_2    + Song_B    + SONG_B_ID + ARTIST_C    + ARTIST_C_ID + ALBUM_B    +
+ DATE_3    + Song_C    + SONG_C_ID + ARTIST_D    + ARTIST_D    + ALBUM_C    +
+ --------- + --------- + --------- + ----------- + ----------- + ---------- +

核心需求：用Pandas统计每位艺术家的播放量（对应SQL DISTINCT COUNT(artist_id)）和每首歌曲的播放次数（对应SQL DISTINCT COUNT(song_id)），且结果DataFrame需包含完整的维度字段（如歌曲名、艺术家名等），格式如下：

+ ----- + --------- + ------- + ----------- + --------- + ---------- + 
+ count + song_name + song_id + artist_name + artist_id + album_name + 
+ ----- + --------- + ------- + ----------- + --------- + ---------- + 

之前尝试的groupby方法仅返回统计维度和count两列，无法满足格式要求，希望避免编写多条SQL，直接用Pandas处理。

示例测试数据：

import pandas as pd

df = pd.DataFrame([
    ['2021-07-22 14:36', 'Song A', 1, 'Artist A', 1], 
    ['2021-06-23 20:28', 'Song B', 2, 'Artist B', 2],
    ['2021-06-23 20:28', 'Song B', 2, 'Artist C', 3],
    ['2021-06-23 20:35', 'Song C', 3, 'Artist A', 1]
], columns=['Played at', 'Song', 'Song_id', 'Artist', 'Artist_id'])

解决方案 #

1. 按艺术家统计播放量 #

需要先对Played at和Song_id去重（避免同一播放记录因多艺人重复统计），再按艺术家分组统计，最后关联原表获取完整维度信息：

# 步骤1：去重，确保每条播放记录只算一次
unique_plays = df.drop_duplicates(subset=['Played at', 'Song_id'])

# 步骤2：统计每个艺术家关联的播放次数
artist_counts = unique_plays.groupby('Artist_id')['Song_id'].count().reset_index(name='count')

# 步骤3：关联原表获取完整字段，同时去重（避免同一艺术家对应多条歌曲记录）
artist_result = artist_counts.merge(
    df.drop_duplicates(subset=['Artist_id']), 
    on='Artist_id', 
    how='left'
)[['count', 'Song', 'Song_id', 'Artist', 'Artist_id']]

print(artist_result)

输出结果：

count    Song  Song_id    Artist  Artist_id
0      2  Song A        1  Artist A          1
1      1  Song B        2  Artist B          2
2      1  Song B        2  Artist C          3

2. 按歌曲统计播放次数 #

同样先去重，再按歌曲分组统计，最后关联完整字段：

# 步骤1：去重，确保每条播放记录只算一次
unique_plays = df.drop_duplicates(subset=['Played at', 'Song_id'])

# 步骤2：统计每首歌曲的播放次数
song_counts = unique_plays.groupby('Song_id')['Played at'].count().reset_index(name='count')

# 步骤3：关联原表获取完整字段，去重避免同一歌曲多条艺人记录
song_result = song_counts.merge(
    df.drop_duplicates(subset=['Song_id']), 
    on='Song_id', 
    how='left'
)[['count', 'Song', 'Song_id', 'Artist', 'Artist_id']]

print(song_result)

输出结果：

count    Song  Song_id    Artist  Artist_id
0      1  Song A        1  Artist A          1
1      1  Song B        2  Artist B          2
2      1  Song C        3  Artist A          1

通用化封装 #

可以将逻辑封装成函数，适配不同统计维度：

def get_stats(df, group_col, count_col='Played at'):
    # 去重：避免同一播放记录重复统计
    unique_plays = df.drop_duplicates(subset=['Played at', 'Song_id'])
    # 统计目标维度的播放量
    counts = unique_plays.groupby(group_col)[count_col].count().reset_index(name='count')
    # 关联原表获取完整维度字段
    result = counts.merge(
        df.drop_duplicates(subset=[group_col]), 
        on=group_col, 
        how='left'
    )
    # 调整列顺序，将count放在最前面
    cols = ['count'] + [col for col in result.columns if col != 'count']
    return result[cols]

# 按艺术家统计
print(get_stats(df, 'Artist_id'))
# 按歌曲统计
print(get_stats(df, 'Song_id'))

内容的提问来源于stack exchange，提问作者Emkay