-- The SQL counts under the collation of the query they stand for (DESIGN.md
-- §9, "SQL surface").
--
-- lion_index_count(idx, key) is `SELECT count(*) ... WHERE col = key`, and
-- that query compares under the key's collation when the key brings one of
-- its own, and under the column's otherwise.  The functions used to look
-- every key up under the INDEX's collation, whatever the call said, which is
-- the same thing only while both collations are deterministic: a
-- case-insensitive index counted 'ABC' under 'abc' where the query under the
-- column's own collation does not.  A count under another collation than the
-- index's is now refused when either of the two is nondeterministic.
--
-- Nondeterministic collations need ICU; the second half is skipped without
-- it (test/expected/countcoll_1.out).
\set VERBOSITY terse
SET client_min_messages = warning;
LOAD 'pg_lion';
CREATE EXTENSION IF NOT EXISTS pg_lion;
RESET client_min_messages;
SET max_parallel_workers_per_gather = 0;
-- ---------- two deterministic collations agree on equality: counted ----------
CREATE TABLE lcc_d (id int, t text) WITH (autovacuum_enabled = off);
INSERT INTO lcc_d SELECT g, (ARRAY['abc', 'ABC', 'Abc', 'xyz'])[1 + g % 4]
  FROM generate_series(1, 400) g;
CREATE INDEX lcc_d_c ON lcc_d USING lion (t COLLATE "C");
CREATE INDEX lcc_d_t ON lcc_d USING lion (t);
SELECT lion_index_count('lcc_d_c', 'abc'::text) AS c_index,
	   lion_index_count('lcc_d_t', 'abc'::text COLLATE "C") AS c_key,
	   (SELECT count(*) FROM lcc_d WHERE t = 'abc') AS query;
 c_index | c_key | query 
---------+-------+-------
     100 |   100 |   100
(1 row)

SELECT lion_index_count_any('lcc_d_c', ARRAY['abc', 'xyz']) AS c_index,
	   (SELECT count(*) FROM lcc_d WHERE t = ANY (ARRAY['abc', 'xyz'])) AS query;
 c_index | query 
---------+-------
     200 |   200
(1 row)

SELECT groups, count FROM lion_index_count_group_stats('lcc_d_c');
 groups | count 
--------+-------
      4 |   400
(1 row)

DROP TABLE lcc_d;
-- ---------- a multi-key query is taken apart under the index's collation ----------
-- as the executor takes it apart (lion_multikey_query_mode()), and not under
-- its literal's: `tags @> ARRAY[...]` over a column COLLATE "C" compares
-- under C while the literal keeps the default collation.  Under any two
-- deterministic collations the extraction is the same - an element longer
-- than a key is hashed, and makes the query a superset rechecked in the heap
-- - so this only checks the counts; a nondeterministic collation, which does
-- not hash, would tell the two apart.
CREATE TABLE lcc_a (id int, tags text[] COLLATE "C")
	WITH (autovacuum_enabled = off);
INSERT INTO lcc_a SELECT g, ARRAY['k' || g % 5, repeat('L', 2100) || g % 3]
  FROM generate_series(1, 300) g;
CREATE INDEX lcc_a_tags ON lcc_a USING lion (tags);
VACUUM ANALYZE lcc_a;
SET enable_seqscan = off;
SET enable_bitmapscan = off;
SET pg_lion.enable_plain_scan = off;
EXPLAIN (COSTS OFF) SELECT count(*) FROM lcc_a WHERE tags @> ARRAY['k1'];
                     QUERY PLAN                      
-----------------------------------------------------
 Custom Scan (LionCount)
   Lion Indexes: lcc_a_tags (tags @> '{k1}'::text[])
(2 rows)

SELECT count(*) FROM lcc_a WHERE tags @> ARRAY['k1'];
 count 
-------
    60
(1 row)

SELECT count(*) FROM lcc_a WHERE tags @> ARRAY['k1', repeat('L', 2100) || '1'];
 count 
-------
    20
(1 row)

SELECT count(*) FROM lcc_a WHERE tags @> ARRAY[repeat('L', 2100) || '2'];
 count 
-------
   100
(1 row)

SET pg_lion.enable_count_pushdown = off;
SELECT count(*) FROM lcc_a WHERE tags @> ARRAY['k1'];
 count 
-------
    60
(1 row)

SELECT count(*) FROM lcc_a WHERE tags @> ARRAY['k1', repeat('L', 2100) || '1'];
 count 
-------
    20
(1 row)

SELECT count(*) FROM lcc_a WHERE tags @> ARRAY[repeat('L', 2100) || '2'];
 count 
-------
   100
(1 row)

RESET pg_lion.enable_count_pushdown;
RESET pg_lion.enable_plain_scan;
RESET enable_bitmapscan;
RESET enable_seqscan;
DROP TABLE lcc_a;
-- ---------- a nondeterministic collation agrees with no other ----------
-- the collation, where the build has ICU and the database is UTF8: core
-- lists its "unicode" ICU collation in pg_collation even in a build without
-- ICU, so whether one can be created is the only test
SET client_min_messages = warning;
DO $$
BEGIN
	IF getdatabaseencoding() = 'UTF8' THEN
		CREATE COLLATION lcc_ci (provider = icu, locale = 'und-u-ks-level2',
								 deterministic = false);
	END IF;
EXCEPTION WHEN feature_not_supported THEN
	NULL;						-- a build without ICU
END $$;
RESET client_min_messages;
SELECT NOT EXISTS (SELECT 1 FROM pg_collation WHERE collname = 'lcc_ci')
	   AS lcc_skip \gset
\if :lcc_skip
\echo 'no ICU collations, or not a UTF8 database: skipped'
\else
-- a column of the default collation, and an index under the
-- case-insensitive one: 'abc', 'ABC' and 'Abc' are one key of the index
CREATE TABLE lcc (id int, m int, t text) WITH (autovacuum_enabled = off);
INSERT INTO lcc SELECT g, g % 2, (ARRAY['abc', 'ABC', 'Abc', 'xyz'])[1 + g % 4]
  FROM generate_series(1, 400) g;
CREATE INDEX lcc_ci_i ON lcc USING lion (t COLLATE lcc_ci);
CREATE INDEX lcc_m ON lcc USING lion (m);
-- `t = 'abc'` compares under the column's collation, and counts 100 rows
SELECT count(*) FROM lcc WHERE t = 'abc';
 count 
-------
   100
(1 row)

SELECT lion_index_count('lcc_ci_i', 'abc'::text);
ERROR:  cannot count through index "lcc_ci_i" under collation "default"
SELECT count FROM lion_index_count_stats('lcc_ci_i', 'abc'::text);
ERROR:  cannot count through index "lcc_ci_i" under collation "default"
SELECT lion_index_count_any('lcc_ci_i', ARRAY['abc', 'xyz']);
ERROR:  cannot count through index "lcc_ci_i" under collation "default"
SELECT lion_index_count('lcc_m', 0, 'lcc_ci_i', 'abc'::text);
ERROR:  cannot count through index "lcc_ci_i" under collation "default"
SELECT groups, count FROM lion_index_count_group_stats('lcc_ci_i');
ERROR:  cannot count through index "lcc_ci_i" under collation "default"
-- under the index's own collation, the count is the query's
SELECT lion_index_count('lcc_ci_i', 'abc'::text COLLATE lcc_ci) AS lion,
	   (SELECT count(*) FROM lcc WHERE t COLLATE lcc_ci = 'abc') AS query;
 lion | query 
------+-------
  300 |   300
(1 row)

SELECT count FROM lion_index_count_stats('lcc_ci_i', 'abc'::text COLLATE lcc_ci);
 count 
-------
   300
(1 row)

SELECT lion_index_count_any('lcc_ci_i', ARRAY['abc', 'xyz'] COLLATE lcc_ci) AS lion,
	   (SELECT count(*) FROM lcc WHERE t COLLATE lcc_ci = ANY (ARRAY['abc', 'xyz'])) AS query;
 lion | query 
------+-------
  400 |   400
(1 row)

-- each key of the two-index form under its own collation
SELECT lion_index_count('lcc_ci_i', 'abc'::text COLLATE lcc_ci, 'lcc_m', 0) AS lion,
	   (SELECT count(*) FROM lcc WHERE t COLLATE lcc_ci = 'abc' AND m = 0) AS query;
 lion | query 
------+-------
  200 |   200
(1 row)

-- a column of the case-insensitive collation: its own index counts as its
-- queries do, an index under "C" does not.  A key of the default collation -
-- a literal's - is refused even by its own index: the call cannot tell it
-- from an explicit COLLATE "default", under which the query does not count
-- 'ABC' (countcoll_default.sql).  Named, the column's collation is counted.
CREATE TABLE lcc2 (id int, t text COLLATE lcc_ci) WITH (autovacuum_enabled = off);
INSERT INTO lcc2 SELECT g, (ARRAY['abc', 'ABC', 'Abc', 'xyz'])[1 + g % 4]
  FROM generate_series(1, 400) g;
CREATE INDEX lcc2_i ON lcc2 USING lion (t);
CREATE INDEX lcc2_c ON lcc2 USING lion (t COLLATE "C");
SELECT lion_index_count('lcc2_i', 'abc'::text) AS lion,
	   (SELECT count(*) FROM lcc2 WHERE t = 'abc') AS query;
ERROR:  could not determine which collation to use for a count through index "lcc2_i"
SELECT lion_index_count_any('lcc2_i', ARRAY['abc', 'xyz']) AS lion,
	   (SELECT count(*) FROM lcc2 WHERE t = ANY (ARRAY['abc', 'xyz'])) AS query;
ERROR:  could not determine which collation to use for a count through index "lcc2_i"
SELECT lion_index_count('lcc2_i', 'abc'::text COLLATE lcc_ci) AS lion,
	   (SELECT count(*) FROM lcc2 WHERE t = 'abc') AS query;
 lion | query 
------+-------
  300 |   300
(1 row)

SELECT lion_index_count_any('lcc2_i', ARRAY['abc', 'xyz'] COLLATE lcc_ci) AS lion,
	   (SELECT count(*) FROM lcc2 WHERE t = ANY (ARRAY['abc', 'xyz'])) AS query;
 lion | query 
------+-------
  400 |   400
(1 row)

SELECT groups, count FROM lion_index_count_group_stats('lcc2_i');
 groups | count 
--------+-------
      2 |   400
(1 row)

SELECT lion_index_count('lcc2_i', 'abc'::text COLLATE "C");
ERROR:  cannot count through index "lcc2_i" under collation "C"
SELECT lion_index_count('lcc2_c', 'abc'::text);
ERROR:  cannot count through index "lcc2_c" under collation "lcc_ci"
SELECT groups, count FROM lion_index_count_group_stats('lcc2_c');
ERROR:  cannot count through index "lcc2_c" under collation "lcc_ci"
SELECT lion_index_count('lcc2_c', 'abc'::text COLLATE "C") AS lion,
	   (SELECT count(*) FROM lcc2 WHERE t = 'abc' COLLATE "C") AS query;
 lion | query 
------+-------
  100 |   100
(1 row)

-- ---------- GROUP BY under a named collation ----------
-- The groups are the grouping expression's collation's, not the column's:
-- `GROUP BY t COLLATE c` may only be driven by an index built under c
-- (lion_collect_targets()).  Under a nondeterministic collation none is
-- pushed down - the group key counts as printed, and an entry's key is no
-- spelling the rows need have; under "C" lcc2_c drives them - and every answer
-- is checked against the same query with the pushdown off.  lcc's column is
-- of the default collation and its one index on t is under lcc_ci; lcc2's
-- column is under lcc_ci.
CREATE FUNCTION lcc_groups(q text) RETURNS text
LANGUAGE plpgsql AS $$
DECLARE
	ln text;
	how text := 'not pushed down';
	lion text;
	core text;
BEGIN
	PERFORM set_config('pg_lion.enable_count_pushdown', 'on', true);
	FOR ln IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP
		IF ln LIKE '%Custom Scan (LionCount)%' THEN
			how := 'LionCount';
		END IF;
	END LOOP;
	EXECUTE format('SELECT string_agg(s::text, '' '') FROM (%s) s', q) INTO lion;
	PERFORM set_config('pg_lion.enable_count_pushdown', 'off', true);
	EXECUTE format('SELECT string_agg(s::text, '' '') FROM (%s) s', q) INTO core;
	PERFORM set_config('pg_lion.enable_count_pushdown', 'on', true);
	IF lion IS DISTINCT FROM core THEN
		RETURN format('MISMATCH: %s got %s, the pushdown off %s', how, lion, core);
	END IF;
	RETURN format('%s: %s', how, lion);
END $$;
SET enable_hashagg = off;
SET enable_seqscan = off;
-- under the column's own collation: four groups; lcc_ci_i is not the index for it
SELECT lcc_groups('SELECT count(*) FROM lcc GROUP BY t ORDER BY 1');
                lcc_groups                
------------------------------------------
 not pushed down: (100) (100) (100) (100)
(1 row)

-- under lcc_ci: two groups, from lcc_ci_i
SELECT lcc_groups('SELECT count(*) FROM lcc GROUP BY t COLLATE lcc_ci ORDER BY 1');
          lcc_groups          
------------------------------
 not pushed down: (100) (300)
(1 row)

SELECT lcc_groups('SELECT count(*) FROM lcc GROUP BY t COLLATE lcc_ci HAVING count(*) > 100');
       lcc_groups       
------------------------
 not pushed down: (300)
(1 row)

SELECT lcc_groups('SELECT m, count(*) FROM lcc GROUP BY t COLLATE lcc_ci, m ORDER BY 1, 2');
                lcc_groups                
------------------------------------------
 not pushed down: (0,200) (1,100) (1,100)
(1 row)

-- the column under lcc_ci: its own groups from lcc2_i, and four under "C",
-- which lcc2_i cannot drive and lcc2_c does
SELECT lcc_groups('SELECT count(*) FROM lcc2 GROUP BY t ORDER BY 1');
          lcc_groups          
------------------------------
 not pushed down: (100) (300)
(1 row)

SELECT lcc_groups('SELECT count(*) FROM lcc2 GROUP BY t COLLATE "C" ORDER BY 1');
             lcc_groups             
------------------------------------
 LionCount: (100) (100) (100) (100)
(1 row)

SELECT lcc_groups('SELECT count(*) FROM lcc2 GROUP BY t COLLATE "C" HAVING count(*) > 100');
 lcc_groups  
-------------
 LionCount: 
(1 row)

SELECT lcc_groups('SELECT t COLLATE "C", count(*) FROM lcc2 GROUP BY 1 ORDER BY 1');
                     lcc_groups                     
----------------------------------------------------
 LionCount: (ABC,100) (Abc,100) (abc,100) (xyz,100)
(1 row)

RESET enable_hashagg;
RESET enable_seqscan;
DROP FUNCTION lcc_groups(text);
DROP TABLE lcc, lcc2;
DROP COLLATION lcc_ci;
\endif