-- Tests for Semi and Anti Joins with Join Custom Scan
CREATE EXTENSION IF NOT EXISTS pg_search;
SET max_parallel_workers_per_gather = 0;

DROP TABLE IF EXISTS table_a CASCADE;
CREATE TABLE table_a (
    id bigint NOT NULL PRIMARY KEY,
    category character varying
);

INSERT INTO table_a (id, category)
SELECT
    i,
    CASE WHEN i % 2 = 0 THEN 'target_category' ELSE 'other_category' END
FROM generate_series(1, 2000) as i;

DROP TABLE IF EXISTS table_b CASCADE;
CREATE TABLE table_b (
    id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
    group_id character varying,
    a_id bigint,
    category character varying
);

INSERT INTO table_b (group_id, a_id, category)
SELECT
    CASE WHEN i % 10 = 0 THEN 'group_1' ELSE 'group_2' END,
    i,
    CASE WHEN i % 2 = 0 THEN 'target_category' ELSE 'other_category' END
FROM generate_series(1, 2000) as i;

CREATE INDEX table_a_idx ON table_a
USING bm25 (id, category)
WITH (
    key_field = id,
    text_fields = '{
        "category": {"fast": true, "tokenizer": {"type": "keyword"}, "normalizer": "lowercase"}
    }'
);

CREATE INDEX table_b_group_id_idx ON table_b USING btree (group_id);
CREATE INDEX table_b_group_id_a_id_idx ON table_b USING btree (group_id, a_id);

CREATE INDEX table_b_idx ON table_b
USING bm25 (id, group_id, a_id, category)
WITH (
    key_field = id,
    text_fields = '{
        "group_id": {"fast": true, "tokenizer": {"type": "keyword"}},
        "category": {"fast": true, "tokenizer": {"type": "keyword"}, "normalizer": "lowercase"}
    }',
    numeric_fields = '{
        "a_id": {"fast": true}
    }'
);

-- Query execution
SET paradedb.enable_join_custom_scan TO on;

-- =====================================================================
-- 1. Semi Join Only
-- =====================================================================
EXPLAIN (COSTS OFF, VERBOSE, TIMING OFF)
SELECT id, category
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id @@@ 'category:"target_category"'
ORDER BY id ASC
LIMIT 10;

-- =====================================================================
-- 2a. Anti Join (No Pushdown)
-- =====================================================================
-- This does not trigger an anti join because Postgres cannot prove that group_id is not NULL.
-- TODO: This query only triggers set_rel_pathlist_hook and not set_join_pathlist_hook, and so does
-- not render our warning. See https://github.com/paradedb/paradedb/issues/4236 about resolving this.
SELECT id, category
FROM table_a
WHERE id NOT IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_3', 'group_4')
)
AND id @@@ 'category:"target_category"'
ORDER BY id ASC
LIMIT 10;

-- =====================================================================
-- 2b. Anti Join (Pushdown)
-- =====================================================================
-- This does trigger an anti join because we use NOT EXISTS
EXPLAIN (COSTS OFF, VERBOSE, TIMING OFF)
SELECT id, category
FROM table_a
WHERE NOT EXISTS (
    SELECT 1
    FROM table_b
    WHERE table_b.a_id = table_a.id
    AND table_b.group_id IN ('group_3', 'group_4')
)
AND table_a.id @@@ 'category:"target_category"'
ORDER BY table_a.id ASC
LIMIT 10;

-- =====================================================================
-- 3. Both Semi and Anti Join
-- =====================================================================
EXPLAIN (COSTS OFF, VERBOSE, TIMING OFF)
SELECT id, category
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id NOT IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_3', 'group_4')
)
AND id @@@ 'category:"target_category"'
ORDER BY id ASC
LIMIT 10;

SELECT id, category
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id NOT IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_3', 'group_4')
)
AND id @@@ 'category:"target_category"'
ORDER BY id ASC
LIMIT 10;

-- =====================================================================
-- 4. Tuple Semi Join
-- =====================================================================
EXPLAIN (COSTS OFF, VERBOSE, TIMING OFF)
SELECT id, category
FROM table_a
WHERE (id, category) IN (
    SELECT a_id, category
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id @@@ pdb.all()
ORDER BY id ASC
LIMIT 10;

SELECT id, category
FROM table_a
WHERE (id, category) IN (
    SELECT a_id, category
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id @@@ pdb.all()
ORDER BY id ASC
LIMIT 10;

-- =====================================================================
-- 5. SELECT * and ORDER BY varchar column (RelabelType wrapper reproduction)
-- =====================================================================
EXPLAIN (COSTS OFF, VERBOSE, TIMING OFF)
SELECT *
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id @@@ 'category:"target_category"'
ORDER BY category ASC, id ASC
LIMIT 10;

SELECT *
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
    WHERE group_id IN ('group_1')
)
AND id @@@ 'category:"target_category"'
ORDER BY category ASC, id ASC
LIMIT 10;

-- =====================================================================
-- 6. Semi Join with selective outer filter
-- =====================================================================
-- The outer filter is so selective that this query requires our ability to force the left side
-- to be partitioned (even when that might be less efficient).
EXPLAIN (COSTS OFF, VERBOSE, TIMING OFF)
SELECT id, category
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
)
AND id @@@ 'id:1'
ORDER BY id ASC
LIMIT 10;

SELECT id, category
FROM table_a
WHERE id IN (
    SELECT a_id
    FROM table_b
)
AND id @@@ 'id:1'
ORDER BY id ASC
LIMIT 10;

-- Cleanup
DROP TABLE table_a CASCADE;
DROP TABLE table_b CASCADE;

RESET paradedb.enable_join_custom_scan;
