Pythonic实现Pandas assign后自动删所用列（类R dplyr .keep="unused"）

在R语言的dplyr包中，使用mutate()函数（对应Pandas的assign()方法）创建新列时，可通过.keep="unused"参数自动删除生成新列所用的原始列。比如iris数据集的示例：

> head(iris)
  Sepal.Length Sepal.Width Petal.Length Petal.Width Species
1          5.1         3.5          1.4         0.2  setosa
2          4.9         3.0          1.4         0.2  setosa
3          4.7         3.2          1.3         0.2  setosa
4          4.6         3.1          1.5         0.2  setosa
5          5.0         3.6          1.4         0.2  setosa
6          5.4         3.9          1.7         0.4  setosa

# 自动删除创建new_col用到的列
> mutate(.data = head(iris), new_col = Sepal.Length + Petal.Length * Petal.Width, .keep = "unused")
  Sepal.Width Species new_col
1         3.5  setosa    5.38
2         3.0  setosa    5.18
3         3.2  setosa    4.96
4         3.1  setosa    4.90
5         3.6  setosa    5.28
6         3.9  setosa    6.08

但Pandas的assign()方法没有类似参数，目前常见实现方式是预先定义所用列列表，创建新列后手动删除，示例如下：

from sklearn import datasets
import pandas as pd

used_columns = ['sepal length (cm)', 'petal length (cm)', 'petal width (cm)']

iris = pd.DataFrame(datasets.load_iris().data, columns=datasets.load_iris().feature_names)

iris.assign(new_col = lambda x: x['sepal length (cm)'] + x['petal length (cm)'] * x['petal width (cm)']).drop(used_columns, axis=1)

这种方式需要单独维护列列表，认知负担较大，下面提供几种更Pythonic、易维护的实现方法，允许使用janitor等第三方包。

方案一：自定义封装函数，自动识别所用列 #

通过解析计算逻辑的源代码，自动提取用到的列，无需手动维护列列表：

from sklearn import datasets
import pandas as pd
import inspect
import ast

def assign_keep_unused(df, new_col_name, calc_func):
    # 解析lambda函数，提取用到的列名
    source_code = inspect.getsource(calc_func).strip()
    ast_tree = ast.parse(source_code)
    used_cols = []
    
    for node in ast.walk(ast_tree):
        # 匹配lambda中x['列名']的写法
        if (isinstance(node, ast.Subscript) 
            and isinstance(node.value, ast.Name) 
            and node.value.id == 'x'
            and isinstance(node.slice, ast.Constant)):
            used_cols.append(node.slice.value)
    
    used_cols = list(set(used_cols))
    # 创建新列并删除所用列
    return df.assign(**{new_col_name: calc_func}).drop(used_cols, axis=1)

# 测试使用
iris = pd.DataFrame(datasets.load_iris().data, columns=datasets.load_iris().feature_names)
result = assign_keep_unused(
    iris,
    'new_col',
    lambda x: x['sepal length (cm)'] + x['petal length (cm)'] * x['petal width (cm)']
)
print(result.head())

方案二：结合janitor包实现管道式操作 #

如果习惯使用janitor的管道风格，可以结合列筛选逻辑实现：

from sklearn import datasets
import pandas as pd
import janitor
import inspect
import ast

def janitor_assign_keep_unused(df, new_col_name, calc_func):
    # 复用列提取逻辑
    source_code = inspect.getsource(calc_func).strip()
    ast_tree = ast.parse(source_code)
    used_cols = []
    
    for node in ast.walk(ast_tree):
        if (isinstance(node, ast.Subscript) 
            and isinstance(node.value, ast.Name) 
            and node.value.id == 'x'
            and isinstance(node.slice, ast.Constant)):
            used_cols.append(node.slice.value)
    
    used_cols = list(set(used_cols))
    # 添加新列后筛选保留未用到的原始列+新列
    return (df
            .add_column(new_col_name, calc_func(df))
            .filter(items=[col for col in df.columns if col not in used_cols] + [new_col_name]))

# 测试使用
iris = pd.DataFrame(datasets.load_iris().data, columns=datasets.load_iris().feature_names)
result = janitor_assign_keep_unused(
    iris,
    'new_col',
    lambda x: x['sepal length (cm)'] + x['petal length (cm)'] * x['petal width (cm)']
)
print(result.head())

方案三：手动指定依赖列的简洁写法 #

如果不想用代码解析，也可以将所用列和计算逻辑绑定，减少单独维护列列表的负担：

from sklearn import datasets
import pandas as pd

def assign_drop_used(df, new_col_def):
    # new_col_def是元组：(新列名, 所用列列表, 计算函数)
    col_name, used_cols, calc_func = new_col_def
    return df.assign(**{col_name: calc_func}).drop(used_cols, axis=1)

# 测试使用
iris = pd.DataFrame(datasets.load_iris().data, columns=datasets.load_iris().feature_names)
result = assign_drop_used(
    iris,
    ('new_col',
     ['sepal length (cm)', 'petal length (cm)', 'petal width (cm)'],
     lambda x: x['sepal length (cm)'] + x['petal length (cm)'] * x['petal width (cm)'])
)
print(result.head())

以上几种方法都避免了单独维护列列表的繁琐，其中方案一和方案二的自动列识别方式更贴近dplyr的.keep="unused"逻辑，无需手动指定所用列；方案三则更直观，适合需要明确依赖列的场景。

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