[pgrx开发postgresql数据库扩展]5.自定义函数与SQL组合应用

news2024/12/28 6:10:04

老规矩的声明：

并不是所有场景都需要（或者适合）用rust来写的，绝大部分操作数据库的功能和计算，用SQL就已经足够了！

本系列中，所有的案例，仅用于说明pgrx的能力，而并非是说这样做比用SQL更合适。反之：对于操作数据库本身的部分，大部分能用SQL来实现的东西，都比做一个扩展开发要更加合适。

——如果哪位大神写Rust走火入魔，说啥数据库功能都要用Rust来扩展实现的，不报我的名字，打成半死就行，报我的名字，请打成八成死。

上篇文章做的函数基本上独立使用的，今天我们来看看写好的扩展函数与SQL联合起来怎么来实现一个比较复杂的功能：

例如有以下这样一个场景：

我们需要评估一下公司有多少人已经学完了公司所有课程，可以毕业输出社会了，这样就需要进行如下统计：假设截止到今年的12月30日：

年满35岁，而且在公司服务满10年的人有多少。
年满35岁，而且在公司服务超过5年，但是没有10年的人有多少。
年满35岁，在公司没有服务超过一年，但是没有5年的有多少？
年满35岁，入职没有超过一年的有多少？
不满35岁的，有多少？

反正一句话：

也就是说，需要把20万条记录，分成五类，然后统计数量，一般来说，这种简单的需求，用SQL就可以直接做了，如下所示：

还是那句话，我对SQL不是太熟，这写功能虽然可以实现，但是性能就不要要求了：

WITH graduate AS (
SELECT 
    CASE
    WHEN date_part('year',age('2023-12-30',birthday)) >= 35 
                    AND date_part('year',age('2023-12-30',indate)) >= 10
                    THEN  '35岁以上工作满10年'
    WHEN date_part('year',age('2023-12-30',birthday)) >= 35 
                    AND date_part('year',age('2023-12-30',indate)) >= 5
                    AND date_part('year',age('2023-12-30',indate)) < 10
        THEN  '35岁以上工作满5年不满10年'
    WHEN date_part('year',age('2023-12-30',birthday)) >= 35 
                    AND date_part('year',age('2023-12-30',indate)) >= 1
                    AND date_part('year',age('2023-12-30',indate)) < 5
        THEN  '35岁以上工作满1年不满5年'
    WHEN date_part('year',age('2023-12-30',birthday)) >= 35 
                    AND date_part('year',age('2023-12-30',indate)) < 1
        THEN  '35岁以上工作不满1年'
    
    WHEN date_part('year',age('2023-12-30',birthday)) < 35 
        THEN  '不满35岁'
    END as ga
FROM tab_emps)
SELECT ga,count(ga) FROM graduate
GROUP BY(ga);

计算结果如下：

可以看见，此类计算的输入和输出并不是1：1的，我们以前做数学计算，表中有20万条记录，计算出来，也是20万条记录，而我们这边在20万条记录的基础上，计算出来的是5条记录……

说那么多，不就是聚合么……

好吧，我们后面还会具体讲聚合函数的开发，这里我们用SRF的方式来实现这种操作。

代码如下：

//rust时间计算函数没有周年计算的能力，所以我自己写了一个 
//but 效率非常差……如果哪位大神知道如何高性能计算周年的话，请与我联系 
fn anniversary(now: &DateTime<Local>,date:&DateTime<Local>)-> i32{
    let month = now.month();  
    let day = now.day();  
    let month2 = date.month();  
    let day2 = date.day();  

    let mut year_a = now.year() - date.year();

    if month < month2{
        year_a -=1;
    }
    else if month == month2{
        if day < day2{
            year_a -=1;
        }
    }
    year_a
}

#[pg_extern]
fn cal_graduate(indate:pgrx::Date,birthday:pgrx::Date) -> &'static str{

    let now: DateTime<Local> = Local.with_ymd_and_hms(2023,12,30,23,59,59).unwrap();
    let dt1 = Local.timestamp_opt(indate.to_posix_time(), 0).unwrap();
    let dt2 = Local.timestamp_opt(birthday.to_posix_time(), 0).unwrap();
    
    let work_year = anniversary(&now,&dt1);
    let bir_year= anniversary(&now,&dt2);
    
    if bir_year >=35{
        if work_year >=10{
            "满35岁且工作满10年"
        }
        else if work_year >=5 && work_year< 10 {
            "满35岁且工作满5年但不满10年"
        }
        else if work_year >=1 && work_year < 5 {
            "满35岁且工作满1年但不满5年"
        }
        else{
            "满35岁且工作不满1年"
        }
    }
    else{
       "不满35岁"
    }
}

计算结果如下：