基于尺寸排序的SKU组维度空值线性填充及可扩展性优化问询

需求说明 #

• 按skugroup分组，对每组内的length/width/height/volume字段的NaN值，用该字段最值生成的线性间距值填充
• 支持数值型（如shirt的1-5）和非数值型（如shoes的s/xl等）尺寸，且尺寸需按业务逻辑排序
• 适配未来新增SKU组（裤子、包、帽子等）的自定义尺寸规则

优化后代码 #

import pandas as pd
import numpy as np

data_wm = {
    'sku': [6124, 7343, 7981, 5761, 1570, 7223, 4107, 8187, 4653, 1802, 4079],
    'skugroup': ['shoes', 'shoes', 'shoes', 'shoes', 'shoes', 'shoes', 'shirt', 'shirt', 'shirt', 'shirt', 'shirt'],
    'size': ['s', 'm', 'l', 'xl', '2xl', '3xl', '1', '2', '3', '4', '5'],
    'length': [1.5, np.nan, np.nan, np.nan, np.nan, 6, 1, np.nan, np.nan, np.nan, 4],
    'width': [2, np.nan, np.nan, np.nan, np.nan, 8, 2, np.nan, np.nan, np.nan, 5],
    'height': [2, np.nan, np.nan, np.nan, np.nan, 3, 3, np.nan, np.nan, np.nan, 6],
    'volume': [6, np.nan, np.nan, np.nan, np.nan, 144, 6, np.nan, np.nan, np.nan, 120]
}

df = pd.DataFrame(data_wm)

# 定义各SKU组的尺寸排序规则，新增组直接在这里加即可
SIZE_ORDER_MAP = {
    'shoes': ['s', 'm', 'l', 'xl', '2xl', '3xl'],
    'shirt': ['1', '2', '3', '4', '5'],
    # 未来新增组示例：
    # 'pants': ['xs', 's', 'm', 'l', 'xl'],
    # 'bag': ['mini', 'small', 'medium', 'large']
}

def fill_nulls_based_on_size(df):
    # 遍历每个SKU组
    for skugroup, group_df in df.groupby('skugroup'):
        # 获取当前组的尺寸排序规则
        if skugroup in SIZE_ORDER_MAP:
            size_order = SIZE_ORDER_MAP[skugroup]
        else:
            # 无预设规则时，尝试按数值排序，非数值放最后
            size_order = sorted(group_df['size'].unique(), 
                               key=lambda x: int(x) if x.isdigit() else float('inf'))
        
        # 按尺寸规则对组内数据排序，确保线性填充和尺寸顺序匹配
        sorted_group = group_df.set_index('size').reindex(size_order).reset_index()
        
        # 遍历各维度字段，生成线性填充值并替换NaN
        for col in ['length', 'width', 'height', 'volume']:
            min_val = sorted_group[col].min()
            max_val = sorted_group[col].max()
            # 生成线性间距序列
            linear_vals = np.linspace(min_val, max_val, len(sorted_group), endpoint=True)
            # 填充空值
            sorted_group[col] = sorted_group[col].fillna(pd.Series(linear_vals, index=sorted_group.index))
        
        # 将填充后的数据更新回原DataFrame
        df.loc[sorted_group.index, :] = sorted_group

fill_nulls_based_on_size(df)
print(df)

关键优化点 #

• 可扩展的尺寸规则：用SIZE_ORDER_MAP字典统一管理各SKU组的尺寸排序逻辑，新增SKU组只需在字典中添加对应规则，无需修改核心逻辑
• 修复排序错位问题：先按尺寸规则对组内数据重新排序，再生成线性填充值，确保填充值和尺寸顺序严格对应，避免原代码中填充错位的问题
• 兼容数值/非数值尺寸：无预设规则的组自动尝试按数值排序，非数值尺寸会被放到序列末尾，适配各种尺寸类型
• 安全的最值计算：基于排序后的组数据计算维度最值，确保线性填充的基准值正确

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