Re: Partition pruning with array-contains check and current_setting function

[Marcelo Zabani <mzabani@xxxxxxxxx>]

I managed to get a plan I was hoping for, but it still doesn't prune partitions. I created a new operator #|<(integer[], integer) that is defined in SQL and is basically equivalent to value=ANY(array), and a non-stable tenants() function defined that returns an array from the setting, and with that I could use a scalar subquery without running into type-checking errors. This gives me an InitPlan node:

=> SET my.tenant_id='{1}';EXPLAIN (COSTS OFF) SELECT COUNT(*) FROM tbl WHERE tenant_id #|< (select tenants());
SET
                        QUERY PLAN
----------------------------------------------------------
 Finalize Aggregate
   InitPlan 1 (returns $0)
     ->  Result
   ->  Gather
         Workers Planned: 2
         Params Evaluated: $0
         ->  Partial Aggregate
               ->  Parallel Append
                     ->  Parallel Seq Scan on tbl2 tbl_2
                           Filter: (tenant_id = ANY ($0))
                     ->  Parallel Seq Scan on tbl1 tbl_1
                           Filter: (tenant_id = ANY ($0))

It still doesn't prune even if I EXPLAIN ANALYZE it. I thought maybe I did something wrong with the operator definition, so I tried making tenants() immutable and removing the scalar subquery, and then it does prune:

=> SET my.tenant_id='{1}';EXPLAIN (COSTS OFF) SELECT COUNT(*) FROM tbl WHERE tenant_id #|< tenants();
SET
                      QUERY PLAN
------------------------------------------------------
 Aggregate
   ->  Seq Scan on tbl1 tbl
         Filter: (tenant_id = ANY ('{1}'::integer[]))

Sadly I can't make tenants() immutable because it's a runtime setting, and making tenants() STABLE does not lead to partition pruning with or without the scalar subquery around it.

I'm a bit lost. It seems like postgres is fully capable of pruning partitions for ="" checks, and some strange detail is confusing it in this case. I'm not sure what else to try.

On Wed, Aug 7, 2024 at 6:10 PM Marcelo Zabani <mzabani@xxxxxxxxx> wrote:

Hello all. I am trying to make postgres 16 prune partition for queries with `WHERE tenant_id=ANY(current_setting('my.tenant_id')::integer[])`, but I haven't been able to make it work, and naturally it impacts performance so I thought this list would be appropriate.

Here's the SQL I tried (but feel free to skip to the end as I'm sure all this stuff is obvious to you!):

CREATE TABLE tbl (id SERIAL NOT NULL, tenant_id INT NOT NULL, some_col INT, PRIMARY KEY (tenant_id, id))
PARTITION BY HASH (tenant_id);

CREATE TABLE tbl1 PARTITION OF tbl FOR VALUES WITH (MODULUS 2, REMAINDER 0);
CREATE TABLE tbl2 PARTITION OF tbl FOR VALUES WITH (MODULUS 2, REMAINDER 1);

INSERT INTO tbl (tenant_id, some_col) SELECT 1, * FROM generate_series(1,10000);
INSERT INTO tbl (tenant_id, some_col) SELECT 3, * FROM generate_series(1,10000);

Partition pruning works as expected for this query (still not an array-contains check):
EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=1;

When reading from a setting it also prunes partitions correctly:
SET my.tenant_id=1;
EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=current_setting('my.tenant_id')::integer;

It still does partition pruning if we use a scalar subquery. I can see the (never executed) scans in the plan.
EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=(SELECT current_setting('my.tenant_id')::integer);

But how about an array-contains check? Still prunes, which is nice.
EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=ANY('{1}'::integer[]);

However, it doesn't prune if the array is in a setting:
SET my.tenant_id='{1}';
EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=ANY(current_setting('my.tenant_id')::integer[]);

I actually expected that when in a setting, none of the previous queries would've done partition pruning because I thought `current_setting` is not a stable function. But some of them did, which surprised me.

So I thought maybe if I put it in a scalar query it will give me an InitPlan node, but it looks like method resolution for ="" won't let me try this:

EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=ANY((SELECT current_setting('my.tenant_id')::integer[]));

ERROR:  operator does not exist: integer = integer[]

I tried using UNNEST, but that adds a Hash Semi Join to the plan which also doesn't do partition pruning.
EXPLAIN ANALYZE SELECT COUNT(*) FROM tbl WHERE tenant_id=ANY((SELECT UNNEST(current_setting('my.tenant_id')::integer[])));

My question is if there's a way to do partition pruning based on array-contains operator if the array is in a setting. The use-case is to make Row Level Security policies do partition pruning "automatically" in a setting where users can be in more than one tenant.

It feels like this would work if there were a non-overloaded operator that takes in an array and a single element and tests for array-contains, because then I could use that operator with a scalar subquery and get an InitPlan node. But I'm new to all of this, so apologies if I'm getting it all wrong!

Thanks in advance,
Marcelo.