Create custom hash operator for Postgres partitioning

Question

I want to create a custom hash function that would be used by Postgres (version 13.2) to distribute rows across partitions. Problem is that with the current solution Postgres does not use partition pruning. Here is my code:

-- dummy hash function
CREATE OR REPLACE FUNCTION partition_custom_bigint_hash(value BIGINT, seed
BIGINT)
RETURNS BIGINT AS $$
    SELECT value;
$$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE;

-- operator
CREATE OPERATOR CLASS partition_custom_bigint_hash_op
    FOR TYPE int8
    USING hash AS
    OPERATOR 1 =,
    FUNCTION 2 partition_custom_bigint_hash(BIGINT, BIGINT);

-- table partitioned by hash with custom operator
CREATE TABLE sample(part_id BIGINT) PARTITION BY hash(part_id partition_custom_bigint_hash_op);
CREATE TABLE sample_part_1 PARTITION OF SAMPLE FOR VALUES WITH (modulus 3, remainder 0);
CREATE TABLE sample_part_2 PARTITION OF SAMPLE FOR VALUES WITH (modulus 3, remainder 1);
CREATE TABLE sample_part_3 PARTITION OF SAMPLE FOR VALUES WITH (modulus 3, remainder 2);

Now make sure that partition pruning is enabled and works correctly:

SHOW enable_partition_pruning;
--  enable_partition_pruning
-- --------------------------
--  on

EXPLAIN * FROM sample WHERE part_id = 1::BIGINT;
--                               QUERY PLAN                              
-- ----------------------------------------------------------------------
--  Seq Scan on sample_part_1 sample  (cost=0.00..38.25 rows=11 width=8)
--    Filter: (part_id = '1'::bigint)
-- (2 rows)

So it works fine when using condition part_id=1::BIGINT but if I skip casting to BIGINT I get:

EXPLAIN SELECT * FROM sample WHERE part_id = 1;
--                                   QUERY PLAN
-- ------------------------------------------------------------------------------
--  Append  (cost=0.00..101.36 rows=33 width=8)
--    ->  Seq Scan on sample_part_1 sample_1  (cost=0.00..33.73 rows=11 width=8)
--          Filter: (part_id = 1)
--    ->  Seq Scan on sample_part_2 sample_2  (cost=0.00..33.73 rows=11 width=8)
--          Filter: (part_id = 1)
--    ->  Seq Scan on sample_part_3 sample_3  (cost=0.00..33.73 rows=11 width=8)
--          Filter: (part_id = 1)

Question: What I need to change in order to make partition pruning work with both conditions part_id=1 and part_id=1::BIGINT?

Answer 1

There are several equality operators with bigint on the left side:

SELECT oid,
       oprcode::regproc AS function,
       oprright::regtype AS right_side
FROM pg_operator
WHERE oprname = '='
  AND oprleft = 'bigint'::regtype;

 oid  | function | right_side 
------+----------+------------
  410 | int8eq   | bigint
  416 | int84eq  | integer
 1868 | int82eq  | smallint
(3 rows)

Now the second query uses the second of these operators, but that operator does not belong to your custom operator family, so partition pruning does not take place.

See this comment in match_clause_to_partition_key in src/backend/partitioning/partprune.c:

/*
 * See if the operator is relevant to the partitioning opfamily.
 *
 * Normally we only care about operators that are listed as being part
 * of the partitioning operator family.  But there is one exception:
 * the not-equals operators are not listed in any operator family
 * whatsoever, but their negators (equality) are.  We can use one of
 * those if we find it, but only for list partitioning.
 *
 * Note: we report NOMATCH on failure, in case a later partkey has the
 * same expression but different opfamily.  That's unlikely, but not
 * much more so than duplicate expressions with different collations.
 */

Create an operator family that contains the required operators:

CREATE FUNCTION partition_custom_hash(value int8, seed int8) RETURNS int8
   AS 'SELECT value' LANGUAGE SQL IMMUTABLE PARALLEL SAFE;

CREATE FUNCTION partition_custom_hash(value int4, seed int4) RETURNS int8
   AS 'SELECT value::int8' LANGUAGE SQL IMMUTABLE PARALLEL SAFE;

CREATE FUNCTION partition_custom_hash(value int2, seed int2) RETURNS int8 
   AS 'SELECT value::int8' LANGUAGE SQL IMMUTABLE PARALLEL SAFE;

CREATE OPERATOR FAMILY partition_custom_integer_hash_ops USING hash;

CREATE OPERATOR CLASS partition_custom_int8_hash_ops FOR TYPE int8 USING hash
   FAMILY partition_custom_integer_hash_ops AS
   OPERATOR 1 = (int8, int8),
   FUNCTION 2 partition_custom_hash(int8, int8);

CREATE OPERATOR CLASS partition_custom_int4_hash_ops FOR TYPE int4 USING hash
   FAMILY partition_custom_integer_hash_ops AS
   OPERATOR 1 = (int8, int4),
   FUNCTION 2 partition_custom_hash(int4, int4);

CREATE OPERATOR CLASS partition_custom_int2_hash_ops FOR TYPE int2 USING hash 
   FAMILY partition_custom_integer_hash_ops AS 
   OPERATOR 1 = (int8, int2),
   FUNCTION 2 partition_custom_hash(int2, int2);

Then if you use partition_custom_int8_hash_ops, it should work as you want.