如何用Python与Polars实现订单数据的拆分/合并/逆透视转换？

核心思路 #

利用Polars的分组操作、**宽表转长表（unpivot）和透视（pivot）**能力，一次性完成分组内的数据重组，避免拆分后关联的繁琐操作，适配5万+组的大规模数据场景。

实现代码 #

import polars as pl

def transform_order_data(df: pl.DataFrame) -> pl.DataFrame:
    # 1. 标记分组与行类型：填充父订单ID，区分父/子订单
    df_grouped = df.with_columns(
        pl.col("Order ID").filter(pl.col("Parent Order ID").is_null()).forward_fill().alias("Parent Order ID"),
        pl.when(pl.col("Parent Order ID").is_null())
          .then("parent")
          .when(pl.col("Direction") == "Buy")
          .then("buy_child")
          .when(pl.col("Direction") == "Sell")
          .then("sell_child")
          .alias("row_type")
    )
    
    # 2. 提取父订单核心信息：每组仅保留父行的关键字段
    parent_info = df_grouped.filter(pl.col("row_type") == "parent").select(
        "Parent Order ID",
        pl.col("Direction").alias("Parent Direction"),
        "Price",
        "Some Value"
    )
    
    # 3. 重组供应商列：将分散的Name/Quote转成长表并配对
    provider_cols = [col for col in df.columns if "Provider" in col]
    name_cols = [col for col in provider_cols if "Name" in col]
    quote_cols = [col for col in provider_cols if "Quote" in col]
    
    df_providers = df_grouped.filter(pl.col("row_type").is_in(["buy_child", "sell_child"])).select(
        "Parent Order ID",
        "row_type",
        *[pl.col(name).alias(f"provider_{i+1}_name") for i, name in enumerate(name_cols)],
        *[pl.col(quote).alias(f"provider_{i+1}_quote") for i, quote in enumerate(quote_cols)]
    ).unpivot(
        index=["Parent Order ID", "row_type"],
        variable_name="provider_key",
        value_name="value"
    ).with_columns(
        pl.col("provider_key").str.extract(r"provider_(\d+)_").alias("provider_num"),
        pl.col("provider_key").str.extract(r"_(name|quote)$").alias("value_type")
    ).pivot(
        index=["Parent Order ID", "row_type", "provider_num"],
        columns="value_type",
        values="value"
    ).drop_nulls("name")
    
    # 4. 透视报价列：将买入/卖出报价转为字段，处理父订单方向反转逻辑
    df_quote_pivot = df_providers.pivot(
        index=["Parent Order ID", "provider_num", "name"],
        columns="row_type",
        values="quote"
    ).rename({"buy_child": "Quote Buy Raw", "sell_child": "Quote Sell Raw"})
    
    # 5. 关联父信息并整理最终格式
    final_df = df_quote_pivot.join(parent_info, on="Parent Order ID").with_columns(
        pl.when(pl.col("Parent Direction") == "Buy")
          .then(pl.col("Quote Buy Raw"))
          .otherwise(pl.col("Quote Sell Raw"))
          .alias("Quote Buy"),
        pl.when(pl.col("Parent Direction") == "Buy")
          .then(pl.col("Quote Sell Raw"))
          .otherwise(pl.col("Quote Buy Raw"))
          .alias("Quote Sell")
    ).select(
        pl.col("Parent Order ID").alias("Order ID"),
        "Parent Direction",
        "Price",
        "Some Value",
        pl.col("name").alias("Name Provider"),
        "Quote Buy",
        "Quote Sell"
    ).sort("Order ID", "provider_num")
    
    return final_df

# 测试单个父订单场景
df_original = pl.DataFrame(
    {
        'Order ID': ['A', 'foo', 'bar'],
        'Parent Order ID': [None, 'A', 'A'],
        'Direction': ["Buy", "Buy", "Sell"],
        'Price': [1.21003, None, 1.21003],
        'Some Value': [4, 4, 4],
        'Name Provider 1': ['P8', 'P8', 'P8'],
        'Quote Provider 1': [None, 1.1, 1.3],
        'Name Provider 2': ['P2', 'P2', 'P2'],
        'Quote Provider 2': [None, 1.15, 1.25],
        'Name Provider 3': ['P1', 'P1', 'P1'],
        'Quote Provider 3': [None, 1.0, 1.4],
        'Name Provider 4': ['P5', 'P5', 'P5'],
        'Quote Provider 4': [None, 1.0, 1.4]
    }
)
print("单个父订单转换结果：")
print(transform_order_data(df_original))

# 测试多父订单场景
df_original_two_orders = pl.DataFrame(
    {
        'Order ID': ['A', 'foo', 'bar', 'B', 'baz', 'rar'],
        'Parent Order ID': [None, 'A', 'A', None, 'B', 'B'],
        'Direction': ["Buy", "Buy", "Sell", "Sell", "Sell", "Buy"],
        'Price': [1.21003, None, 1.21003, 1.1384, None, 1.1384],
        'Some Value': [4, 4, 4, 42, 42, 42],
        'Name Provider 1': ['P8', 'P8', 'P8', 'P2', 'P2', 'P2'],
        'Quote Provider 1': [None, 1.1, 1.3, None, 1.10, 1.40],
        'Name Provider 2': ['P2', 'P2', 'P2', 'P1', 'P1', 'P1'],
        'Quote Provider 2': [None, 1.15, 1.25, None, 1.11, 1.39],
        'Name Provider 3': ['P1', 'P1', 'P1', 'P3', 'P3', 'P3'],
        'Quote Provider 3': [None, 1.0, 1.4, None, 1.05, 1.55],
        'Name Provider 4': ['P5', 'P5', 'P5', None, None, None],
        'Quote Provider 4': [None, 1.0, 1.4, None, None, None]
    }
)
print("\n多父订单转换结果：")
print(transform_order_data(df_original_two_orders))

关键步骤说明 #

• 分组标记：用forward_fill填充父订单ID，给每行标记类型（父/买入子/卖出子），确保每组数据关联正确。
• 提取父信息：单独提取每组父订单的核心字段（方向、价格等），避免重复处理。
• 供应商列重组：通过unpivot把分散的供应商Name/Quote列转成长表，再用pivot配对每个供应商的名称和报价，自动过滤空供应商。
• 报价方向调整：根据父订单的方向，动态切换买入/卖出报价的取值（比如父订单是Sell时，子订单的Buy报价对应最终的Sell列）。
• 最终整理：关联父信息，重命名列并排序，得到目标格式。

优势 #

• 全程使用Polars向量化操作，处理5万+组数据性能优异，比循环或拆分关联效率高。
• 自动适配最多15个供应商的场景，无需硬编码列名。
• 处理了父订单方向不同的特殊情况，符合真实业务逻辑。

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