python - 在 Python 中复制 %transpose SAS 宏

Question

我的公司正在从 SAS 转向其他工具，Python 是首选工具。我有很多在 SAS 中运行的遗留代码需要移植过来，有些事情比其他事情更容易做。

我使用的许多数据集都需要转换，以便每个唯一 ID 列有一行，为此我广泛使用了此处的 %transpose 宏:https://raw.githubusercontent.com/art297/transpose/master/transpose.sas。从本质上讲，它允许我进行单个函数调用，我在其中传递输入数据集名称、输出数据集名称、id 列和要旋转的列以获得我想要的输出。

我一直在努力使用 Python 复制功能，我已经非常接近了，但我现在卡住了。

这是一个例子:

import pandas as pd

cust = ['CUST123', 'CUST123', 'CUST123', 'CUST456', 'CUST456']
start_date = ['01/01/2021','01/02/2021','01/03/2021','01/04/2021','01/05/2021']
end_date = ['01/11/2021','01/12/2021','01/13/2021','01/14/2021','01/15/2021']
code = ['ABC123','DEF456','GHI789','JKL123','MNO456']
df = pd.DataFrame(list(zip(cust, start_date, end_date, code )), columns = ['cust', 'start_date', 'end_date', 'code'])

df['start_date'] = pd.to_datetime(df['start_date'])
df['end_date'] = pd.to_datetime(df['end_date'])

这给了我一个看起来像这样的示例数据集:

      cust start_date   end_date    code
0  CUST123 2021-01-01 2021-01-11  ABC123
1  CUST123 2021-01-02 2021-01-12  DEF456
2  CUST123 2021-01-03 2021-01-13  GHI789
3  CUST456 2021-01-04 2021-01-14  JKL123
4  CUST456 2021-01-05 2021-01-15  MNO456

在 SAS 中我会使用:

%transpose(data=INPUT, 
   out=OUTPUT, 
   by=cust, 
   delimiter=_,
   var=start_date end_date code);

将其作为我想要的输出:

     cust start_date_1 end_date_1 code_1 start_date_2 end_date_2 code_2 start_date_3 end_date_3 code_3
0 CUST123   2021-01-01 2021-01-11 ABC123   2021-01-02 2021-01-12 DEF456   2021-01-03 2021-01-13 GHI789
1 CUST456   2021-01-04 2021-01-14 JKL123   2021-01-05 2021-01-15 MNO456

到目前为止，我已经使用了这段代码:

df['idx'] = df.groupby('cust').cumcount()+1
df = df.pivot_table(index='cust', columns='idx', values=['start_date', 'end_date', 'code'], aggfunc='first')
df = df.sort_index(axis=1, level=1)
df.columns = [f'{x}_{y}' for x, y in df.columns]
df.reset_index()

然而这会产生:

      cust  code_1 end_date_1 start_date_1  code_2 end_date_2 start_date_2  code_3 end_date_3 start_date_3
0  CUST123  ABC123 2021-01-11   2021-01-01  DEF456 2021-01-12   2021-01-02  GHI789 2021-01-13   2021-01-03
1  CUST456  JKL123 2021-01-14   2021-01-04  MNO456 2021-01-15   2021-01-05

这与我所追求的非常接近。唯一的问题是 SAS 宏按照我在函数调用中指定的顺序保留旋转列。无论我指定什么，Python 代码都会按字母顺序排列它们，因为它就是这样工作的。

我可以插入一个 for 循环，在我的数据透视字段前添加 [a...b...c...etc...]，然后对它们进行排序和旋转，然后遍历列标题以将它们更改为包含除了第一个角色之外的所有内容，但我无法想象这是最好的方法。这很丑陋，而且我几乎可以肯定还有一些我在所有搜索中都没有找到的其他方法更可取。

我是否坚持使用上述丑陋的方法，或者我是否正确认为有更好的方法？

Answer 1

在 sort_index 函数中传递 sort_remaining=False，以保持您的列不变。

result = (
    df.assign(idx=df.groupby("cust").cumcount() + 1)
    .pivot("cust", "idx")
    .sort_index(axis="columns", level=1, sort_remaining=False)
)

result.columns = result.columns.map(lambda x: f"{x[0]}_{x[1]}")