citus/src/test/regress/sql/columnar_test_helpers.sql

SET client_min_messages TO WARNING;
DROP SCHEMA IF EXISTS columnar_test_helpers CASCADE;
RESET client_min_messages;

CREATE SCHEMA columnar_test_helpers;
SET search_path TO columnar_test_helpers;

CREATE OR REPLACE FUNCTION columnar_storage_info(
    rel regclass,
    version_major OUT int4,
    version_minor OUT int4,
    storage_id OUT int8,
    reserved_stripe_id OUT int8,
    reserved_row_number OUT int8,
    reserved_offset OUT int8)
  STRICT
  LANGUAGE c AS 'citus', $$columnar_storage_info$$;

CREATE FUNCTION compression_type_supported(type text) RETURNS boolean
AS $$
BEGIN
   EXECUTE 'SET LOCAL columnar.compression TO ' || quote_literal(type);
   return true;
EXCEPTION WHEN invalid_parameter_value THEN
   return false;
END;
$$ LANGUAGE plpgsql;

-- are chunk groups and chunks consistent?
CREATE view chunk_group_consistency AS
WITH a as (
   SELECT storage_id, stripe_num, chunk_group_num, min(value_count) as row_count
   FROM columnar_internal.chunk
   GROUP BY 1,2,3
), b as (
   SELECT storage_id, stripe_num, chunk_group_num, max(value_count) as row_count
   FROM columnar_internal.chunk
   GROUP BY 1,2,3
), c as (
   (TABLE a EXCEPT TABLE b) UNION (TABLE b EXCEPT TABLE a) UNION
   (TABLE a EXCEPT TABLE columnar_internal.chunk_group) UNION (TABLE columnar_internal.chunk_group EXCEPT TABLE a)
), d as (
   SELECT storage_id, stripe_num, count(*) as chunk_group_count
   FROM columnar_internal.chunk_group
   GROUP BY 1,2
), e as (
   SELECT storage_id, stripe_num, chunk_group_count
   FROM columnar_internal.stripe
), f as (
   (TABLE d EXCEPT TABLE e) UNION (TABLE e EXCEPT TABLE d)
)
SELECT (SELECT count(*) = 0 FROM c) AND
       (SELECT count(*) = 0 FROM f) as consistent;

CREATE FUNCTION columnar_metadata_has_storage_id(input_storage_id bigint) RETURNS boolean
AS $$
DECLARE
   union_storage_id_count integer;
BEGIN
   SELECT count(*) INTO union_storage_id_count FROM
   (
   SELECT storage_id FROM columnar_internal.stripe UNION ALL
   SELECT storage_id FROM columnar_internal.chunk UNION ALL
   SELECT storage_id FROM columnar_internal.chunk_group
   ) AS union_storage_id
   WHERE storage_id=input_storage_id;

   IF union_storage_id_count > 0 THEN
   RETURN true;
   END IF;

   RETURN false;
END;
$$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION columnar_store_memory_stats(
                  OUT TopMemoryContext BIGINT,
		          OUT TopTransactionContext BIGINT,
		          OUT WriteStateContext BIGINT)
    RETURNS RECORD
    LANGUAGE C STRICT VOLATILE
    AS 'citus', $$columnar_store_memory_stats$$;

CREATE FUNCTION top_memory_context_usage()
	RETURNS BIGINT AS $$
		SELECT TopMemoryContext FROM columnar_test_helpers.columnar_store_memory_stats();
	$$ LANGUAGE SQL VOLATILE;

CREATE OR REPLACE FUNCTION uses_index_scan(command text)
RETURNS BOOLEAN AS $$
DECLARE
  query_plan text;
BEGIN
  FOR query_plan IN EXECUTE 'EXPLAIN' || command LOOP
    IF query_plan ILIKE '%Index Only Scan using%' OR
       query_plan ILIKE '%Index Scan using%'
    THEN
        RETURN true;
    END IF;
  END LOOP;
  RETURN false;
END; $$ language plpgsql;

CREATE OR REPLACE FUNCTION uses_custom_scan(command text)
RETURNS BOOLEAN AS $$
DECLARE
  query_plan text;
BEGIN
  FOR query_plan IN EXECUTE 'EXPLAIN' || command LOOP
    IF query_plan ILIKE '%Custom Scan (ColumnarScan)%'
    THEN
        RETURN true;
    END IF;
  END LOOP;
  RETURN false;
END; $$ language plpgsql;

CREATE OR REPLACE FUNCTION uses_seq_scan(command text)
RETURNS BOOLEAN AS $$
DECLARE
  query_plan text;
BEGIN
  FOR query_plan IN EXECUTE 'EXPLAIN' || command LOOP
    IF query_plan ILIKE '%Seq Scan on %'
    THEN
        RETURN true;
    END IF;
  END LOOP;
  RETURN false;
END; $$ language plpgsql;

CREATE OR REPLACE FUNCTION pg_waitpid(p_pid integer)
RETURNS VOID AS $$
BEGIN
  WHILE EXISTS (SELECT * FROM pg_stat_activity WHERE pid=p_pid)
  LOOP
    PERFORM pg_sleep(0.001);
  END LOOP;
END; $$ language plpgsql;

-- This function formats EXPLAIN output to conform to how pg <= 16 EXPLAIN
-- shows ANY <subquery> in an expression the pg version >= 17. When 17 is
-- the minimum supported pgversion this function can be retired. The commit
-- that changed how ANY <subquery> exrpressions appear in EXPLAIN is:
-- https://git.postgresql.org/gitweb/?p=postgresql.git;a=commitdiff;h=fd0398fcb
CREATE OR REPLACE FUNCTION explain_with_pg16_subplan_format(explain_command text, out query_plan text)
RETURNS SETOF TEXT AS $$
DECLARE
  pgversion int = 0;
BEGIN
  pgversion = substring(version(), '\d+')::int ;
  FOR query_plan IN execute explain_command LOOP
    IF pgversion >= 17 THEN
      IF query_plan ~ 'SubPlan \d+\).col' THEN
    	  query_plan = regexp_replace(query_plan, '\(ANY \(\w+ = \(SubPlan (\d+)\).col1\)\)', '(SubPlan \1)', 'g');
      END IF;
    END IF;
    RETURN NEXT;
  END LOOP;
END; $$ language plpgsql;

CREATE OR REPLACE FUNCTION _plan_json(q text)
RETURNS jsonb
LANGUAGE plpgsql AS $$
DECLARE j jsonb;
BEGIN
  EXECUTE format('EXPLAIN (FORMAT JSON, COSTS OFF, ANALYZE OFF) %s', q) INTO j;
  RETURN j;
END $$;