如何在Snowflake中调整列值并保持列总和不变？

问题背景 #

在Snowflake中需要对PRICE列生成新列，要求：

• 新列保持原PRICE的大小顺序（相同PRICE的新值一致，最小值对应新列最小值，依此类推）
• 新列的总和与原PRICE列的总和完全一致

之前尝试使用WIDTH_BUCKET函数，但该函数仅能生成等宽分桶的编号，无法满足总和不变的要求。

原代码及输出 #

原执行代码：

SELECT 
    sale_date, 
    price,
    WIDTH_BUCKET(price, 200000, 600000, 5) AS "SALES GROUP"
  FROM home_sales
  ORDER BY sale_date;

原输出：

+------------+-----------+-------------+
| SALE_DATE  |     PRICE | SALES GROUP |
|------------+-----------+-------------|
| 2013-08-01 | 10        |           1 |
| 2014-02-01 | 20        |           2 |
| 2015-04-01 | 30        |           3 |
| 2016-04-01 | 10        |           1 |
| 2017-04-01 | 50        |           4 |
| 2018-04-01 | 60        |           5 |
+------------+-----------+-------------+

期望输出（核心要求） #

新列总和与原PRICE列总和180一致，且保持原列顺序：

+------------+-----------+-------------+
| SALE_DATE  |     PRICE | SALES GROUP |
|------------+-----------+-------------|
| 2013-08-01 | 10        |          12 |
| 2014-02-01 | 20        |          24 |
| 2015-04-01 | 30        |          28 |
| 2016-04-01 | 10        |          12 |
| 2017-04-01 | 50        |          47 |
| 2018-04-01 | 60        |          57 |
+------------+-----------+-------------+

解决方案思路 #

要实现需求，需通过以下关键步骤：

1. 对PRICE列进行稠密排序，保证相同价格的行获得相同排名（维持顺序一致性）
2. 基于排名比例计算缩放后的基础值，确保总和与原列一致
3. 处理小数余数，将额外金额分配给高排名的组，保证最终值符合顺序且总和不变

完整SQL代码 #

WITH ranked_data AS (
    -- 生成PRICE的稠密排名，同时计算原列总金额
    SELECT 
        sale_date,
        price,
        DENSE_RANK() OVER (ORDER BY price) AS price_rank,
        SUM(price) OVER () AS total_price
    FROM home_sales
),
rank_stats AS (
    -- 统计每个排名的行数、所有排名的总和
    SELECT 
        price_rank,
        COUNT(*) AS rank_count,
        SUM(price_rank) OVER () AS total_rank,
        total_price
    FROM ranked_data
    GROUP BY price_rank, total_price
),
rank_scaled AS (
    -- 计算每个排名的总分配金额，拆分整数和小数部分
    SELECT 
        *,
        (total_price * price_rank) / total_rank AS rank_total_scaled,
        FLOOR((total_price * price_rank) / total_rank) AS rank_total_floor,
        (total_price * price_rank) / total_rank - FLOOR((total_price * price_rank) / total_rank) AS rank_frac_total
    FROM rank_stats
),
total_remainder AS (
    -- 计算需要额外分配的整数余数总和
    SELECT SUM(rank_frac_total) AS remainder FROM rank_scaled
),
rank_adjusted AS (
    -- 按排名从高到低分配余数，保证高排名值更高
    SELECT 
        rs.*,
        rs.rank_total_floor + CASE WHEN ROW_NUMBER() OVER (ORDER BY rs.price_rank DESC) <= tr.remainder THEN 1 ELSE 0 END AS rank_total_adjusted
    FROM rank_scaled rs, total_remainder tr
),
final_data AS (
    -- 将排名总金额分摊到每一行，保证同排名值一致
    SELECT 
        rd.sale_date,
        rd.price,
        ra.rank_total_adjusted / ra.rank_count AS adjusted_price
    FROM ranked_data rd
    JOIN rank_adjusted ra ON rd.price_rank = ra.price_rank
)
-- 输出结果，按需转为整数
SELECT 
    sale_date,
    price,
    CAST(adjusted_price AS INTEGER) AS "SALES GROUP"
FROM final_data
ORDER BY sale_date;

代码说明 #

1. ranked_data：使用DENSE_RANK()确保相同价格的行排名相同，窗口函数SUM(price) OVER ()计算原列总金额。
2. rank_stats：按排名分组，统计每个排名的行数和所有排名的总和，用于后续比例计算。
3. rank_scaled：根据排名占比计算每个排名应分配的总金额，拆分出整数和小数部分。
4. total_remainder：统计所有小数部分的总和，这部分需要转化为整数额外分配。
5. rank_adjusted：将余数按排名从高到低分配，保证高排名的行值始终高于低排名，维持顺序。
6. final_data：将每个排名的总调整金额分摊到该行，确保相同价格的行新值一致。

如果不需要整数结果，可去掉CAST(adjusted_price AS INTEGER)保留小数形式。

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