能否优化PostgreSQL中Bitstring的字典序后继函数？

背景说明 #

PostgreSQL中对bitstring执行ORDER BY时遵循字典序。例如，1到4位的所有bitstring排序结果如下：

0
 00
 000
 001
 01
 010
 011
 1
 10
 100
 101
 11
 110
 111

我们需要实现一个最多k位的字典序后继函数succ_lex(x,k)，返回字典序的下一个值，典型示例：

• succ_lex('0',3)=b'00'
• succ_lex('01',3)=b'010'
• succ_lex('010',3)=b'011'

现有朴素实现（性能瓶颈） #

当前有一个基于递归的朴素实现，但递归调用会带来额外性能开销，处理大量数据时效率较低：

CREATE FUNCTION vbit_cutlast1s(p_x varbit) RETURNS varbit AS $f$
  SELECT CASE
    WHEN len>0 AND get_bit(p_x,len-1)::boolean THEN vbit_cutlast1s( substring(p_x,1,len-1) )
    ELSE CASE WHEN len=0 THEN ''::varbit ELSE substring(p_x,1,len-1) || '1'::bit(1) END
    END
  FROM (SELECT length(p_x)) t(len)
$f$ LANGUAGE SQL IMMUTABLE;
COMMENT ON FUNCTION vbit_cutlast1s(varbit)
  IS '截断末尾连续的1，将剩余部分的最后一个0（或空值）替换为1。'
;
CREATE FUNCTION vbit_succ_lex(p_x varbit, p_bits int DEFAULT 64) RETURNS varbit AS $f$
  SELECT CASE
    WHEN length(p_x)<p_bits THEN p_x || '0'::bit(1)
    ELSE vbit_cutlast1s(p_x)
    END
$f$ LANGUAGE SQL IMMUTABLE;
COMMENT ON FUNCTION vbit_succ_lex(varbit,int)
  IS '基于cutlast1s()的、受k位限制的字典序后继函数。'
;

测试代码 #

以下测试用例用于验证函数逻辑正确性：

DO $tests$
begin
  ASSERT vbit_succ_lex('',3) = b'0',      'T1';
  ASSERT vbit_succ_lex('0',3) = b'00',    'T2';
  ASSERT vbit_succ_lex('01',3) = b'010',  'T3';
  ASSERT vbit_succ_lex('010',3) = b'011', 'T4';
  ASSERT vbit_succ_lex('00111',5) = b'01','T5';
  ASSERT vbit_succ_lex('01111',5) = b'1', 'T6';
  ASSERT vbit_succ_lex('111',3) = b'',    'T7';
end;
$tests$ LANGUAGE plpgsql;

优化后的快速实现 #

通过非递归方式，利用字符串定位和直接位操作替代递归，大幅提升性能：

CREATE FUNCTION succ_lex_fast(p_x varbit, p_bits int DEFAULT 64) RETURNS varbit AS $f$
DECLARE
  len int := length(p_x);
  pos int;
BEGIN
  -- 若当前bitstring长度小于k，直接追加0得到后继
  IF len < p_bits THEN
    RETURN p_x || '0'::bit(1);
  END IF;
  
  -- 反转字符串后找到第一个0的位置（从左数，对应原字符串从右数第一个0）
  pos := position('0' IN reverse(p_x));
  
  -- 若全为1，返回空串（已到字典序末尾）
  IF pos = 0 THEN
    RETURN ''::varbit;
  END IF;
  
  -- 截断到第一个0的前一位，将该0改为1，得到后继值
  RETURN substring(p_x, 1, len - pos) || '1'::bit(1);
END;
$f$ LANGUAGE plpgsql IMMUTABLE;

COMMENT ON FUNCTION succ_lex_fast(varbit, int)
  IS '非递归实现的、受k位限制的bitstring字典序后继函数，性能更优。';

验证优化实现 #

替换测试用例中的函数名，验证优化后实现的正确性：

DO $tests_fast$
begin
  ASSERT succ_lex_fast('',3) = b'0',      'T1';
  ASSERT succ_lex_fast('0',3) = b'00',    'T2';
  ASSERT succ_lex_fast('01',3) = b'010',  'T3';
  ASSERT succ_lex_fast('010',3) = b'011', 'T4';
  ASSERT succ_lex_fast('00111',5) = b'01','T5';
  ASSERT succ_lex_fast('01111',5) = b'1', 'T6';
  ASSERT succ_lex_fast('111',3) = b'',    'T7';
end;
$tests_fast$ LANGUAGE plpgsql;

优化说明 #

• 移除递归调用，避免了函数栈的反复创建与销毁开销
• 使用reverse+position直接定位目标位置，逻辑更直观，执行步骤更少
• 适合处理长bitstring场景，在批量处理时性能提升明显

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