在Python中复制%transpose宏

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]}")