我可以使用带有变量名的 Pandas Dataframe.assign(...) [重复]答案

【问题标题】：Can I use Pandas Dataframe.assign(...) with a variable name [duplicate]我可以使用带有变量名的 Pandas Dataframe.assign(...) [重复]
【发布时间】：2020-05-08 22:48:31
【问题描述】：

我编写了这段代码，以便对任何 Pandas DataFrame 进行分组，并快速获取分组大小和数据帧的示例行。

效果很好，但有一个问题：新列/索引“大小”的名称是固定，因为.assign( ... ) 命令不接受变量。因此，如果我的 DataFrame 有一个名为“Size”的列，它就会丢失。

我的计划是检查是否存在名为“Size”的列，如果存在，为索引使用不同的名称。我可以将assign 命令与字段名称的变量，而不是固定文本？

我想避免像多次重命名列这样的 hacky 解决方案。

import pandas as pd
try:
    from pandas.api.extensions import register_dataframe_accessor
except ImportError:
    raise ImportError('Pandas 0.24 or better needed')

@register_dataframe_accessor("cgrp")
class CustomGrouper:
    """Extra methods for dataframes."""

    def __init__(self, df):
        self._df = df

    def group_sample(self, by, subset=None):
        result = (self._df.groupby(by).apply(lambda x: x.sample(1).assign(Size = len(x)))).set_index('Size').sort_index(ascending=False)
        return result

我可以这样称呼

df.cgrp.group_sample(by=['column1', ... ])

并获得带有索引“Size”的结果

【问题讨论】：

是什么阻止您创建函数并在应用中使用它？
@G.Anderson - 是的，这就是我要找的，谢谢

标签： python pandas

【解决方案1】：

基本思想是使用字典解包。而不是在assign 函数中硬编码名称：

.assign(Size = len(x))

你可以使用字典解包来指定一个变量名：

.assign(**{col_name: len(x)})

我冒昧地修改了您的 group_sample 函数，具有 2 个功能：允许用户指定自定义名称并从默认列表中选择（如果他们不这样做）：

def group_sample(self, by, subset=None, col_name=None):
    _col_name = None

    if col_name is not None:
        # If a user specify a column name, use it
        # Raise error if the column already exists
        if col_name in self._df.columns:
            raise ValueError(f"Dataframe already has column '{col_name}'")
        else:
            _col_name = col_name
    else:
        # Choose from a list of default names
        _col_name = next((name for name in ['Size', 'Size_', 'Size__'] if name not in self._df.columns), None)

        if _col_name is None:
            raise ValueError('Cannot determine a default name for the size column. Please specify one manually')

    result = (self._df.groupby(by).apply(lambda x: x.sample(1).assign(**{_col_name: len(x)}))).set_index(_col_name).sort_index(ascending=False)
    return result

用法：

df1 = pd.DataFrame(np.random.randint(1, 5, (3, 2)), columns=['A','B'])
df1.cgrp.group_sample(by=['A'])     # the column name is Size

df2 = pd.DataFrame(np.random.randint(1, 5, (3, 2)), columns=['A','Size'])
df2.cgrp.group_sample(by=['A'])     # the column name is Size_

df3 = pd.DataFrame(np.random.randint(1, 5, (3, 2)), columns=['A','B'])
df3.cgrp.group_sample(by=['A'], col_name='B')  # error, B already exists

df4 = pd.DataFrame(np.random.randint(1, 5, (3, 2)), columns=['A','B'])
df4.cgrp.group_sample(by=['A'], col_name='MySize')  # custom column name

【讨论】：