PL/PGSQL返回查询中如何赋值变量复用距离计算结果？

问题说明 #

我编写了一个用于获取用户推荐好友的PostgreSQL函数，但代码无法运行——原因是我在WHERE子句中直接执行赋值操作isFriendNearTheUser := ST_DWithin(...)。我的核心需求是仅计算一次该距离判断结果，将其传入is_friend_recommendable函数复用，避免重复计算导致查询性能下降。请问这种需求能否在当前查询中实现？是否需要重写查询？

原函数代码：

create function get_recommended_friends(id VARCHAR)
    RETURNS TABLE(
        friend_id VARCHAR,
        is_friend_recommendable BOOLEAN
    )
    LANGUAGE plpgsql AS $$
DECLARE
    user user%rowtype;
BEGIN
    SELECT * INTO user from user where id = id;

    isFriendNearTheUser BOOLEAN;

    RETURN QUERY
        SELECT
            friend_id,
            is_friend_recommendable(player, subject.*, isFriendNearTheUser) AS is_friend_recommendable,
        FROM friend friend
        WHERE friend_id != id
          isFriendNearTheUser := ST_DWithin(user.current_point, friend.current_point, user.maximum_wanted_distance)
          ....
      .....
      AND ....
          AND...
END
$$;

解决方案 #

可以实现，无需完全重写函数，只需调整查询逻辑即可。以下是几种高效的实现方式：

方法1：用CTE预计算距离判断结果 #

通过公共表表达式（CTE）提前计算出距离判断结果，后续在SELECT和WHERE中直接复用，确保仅计算一次：

create function get_recommended_friends(id VARCHAR)
    RETURNS TABLE(
        friend_id VARCHAR,
        is_friend_recommendable BOOLEAN
    )
    LANGUAGE plpgsql AS $$
DECLARE
    current_user user%rowtype; -- 避免变量名与表名`user`冲突
BEGIN
    -- 修正原代码的恒真查询问题，正确匹配参数ID
    SELECT * INTO current_user from user where id = $1;

    RETURN QUERY
    WITH friend_distance_check AS (
        SELECT
            friend_id,
            ST_DWithin(current_user.current_point, friend.current_point, current_user.maximum_wanted_distance) AS is_friend_near
        FROM friend
        WHERE friend_id != $1
    )
    SELECT
        friend_id,
        is_friend_recommendable(player, fd.*, fd.is_friend_near) AS is_friend_recommendable
    FROM friend_distance_check fd
    WHERE 
        fd.is_friend_near -- 复用预计算的结果
        AND ... -- 你的其他过滤条件
END
$$;

方法2：用LATERAL子查询计算单次结果 #

如果需要关联其他表，LATERAL子查询是更灵活的选择，同样保证距离判断仅执行一次：

create function get_recommended_friends(id VARCHAR)
    RETURNS TABLE(
        friend_id VARCHAR,
        is_friend_recommendable BOOLEAN
    )
    LANGUAGE plpgsql AS $$
DECLARE
    current_user user%rowtype;
BEGIN
    SELECT * INTO current_user from user where id = $1;

    RETURN QUERY
    SELECT
        f.friend_id,
        is_friend_recommendable(player, f.*, distance_check.is_friend_near) AS is_friend_recommendable
    FROM friend f
    LATERAL (
        SELECT ST_DWithin(current_user.current_point, f.current_point, current_user.maximum_wanted_distance) AS is_friend_near
    ) AS distance_check
    WHERE 
        f.friend_id != $1
        AND distance_check.is_friend_near
        AND ... -- 其他过滤条件
END
$$;

方法3：plpgsql循环处理（适合复杂业务逻辑） #

如果过滤逻辑涉及较多分支判断，可以用游标循环逐行处理，先计算距离结果再复用：

create function get_recommended_friends(id VARCHAR)
    RETURNS TABLE(
        friend_id VARCHAR,
        is_friend_recommendable BOOLEAN
    )
    LANGUAGE plpgsql AS $$
DECLARE
    current_user user%rowtype;
    rec friend%rowtype;
    is_friend_near BOOLEAN;
BEGIN
    SELECT * INTO current_user from user where id = $1;

    FOR rec IN SELECT * FROM friend WHERE friend_id != $1 LOOP
        -- 仅计算一次距离判断
        is_friend_near := ST_DWithin(current_user.current_point, rec.current_point, current_user.maximum_wanted_distance);
        -- 加上你的其他过滤条件
        IF is_friend_near AND ... THEN
            RETURN NEXT (
                rec.friend_id,
                is_friend_recommendable(player, rec, is_friend_near)
            );
        END IF;
    END LOOP;
END
$$;

总结 #

以上三种方式都能满足单次计算、多次复用的需求。优先推荐CTE或LATERAL子查询的方式，它们基于集合查询，性能比循环更优，尤其适合数据量较大的场景。

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