diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 6f12121..7e02165 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -319,6 +319,90 @@ mod tests { assert!(rows[0].1 > rows[1].1); } + #[pg_test] + fn scoring_and_inspection_respect_partial_index_predicates() { + Spi::run( + "CREATE TABLE lite_partial_score (id int, body text, active boolean); + INSERT INTO lite_partial_score VALUES + (1, 'beer', true), (2, 'wine', true), + (3, 'wine', NULL), (4, NULL, true); + INSERT INTO lite_partial_score + SELECT n, 'wine', false FROM generate_series(5, 104) AS n; + CREATE INDEX lite_partial_score_idx ON lite_partial_score + USING tin (body) WHERE active; + CREATE TABLE lite_partial_score_control AS + SELECT id, body FROM lite_partial_score WHERE active; + CREATE INDEX lite_partial_score_control_idx ON lite_partial_score_control + USING tin (body);", + ) + .unwrap(); + let partial = Spi::get_one::( + "SELECT tin.full_score(ctid) FROM lite_partial_score + WHERE active AND body ==> 'beer'", + ) + .unwrap() + .unwrap(); + let control = Spi::get_one::( + "SELECT tin.full_score(ctid) FROM lite_partial_score_control + WHERE body ==> 'beer'", + ) + .unwrap() + .unwrap(); + assert!(control > 0.0); + assert_eq!(partial, control); + + // In the indexed population, beer occurs in half the documents and + // must be elided at the default dense ratio, despite the excluded rows. + assert_eq!( + Spi::get_one::( + "SELECT count(*) FROM tin.score_inspect('lite_partial_score_idx', 'beer')" + ) + .unwrap(), + Some(0) + ); + assert_eq!( + Spi::get_one::( + "SELECT tin.score(ctid) FROM lite_partial_score + WHERE active AND body ==> 'beer'" + ) + .unwrap(), + Some(0.0) + ); + } + + #[pg_test] + fn scoring_respects_partial_expression_index_predicates() { + Spi::run( + "CREATE TABLE lite_partial_expression (id int, body text, active boolean); + INSERT INTO lite_partial_expression VALUES + (1, 'BEER', true), (2, 'wine wine', true), + (3, 'BEER BEER', false), (4, 'excluded', false), + (5, 'excluded', NULL), (6, NULL, true); + CREATE INDEX lite_partial_expression_idx ON lite_partial_expression + USING tin (lower(body)) WHERE active OR id = 3; + CREATE TABLE lite_partial_expression_control AS + SELECT id, body FROM lite_partial_expression WHERE active OR id = 3; + CREATE INDEX lite_partial_expression_control_idx + ON lite_partial_expression_control USING tin (lower(body));", + ) + .unwrap(); + let partial = Spi::get_one::>( + "SELECT array_agg(tin.full_score(ctid) ORDER BY id) + FROM lite_partial_expression + WHERE (active OR id = 3) AND lower(body) ==> 'beer'", + ) + .unwrap() + .unwrap(); + let control = Spi::get_one::>( + "SELECT array_agg(tin.full_score(ctid) ORDER BY id) + FROM lite_partial_expression_control WHERE lower(body) ==> 'beer'", + ) + .unwrap() + .unwrap(); + assert_eq!(control.len(), 2); + assert_eq!(partial, control); + } + #[pg_test] fn highlighting_supports_explicit_and_implicit_queries() { assert_eq!( diff --git a/postgres/src/score.rs b/postgres/src/score.rs index 2594575..6dda5e9 100644 --- a/postgres/src/score.rs +++ b/postgres/src/score.rs @@ -230,10 +230,11 @@ fn load_documents(heap_oid: pg_sys::Oid, index_oid: pg_sys::Oid) -> Vec pgrx::error!("tin score relation no longer exists"); } let qualified = pg_sys::quote_qualified_identifier(namespace, relname); - let expression_sql = format!( + let index_sql = format!( "SELECT CASE WHEN i.indkey[0] = 0 \ THEN pg_catalog.pg_get_expr(i.indexprs, i.indrelid) \ - ELSE pg_catalog.quote_ident(a.attname) END \ + ELSE pg_catalog.quote_ident(a.attname) END, \ + pg_catalog.pg_get_expr(i.indpred, i.indrelid) \ FROM pg_catalog.pg_index i \ LEFT JOIN pg_catalog.pg_attribute a \ ON a.attrelid=i.indrelid AND a.attnum=i.indkey[0] \ @@ -241,13 +242,15 @@ fn load_documents(heap_oid: pg_sys::Oid, index_oid: pg_sys::Oid) -> Vec index_oid.to_u32(), heap_oid.to_u32(), ); - let expression = Spi::get_one::(&expression_sql) - .unwrap_or_else(|error| { - pgrx::error!("tin score index expression lookup failed: {error}") - }) + let (expression, predicate) = Spi::get_two::(&index_sql) + .unwrap_or_else(|error| pgrx::error!("tin score index lookup failed: {error}")); + let expression = expression .unwrap_or_else(|| pgrx::error!("tin score index expression no longer exists")); + let predicate = predicate + .map(|predicate| format!(" AND ({predicate})")) + .unwrap_or_default(); let sql = format!( - "SELECT ({expression})::text FROM {} WHERE ({expression}) IS NOT NULL", + "SELECT ({expression})::text FROM {} WHERE ({expression}) IS NOT NULL{predicate}", CStr::from_ptr(qualified).to_string_lossy(), ); Spi::connect(|client| {