--ALTER SYSTEM SET max_parallel_workers=50;
--ALTER SYSTEM SET max_worker_processes=100;

DROP FUNCTION IF EXISTS partition_by_hash(VARCHAR, INT);
CREATE FUNCTION partition_by_hash(table_name VARCHAR, num_partitions INT) RETURNS VOID AS $$
DECLARE
  n INT;
BEGIN
    FOR n IN 1..num_partitions LOOP
        EXECUTE '
            CREATE TABLE ' || table_name || '__part_' || n || '
            PARTITION OF ' || table_name || '
            FOR VALUES WITH (MODULUS ' || num_partitions || ', REMAINDER ' || n - 1  || ')';
    END LOOP;
END;
$$ LANGUAGE plpgsql;

DROP FUNCTION get_explain(TEXT);
CREATE FUNCTION get_explain(query TEXT) RETURNS SETOF TEXT AS $q$
BEGIN
    RETURN QUERY EXECUTE 'EXPLAIN (FORMAT YAML, ANALYZE) ' || query;
END;
$q$ LANGUAGE plpgsql;

DROP FUNCTION IF EXISTS run_query_count(TEXT, TEXT, INT);
CREATE FUNCTION run_query_count(name TEXT, query TEXT, count INT) RETURNS SETOF TEXT AS $q$
DECLARE
    function_count TEXT;
    loop_start TIMESTAMP;
    loop_end TIMESTAMP;
BEGIN
    SELECT coalesce(substring(substring(output from 'Functions:\s+[0-9]+'), '[0-9]+$'), 'none') INTO function_count
        FROM get_explain(query) AS output;
    SELECT clock_timestamp() INTO loop_start;
    FOR i in 1..count LOOP
        EXECUTE query;
    END LOOP;
    SELECT clock_timestamp() INTO loop_end;
    RETURN NEXT name || ': '
        || ((loop_end - loop_start) / count)
        || ' #runs=' || count || ' #JIT=' || function_count;
END;
$q$ LANGUAGE plpgsql;

DROP FUNCTION IF EXISTS run_queries();
CREATE FUNCTION run_queries() RETURNS SETOF TEXT AS $q$
BEGIN
    RETURN NEXT 'parameters: jit=' || current_setting('jit') || ' workers=' || current_setting('max_parallel_workers_per_gather') || ' jit-inline=' || current_setting('jit_inline_above_cost') || ' jit-optimize=' || current_setting('jit_optimize_above_cost');
    RETURN QUERY SELECT run_query_count('query1', 'SELECT SUM(col1) / 7.0 as avg_yearly, COUNT(*) AS cnt FROM parent WHERE col1 < 10000 AND col1 != 1 AND col1 != 2 AND col1 != 3 AND col1 != 4 AND col1 != 5', 5);
    RETURN QUERY SELECT run_query_count('query2', 'SELECT COUNT(*) FROM (SELECT col0 AS key, SUM(col1) / 7.0 as avg_yearly, COUNT(*) AS cnt FROM parent WHERE col1 < 10000 AND col1 != 1 AND col1 != 2 AND col1 != 3 AND col1 != 4 AND col1 != 5 GROUP BY col0) a', 5);
    RETURN QUERY SELECT run_query_count('query3', 'SELECT COUNT(*) FROM (SELECT col0 AS key, SUM(col1) / 7.0 as avg_yearly, COUNT(*) AS cnt FROM single_big WHERE col1 < 10000 AND col1 != 1 AND col1 != 2 AND col1 != 3 AND col1 != 4 AND col1 != 5 GROUP BY col0) a', 5);
    RETURN QUERY SELECT run_query_count('query4', 'SELECT COUNT(*) FROM (SELECT col0 AS key, SUM(col1) / 7.0 as avg_yearly, COUNT(*) AS cnt FROM single_small WHERE col1 < 10000 AND col1 != 1 AND col1 != 2 AND col1 != 3 AND col1 != 4 AND col1 != 5 GROUP BY col0) a', 100);
END;
$q$ LANGUAGE plpgsql;


CREATE TABLE parent (col0 INT NOT NULL, col1 double precision NOT NULL) PARTITION BY HASH(col0);
SELECT partition_by_hash('parent', 500);
INSERT INTO parent SELECT n/1000, n+10 FROM generate_series(1, 10000000) n;
CREATE TABLE single_big AS SELECT * FROM parent;
CREATE TABLE single_small AS SELECT * FROM single_big LIMIT 10;

ANALYZE parent;
ANALYZE single_big;
ANALYZE single_small;

-- max out the parallel workers for the tables
ALTER TABLE parent__part_1 SET (parallel_workers=50);
ALTER TABLE single_big SET (parallel_workers=50);
ALTER TABLE single_small SET (parallel_workers=50);

-- enable all jit options to emulate a more expensive query
SET jit_above_cost = 0;
SET jit_inline_above_cost = 0;
SET jit_optimize_above_cost = 0;
SET jit_tuple_deforming = on;
SET jit_expressions = on;
SET enable_partitionwise_aggregate = on;
SET enable_partitionwise_join = on;


SET jit=off;
SET max_parallel_workers_per_gather=5;
SELECT run_queries();
SET max_parallel_workers_per_gather=50;
SELECT run_queries();

SET jit=on;
SET max_parallel_workers_per_gather=5;
SELECT run_queries();
SET max_parallel_workers_per_gather=50;
SELECT run_queries();
SET jit_inline_above_cost   = 1000000;
SELECT run_queries();
SET jit_inline_above_cost   = 0;
SET jit_optimize_above_cost = 1000000;
SELECT run_queries();
SET jit_inline_above_cost   = 1000000;
SELECT run_queries();

