创建DataFrame两列时触发too many values to unpack错误的修复方案

输入DataFrame #

Id   Status
0  Id001   online
1  Id002  running
2  Id002      off
3  Id003   online
4  Id003    valid
5  Id003  running
6  Id004      off
7  Id004      off

期望输出DataFrame #

Id   Status   Type  Values
0  Id001   online  green  yellow
1  Id002  running    NaN     NaN
2  Id002      off    red   white
3  Id003   online  green  yellow
4  Id003    valid    NaN     NaN
5  Id003  running    NaN     NaN
6  Id004      off    red   white
7  Id004      off    red   white

尝试的代码 #

import pandas as pd

df = pd.DataFrame({'Id': ['Id001', 'Id002', 'Id002', 'Id003', 'Id003', 'Id003', 'Id004', 'Id004'],
                   'Status': ['online', 'running', 'off', 'online', 'valid', 'running', 'off', 'off']})

def test_func(df):
    if df['Status']=='online':
        return ['green', 'yellow']
    elif df['Status']=='off':
        return ['red', 'white']
    
df['Type'], df['Values'] = df.apply(test_func, axis=1)

报错信息 #

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [70], in <cell line: 7>()
      4     elif df['Status']=='off':
      5         return ['red', 'white']
----> 7 df['Type'], df['Values'] = df.apply(test_func, axis=1)

ValueError: too many values to unpack (expected 2)

解决方案 #

错误根源是：当Status为running或valid时，test_func无返回值，默认返回None。apply返回的Series中混合了列表和None，无法直接拆分成两列赋值。

提供两种修复方式：

方式一：让函数始终返回长度为2的列表 #

导入numpy，让函数在无匹配值时返回[np.nan, np.nan]，再通过result_type='expand'将结果转为DataFrame赋值：

import pandas as pd
import numpy as np

df = pd.DataFrame({'Id': ['Id001', 'Id002', 'Id002', 'Id003', 'Id003', 'Id003', 'Id004', 'Id004'],
                   'Status': ['online', 'running', 'off', 'online', 'valid', 'running', 'off', 'off']})

def test_func(row):
    if row['Status'] == 'online':
        return ['green', 'yellow']
    elif row['Status'] == 'off':
        return ['red', 'white']
    else:
        return [np.nan, np.nan]

result = df.apply(test_func, axis=1, result_type='expand')
df[['Type', 'Values']] = result

方式二：让函数返回指定索引的Series #

修改函数返回带列名索引的Series，apply会自动匹配列名生成对应列：

import pandas as pd

df = pd.DataFrame({'Id': ['Id001', 'Id002', 'Id002', 'Id003', 'Id003', 'Id003', 'Id004', 'Id004'],
                   'Status': ['online', 'running', 'off', 'online', 'valid', 'running', 'off', 'off']})

def test_func(row):
    if row['Status'] == 'online':
        return pd.Series(['green', 'yellow'], index=['Type', 'Values'])
    elif row['Status'] == 'off':
        return pd.Series(['red', 'white'], index=['Type', 'Values'])
    else:
        return pd.Series([pd.NA, pd.NA], index=['Type', 'Values'])

df[['Type', 'Values']] = df.apply(test_func, axis=1)

两种方式都能解决ValueError，生成符合预期的结果。

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