基于DataFrame动态映射生成AcctId列的实现需求问询

看起来你需要根据Name列的不同取值，动态从对应的AcctId系列列里提取值来生成新的AcctId列，对吧？我给你分Pandas和PySpark两种常用的DataFrame工具场景来讲解实现方法，你可以根据自己的实际环境来选～

先看Pandas的实现方式 #

咱们先把你提供的样例DataFrame构造出来，然后一步步实现映射逻辑：

步骤1：构造样例DataFrame #

import pandas as pd

data = {
    "ID": [1,2,3,4,5,6],
    "Name": ["IdName", "IdLoc", "IdPop", "IdTop", "IdPin", "IdTrin"],
    "AcctIdName": [1,0,0,0,0,0],
    "AcctIdLoc": [0,-1,0,0,0,0],
    "AcctIdPop": [0,0,0,0,0,0],
    "AcctIdTop": [0,0,3,2,0,0],
    "AcctIdPin": [0,0,0,0,7,8]
}
df = pd.DataFrame(data)

步骤2：定义映射规则 #

先把Name的取值和对应的目标列名做成一个字典，方便后续调用：

name_to_col = {
    "IdName": "AcctIdName",
    "IdLoc": "AcctIdLoc",
    "IdPop": "AcctIdPop",
    "IdTop": "AcctIdTop",
    "IdPin": "AcctIdPin"
}

步骤3：实现映射逻辑 #

这里给你两种方法，按需选择：

方法一：逐行处理（适合小数据集）

用apply函数自定义一个提取值的逻辑，代码直观易懂：

def get_acct_id(row):
    # 先查映射字典，找不到就用默认的AcctIdName
    target_col = name_to_col.get(row["Name"], "AcctIdName")
    return row[target_col]

# 生成新的AcctId列
df["AcctId"] = df.apply(get_acct_id, axis=1)

方法二：批量向量处理（适合大数据集，效率更高）

用numpy的np.select来批量处理，比逐行apply快很多：

import numpy as np

# 构建条件列表和对应的取值列列表
conditions = [df["Name"] == name for name in name_to_col.keys()]
values = [df[col] for col in name_to_col.values()]
# 定义默认值：所有条件不满足时取AcctIdName
default_val = df["AcctIdName"]

df["AcctId"] = np.select(conditions, values, default=default_val)

运行完之后，你就能得到符合要求的AcctId列了，比如最后一行Name为IdTrin的行，AcctId会取AcctIdName的0值。

如果用的是PySpark怎么办？ #

要是你用的是分布式的PySpark DataFrame，也可以用条件判断链来实现：

步骤1：构造PySpark DataFrame #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when

spark = SparkSession.builder.appName("AcctIdMapper").getOrCreate()

data = [
    (1, "IdName", 1, 0, 0, 0, 0),
    (2, "IdLoc", 0, -1, 0, 0, 0),
    (3, "IdPop", 0, 0, 0, 3, 0),
    (4, "IdTop", 0, 0, 0, 2, 0),
    (5, "IdPin", 0, 0, 0, 0, 7),
    (6, "IdTrin", 0, 0, 0, 0, 8)
]
columns = ["ID", "Name", "AcctIdName", "AcctIdLoc", "AcctIdPop", "AcctIdTop", "AcctIdPin"]
df = spark.createDataFrame(data, columns)

步骤2：构建条件表达式生成AcctId列 #

# 先初始化表达式为默认值AcctIdName
acct_id_expr = col("AcctIdName")

# 遍历映射字典，逐个添加when条件
name_to_col = {
    "IdName": "AcctIdName",
    "IdLoc": "AcctIdLoc",
    "IdPop": "AcctIdPop",
    "IdTop": "AcctIdTop",
    "IdPin": "AcctIdPin"
}

for name, col_name in name_to_col.items():
    acct_id_expr = when(col("Name") == name, col(col_name)).otherwise(acct_id_expr)

# 添加新列
df = df.withColumn("AcctId", acct_id_expr)

# 查看结果
df.show()

这样就能在PySpark里得到符合要求的结果啦。

备注：内容来源于stack exchange，提问作者Maria