选择并在Postgres数据库条件计算

Question

我有一个名为rain_tanks一个Postgres表更新，如下图所示：

id hour  rain demand current_volume unmet_demand 
1 00:00 4.0  2.0  2    0.0 
2 00:10 3.0  4.0  [null]   [null] 
3 00:20 1.0  6.0  [null]   [null] 
4 00:30 7.0  3.0  [null]   [null] 

我想做到这一点的计算和更新current_volume和unmet_demand列（此代码只是为了。显示需要做什么，我想做到这一点，而无需使用Python中的函数或代码行）：

a = lag(current_volume) + rain - demand 
if a < 0: 
    unmet_demand = current_volume  
    current_volume = 0 
else: 
    unmet_demand = 0 
    current_volume = a 

预计表：

id hour  rain demand current_volume unmet_demand 
1 00:00 4.0  2.0  2    0 
2 00:10 3.0  4.0  1    0 
3 00:20 1.0  6.0  0    -4 
4 00:30 7.0  3.0  4    0 

我想我需要的是先选择和更新列。我尝试了以下的选择，但它不工作：

import psycopg2 as p 
conn = p.connect("dbname = 'test' user = 'postgres' password = 'pass' host = 'localhost'") 
cur = conn.cursor() 

cur.execute("""SELECT Id,rain,demand,current_volume,unmet_demand, 
         CASE WHEN (rain - demand + lag(current_volume) over(
           order by Id)) >= 0 
          THEN (rain - demand + lag(current_volume) over(
           order by Id)) ELSE 0 END 
       FROM rain_tanks ORDER BY Id""") 

任何帮助将非常感激。

编辑（添加与性能相关的问题）：我决定在postgres数据库中执行这些计算的原因是为了查看在Python中使用numpy数组是否有速度改进。我有大约1000万点的降雨和需求列，这里提供的答案需要花费比雨天和需求数量庞大的等效python函数更长的时间。是否仍有空间来提高查询的性能？

Answer 1

您应该使用递归cte来获得所需的结果。由于每行都取决于以前的计算，因此无法通过lag获得。

WITH RECURSIVE CTE(id,rain,demand,new_current_volume,new_unmet_demand) as 
(SELECT Id,rain,demand 
,case when rain-demand <0 then 0 else rain-demand end as new_current_volume 
,case when rain-demand <0 then rain-demand else 0 end as new_unmet_demand 
FROM rain_tanks 
WHERE id = 1 
UNION ALL 
SELECT r2.Id,r2.rain,r2.demand 
,case when r1.new_current_volume+r2.rain-r2.demand <0 then 0 else r1.new_current_volume+r2.rain-r2.demand end 
,case when r1.new_current_volume+r2.rain-r2.demand <0 then r1.new_current_volume+r2.rain-r2.demand else 0 end 
FROM cte r1 
JOIN rain_tanks r2 ON r2.id=r1.id+1 
) 
SELECT * FROM CTE; 

Sample Demo

编辑：

为update根据新计算出的值的表，包括2分null列从表中的列名，与计算出的那些一起。那么你可以在cte之后用这些计算update。

WITH RECURSIVE CTE(id,rain,demand,current_volume,unmet_demand,new_current_volume,new_unmet_demand) as 
(SELECT Id,rain,demand,null,null 
,case when rain-demand <0 then 0 else rain-demand end as new_current_volume 
,case when rain-demand <0 then rain-demand else 0 end as new_unmet_demand 
FROM rain_tanks 
WHERE id = 1 
UNION ALL 
SELECT r2.Id,r2.rain,r2.demand,null,null 
,case when r1.new_current_volume+r2.rain-r2.demand <0 then 0 else r1.new_current_volume+r2.rain-r2.demand end 
,case when r1.new_current_volume+r2.rain-r2.demand <0 then r1.new_current_volume+r2.rain-r2.demand else 0 end 
FROM cte r1 
JOIN rain_tanks r2 ON r2.id=r1.id+1 
) 
UPDATE rain_tanks r 
SET current_volume=c.new_current_volume 
    ,unmet_demand=c.new_unmet_demand 
FROM cte c 
WHERE r.id=c.id;