diff --git a/postgres/src/highlight_udfs.rs b/postgres/src/highlight_udfs.rs index 9ee916c..b4180da 100644 --- a/postgres/src/highlight_udfs.rs +++ b/postgres/src/highlight_udfs.rs @@ -77,37 +77,6 @@ fn highlight_ansi( render_highlight_ansi(text, wrap_to, query) } -struct VarContext { - varno: i32, - seen: bool, - valid: bool, -} - -#[pg_guard] -unsafe extern "C-unwind" fn collect_varno(node: *mut pg_sys::Node, context: *mut c_void) -> bool { - if node.is_null() { - return false; - } - let context = unsafe { &mut *context.cast::() }; - if unsafe { (*node).type_ } == pg_sys::NodeTag::T_Var { - let var = unsafe { &*node.cast::() }; - if var.varlevelsup != 0 || (context.seen && context.varno != var.varno) { - context.valid = false; - } else { - context.varno = var.varno; - context.seen = true; - } - return false; - } - unsafe { - pg_sys::expression_tree_walker( - node, - Some(collect_varno), - (context as *mut VarContext).cast(), - ) - } -} - struct QueryContext { document: *mut pg_sys::Node, queries: Vec<*mut pg_sys::Node>, @@ -115,7 +84,10 @@ struct QueryContext { #[pg_guard] unsafe extern "C-unwind" fn collect_queries(node: *mut pg_sys::Node, context: *mut c_void) -> bool { - if node.is_null() { + if node.is_null() + || unsafe { (*node).type_ } == pg_sys::NodeTag::T_Query + || crate::score::is_negation(node) + { return false; } let context = unsafe { &mut *context.cast::() }; @@ -147,40 +119,73 @@ fn unhandled() -> Internal { Internal::from(Some(pg_sys::Datum::from(0_usize))) } +unsafe fn text_const(value: &str) -> *mut pg_sys::Node { + let datum = value.into_datum().expect("&str is never NULL"); + unsafe { + pg_sys::makeConst( + pg_sys::TEXTOID, + -1, + pg_sys::DEFAULT_COLLATION_OID, + -1, + datum, + false, + false, + ) + .cast() + } +} + +unsafe fn concatenate(left: *mut pg_sys::Node, right: *mut pg_sys::Node) -> *mut pg_sys::Node { + let mut args = PgList::::new(); + args.push(left); + args.push(right); + unsafe { + pg_sys::makeFuncExpr( + pg_sys::Oid::from(pg_sys::F_TEXTCAT), + pg_sys::TEXTOID, + args.into_pg(), + pg_sys::DEFAULT_COLLATION_OID, + pg_sys::DEFAULT_COLLATION_OID, + pg_sys::CoercionForm::COERCE_EXPLICIT_CALL, + ) + .cast() + } +} + +/// Combines the search expressions of every binding qual into one query. +/// Constant expressions fold at plan time; parameters and other run-time +/// expressions are concatenated by the plan instead of being dropped. unsafe fn combined_query(queries: &[*mut pg_sys::Node]) -> *mut pg_sys::Node { if queries.len() < 2 { return unsafe { pg_sys::copyObjectImpl(queries[0].cast()).cast() }; } + if let Some(text) = unsafe { constant_texts(queries) } { + return unsafe { text_const(&text.join(" OR ")) }; + } + let mut combined = unsafe { text_const("(") }; + for (position, &query) in queries.iter().enumerate() { + if position > 0 { + combined = unsafe { concatenate(combined, text_const(") OR (")) }; + } + combined = unsafe { concatenate(combined, pg_sys::copyObjectImpl(query.cast()).cast()) }; + } + unsafe { concatenate(combined, text_const(")")) } +} + +unsafe fn constant_texts(queries: &[*mut pg_sys::Node]) -> Option> { let mut text = Vec::with_capacity(queries.len()); for &query in queries { if query.is_null() || unsafe { (*query).type_ } != pg_sys::NodeTag::T_Const { - return unsafe { pg_sys::copyObjectImpl(queries[0].cast()).cast() }; + return None; } let value = unsafe { &*query.cast::() }; if value.constisnull || value.consttype != pg_sys::TEXTOID { - return unsafe { pg_sys::copyObjectImpl(queries[0].cast()).cast() }; + return None; } - let Some(value) = (unsafe { String::from_datum(value.constvalue, false) }) else { - return unsafe { pg_sys::copyObjectImpl(queries[0].cast()).cast() }; - }; + let value = unsafe { String::from_datum(value.constvalue, false) }?; text.push(format!("({value})")); } - let datum = text - .join(" OR ") - .into_datum() - .expect("String is never NULL"); - unsafe { - pg_sys::makeConst( - pg_sys::TEXTOID, - -1, - pg_sys::DEFAULT_COLLATION_OID, - -1, - datum, - false, - false, - ) - .cast() - } + Some(text) } #[pg_extern(immutable, parallel_unsafe)] @@ -221,20 +226,14 @@ fn highlight_support(request: Internal) -> Internal { return unhandled(); } let document = pg_sys::list_nth((*request.fcall).args, 0).cast::(); - let mut vars = VarContext { - varno: 0, - seen: false, - valid: true, - }; - collect_varno(document, (&mut vars as *mut VarContext).cast()); - if !vars.valid || !vars.seen { + let Some(varno) = crate::score::single_varno(document) else { return unhandled(); - } + }; let parse = (*request.root).parse; - let rte = pg_sys::list_nth((*parse).rtable, vars.varno - 1).cast::(); + let rte = pg_sys::list_nth((*parse).rtable, varno - 1).cast::(); if rte.is_null() || (*rte).rtekind != pg_sys::RTEKind::RTE_RELATION - || crate::score::find_matching_tin_index((*rte).relid, vars.varno, document).is_none() + || crate::score::find_matching_tin_index((*rte).relid, varno, document).is_none() { return unhandled(); } @@ -242,8 +241,9 @@ fn highlight_support(request: Internal) -> Internal { document, queries: Vec::new(), }; + // Pulled-up subqueries leave their quals in nested FromExpr nodes. collect_queries( - (*(*parse).jointree).quals.cast::(), + (*parse).jointree.cast::(), (&mut binding as *mut QueryContext).cast(), ); if binding.queries.is_empty() { @@ -282,6 +282,44 @@ mod tests { use super::*; use pgrx::pg_test; + fn create_highlight_table() { + pgrx::Spi::run( + "CREATE TABLE lite_highlight_quals (id int, title text); + INSERT INTO lite_highlight_quals VALUES (1, 'alpha zeta filler'); + CREATE INDEX ON lite_highlight_quals USING tin (title);", + ) + .unwrap(); + } + + #[pg_test] + fn parameterized_quals_highlight_every_term() { + create_highlight_table(); + pgrx::Spi::run( + "PREPARE lite_highlight(text, text) AS + SELECT tin.highlight(title) FROM lite_highlight_quals + WHERE title ==> $1 AND title ==> $2", + ) + .unwrap(); + let marked = pgrx::Spi::get_one::("EXECUTE lite_highlight('alpha', 'zeta')") + .unwrap() + .unwrap(); + assert_eq!(marked, "alpha zeta filler"); + } + + #[pg_test] + fn highlighting_survives_subquery_pullup() { + create_highlight_table(); + let marked = pgrx::Spi::get_one::( + "SELECT h FROM ( + SELECT tin.highlight(title) AS h FROM lite_highlight_quals + WHERE title ==> 'zeta' + ) AS marked", + ) + .unwrap() + .unwrap(); + assert_eq!(marked, "alpha zeta filler"); + } + #[pg_test] fn explicit_html_and_ansi_highlighting_render_matches() { assert_eq!( diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 8721b96..0620ad1 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -266,6 +266,116 @@ mod tests { assert_eq!(inspected, Some(vec!["rare".to_owned()])); } + #[pg_test] + fn parameterized_queries_score_like_literals() { + Spi::run( + "CREATE TABLE lite_param_score (id int PRIMARY KEY, title text); + INSERT INTO lite_param_score VALUES (1, 'lorem ipsum'), (2, 'lorem ipsun'); + CREATE INDEX ON lite_param_score USING tin (title); + PREPARE lite_ranked(text, text) AS + SELECT tin.score(ctid) FROM lite_param_score + WHERE title ==> $1 AND title ==> $2 ORDER BY id", + ) + .unwrap(); + let literal = Spi::get_one::( + "SELECT tin.score(ctid) FROM lite_param_score + WHERE title ==> 'lorem^4' + AND title ==> '(ipsum^4 OR ipsum~1^1.4 OR ipsum*^2)' + ORDER BY id", + ) + .unwrap(); + let prepared = Spi::get_one::( + "EXECUTE lite_ranked('lorem^4', '(ipsum^4 OR ipsum~1^1.4 OR ipsum*^2)')", + ) + .unwrap(); + assert_eq!(prepared, literal); + } + + #[pg_test] + fn score_survives_subquery_aggregate() { + Spi::run( + "CREATE TABLE lite_sub_score (id int PRIMARY KEY, title text); + INSERT INTO lite_sub_score VALUES (1, 'lorem ipsum'), (2, 'lorem ipsun'); + CREATE INDEX ON lite_sub_score USING tin (title);", + ) + .unwrap(); + let direct = Spi::get_one::( + "SELECT max(tin.score(ctid)) FROM lite_sub_score WHERE title ==> 'lorem^4'", + ) + .unwrap(); + let nested = Spi::get_one::( + "SELECT max(score) FROM ( + SELECT tin.score(ctid) AS score FROM lite_sub_score WHERE title ==> 'lorem^4' + ) AS matches", + ) + .unwrap(); + assert_eq!(nested, direct); + } + + #[pg_test(error = "tin.score_bound() is missing: the installed tin SQL predates this build")] + fn scoring_reports_an_outdated_installed_extension() { + Spi::run( + "CREATE TABLE lite_stale (title text); + CREATE INDEX ON lite_stale USING tin (title); + ALTER FUNCTION tin.score_bound(text, text[], int4, int4, int4, + float4, float4, float4, text[], text[]) RENAME TO score_bound_stale;", + ) + .unwrap(); + Spi::run("SELECT tin.score(ctid) FROM lite_stale WHERE title ==> 'lorem'").unwrap(); + } + + #[pg_test] + fn quals_on_other_relations_do_not_bind_to_the_scored_relation() { + Spi::run( + "CREATE TABLE lite_cross_a (id int, title text); + INSERT INTO lite_cross_a SELECT g, 'filler word number ' || g + FROM generate_series(1, 30) g; + UPDATE lite_cross_a SET title = 'alpha zeta filler' WHERE id = 1; + CREATE INDEX ON lite_cross_a USING tin (title); + CREATE TABLE lite_cross_b (id int, title text); + INSERT INTO lite_cross_b SELECT g, 'padding text number ' || g + FROM generate_series(1, 30) g; + UPDATE lite_cross_b SET title = 'kappa padding' WHERE id = 1; + CREATE INDEX ON lite_cross_b USING tin (title);", + ) + .unwrap(); + let alone = Spi::get_one::( + "SELECT tin.score(b.ctid) FROM lite_cross_b b WHERE b.title ==> 'kappa'", + ) + .unwrap(); + assert!(alone.is_some_and(|score| score > 0.0), "{alone:?}"); + let joined = Spi::get_one::( + "SELECT tin.score(b.ctid) FROM lite_cross_a a, lite_cross_b b + WHERE a.title ==> 'zeta' AND b.title ==> 'kappa'", + ) + .unwrap(); + assert_eq!(joined, alone); + } + + #[pg_test] + fn negated_quals_do_not_contribute_to_scores() { + Spi::run( + "CREATE TABLE lite_negated (id int, title text); + INSERT INTO lite_negated SELECT g, 'filler word number ' || g + FROM generate_series(1, 30) g; + UPDATE lite_negated SET title = 'alpha zeta filler' WHERE id = 1; + UPDATE lite_negated SET title = 'alpha omega omega omega' WHERE id = 2; + CREATE INDEX ON lite_negated USING tin (title);", + ) + .unwrap(); + let baseline = Spi::get_two::( + "SELECT tin.max_score(ctid), tin.score(ctid) FROM lite_negated + WHERE title ==> 'zeta'", + ) + .unwrap(); + let negated = Spi::get_two::( + "SELECT tin.max_score(ctid), tin.score(ctid) FROM lite_negated + WHERE title ==> 'zeta' AND NOT (title ==> 'omega')", + ) + .unwrap(); + assert_eq!(negated, baseline); + } + #[pg_test] fn max_score_excludes_nonmatching_documents() { for (name, matching, nonmatching, query) in [ diff --git a/postgres/src/score.rs b/postgres/src/score.rs index 30e2959..7e4eb65 100644 --- a/postgres/src/score.rs +++ b/postgres/src/score.rs @@ -20,8 +20,7 @@ use crate::bm25::{ }; use pgrx::iter::TableIterator; use pgrx::{ - FromDatum, Internal, IntoDatum, PgList, PgRelation, Spi, default, name, pg_extern, pg_guard, - pg_sys, + Internal, IntoDatum, PgBox, PgList, PgRelation, Spi, default, name, pg_extern, pg_guard, pg_sys, }; use rustc_hash::FxHashMap; use std::cell::RefCell; @@ -104,7 +103,7 @@ fn bits(value: Option) -> Option { )] fn score_bound( document: &str, - query: &str, + query: Vec>, heap_oid: i32, index_oid: i32, mode: i32, @@ -114,12 +113,19 @@ fn score_bound( term_add: Option>, term_replace: Option>, ) -> f32 { + let mut parts = Vec::with_capacity(query.len()); + for part in query { + // A NULL search expression matches no rows, so nothing needs a score. + let Some(part) = part else { return 0.0 }; + parts.push(format!("({part})")); + } + let query = parts.join(" OR "); let key = CacheKey { transaction: unsafe { pg_sys::GetTopTransactionIdIfAny().into_inner() }, command: unsafe { pg_sys::GetCurrentCommandId(false) }, heap_oid: heap_oid as u32, index_oid: index_oid as u32, - query: query.to_owned(), + query, full: mode == 1 || mode == 3, dense: dense_ratio.unwrap_or(DenseRatio::DEFAULT).to_bits(), k1: bits(k1), @@ -427,9 +433,18 @@ struct QualBinding { matches: Vec<(*mut pg_sys::Node, *mut pg_sys::Node)>, } +/// Search expressions under a `NOT` exclude rows rather than describe them, so +/// they contribute neither scoring terms nor highlight marks. +pub(crate) fn is_negation(node: *mut pg_sys::Node) -> bool { + unsafe { + (*node).type_ == pg_sys::NodeTag::T_BoolExpr + && (*node.cast::()).boolop == pg_sys::BoolExprType::NOT_EXPR + } +} + #[pg_guard] unsafe extern "C-unwind" fn find_qual(node: *mut pg_sys::Node, context: *mut c_void) -> bool { - if node.is_null() { + if node.is_null() || unsafe { (*node).type_ } == pg_sys::NodeTag::T_Query || is_negation(node) { return false; } let binding = unsafe { &mut *context.cast::() }; @@ -450,11 +465,53 @@ unsafe extern "C-unwind" fn find_qual(node: *mut pg_sys::Node, context: *mut c_v unsafe { pg_sys::expression_tree_walker(node, Some(find_qual), context) } } +struct VarnoBinding { + varno: i32, + seen: bool, + valid: bool, +} + +#[pg_guard] +unsafe extern "C-unwind" fn collect_varno(node: *mut pg_sys::Node, context: *mut c_void) -> bool { + if node.is_null() { + return false; + } + let binding = unsafe { &mut *context.cast::() }; + if unsafe { (*node).type_ } == pg_sys::NodeTag::T_Var { + let var = unsafe { &*node.cast::() }; + if var.varlevelsup != 0 || (binding.seen && binding.varno != var.varno) { + binding.valid = false; + } else { + binding.varno = var.varno; + binding.seen = true; + } + return false; + } + unsafe { pg_sys::expression_tree_walker(node, Some(collect_varno), context) } +} + +/// Returns the range table entry that every `Var` in `node` belongs to, or +/// `None` when the node spans several entries, an outer query level, or none. +pub(crate) unsafe fn single_varno(node: *mut pg_sys::Node) -> Option { + let mut binding = VarnoBinding { + varno: 0, + seen: false, + valid: true, + }; + unsafe { collect_varno(node, (&raw mut binding).cast()) }; + (binding.valid && binding.seen).then_some(binding.varno) +} + pub(crate) unsafe fn find_matching_tin_index( heap_oid: pg_sys::Oid, query_varno: i32, operand: *mut pg_sys::Node, ) -> Option { + // The operand is normalized to varno 1 below to compare it against stored + // index expressions, so quals on other relations have to be rejected here. + if unsafe { single_varno(operand) } != Some(query_varno) { + return None; + } let tin_name = CString::new("tin").expect("static access method name is valid"); let tin_am = unsafe { pg_sys::get_index_am_oid(tin_name.as_ptr(), false) }; let normalized = unsafe { pg_sys::copyObjectImpl(operand.cast()).cast::() }; @@ -580,8 +637,10 @@ fn score_support(request: Internal) -> Internal { let mut binding = QualBinding { matches: Vec::new(), }; - let quals = (*(*parse).jointree).quals.cast::(); - find_qual(quals, (&mut binding as *mut QualBinding).cast()); + // Pulled-up subqueries leave their quals in nested FromExpr nodes, so + // walk the whole jointree rather than only its top-level quals. + let jointree = (*parse).jointree.cast::(); + find_qual(jointree, (&mut binding as *mut QualBinding).cast()); let rte = pg_sys::list_nth((*parse).rtable, (ctid.varno - 1) as i32) .cast::(); if rte.is_null() || (*rte).rtekind != pg_sys::RTEKind::RTE_RELATION { @@ -625,9 +684,7 @@ fn score_support(request: Internal) -> Internal { .copied() .filter(|(candidate, _)| pg_sys::equal((*candidate).cast(), document.cast())) .collect::>(); - let combined_query = combine_constant_queries(&same_expression) - .unwrap_or_else(|| pg_sys::copyObjectImpl(first_query.cast()).cast()); - args.push(combined_query); + args.push(make_query_array(&same_expression, first_query)); args.push(make_int4_const((*rte).relid.to_u32() as i32).cast()); args.push(make_int4_const(index_oid.to_u32() as i32).cast()); args.push(make_int4_const(mode).cast()); @@ -673,41 +730,31 @@ fn score_support(request: Internal) -> Internal { } } -unsafe fn combine_constant_queries( +/// Builds the `text[]` of search expressions that the scorer combines. Passing +/// the expressions as an array keeps parameters and other non-constant nodes, +/// which cannot be combined at plan time, contributing to the scores. +unsafe fn make_query_array( matches: &[(*mut pg_sys::Node, *mut pg_sys::Node)], -) -> Option<*mut pg_sys::Node> { - if matches.len() < 2 { - return None; - } - let mut queries = Vec::with_capacity(matches.len()); + first_query: *mut pg_sys::Node, +) -> *mut pg_sys::Node { + let mut elements = PgList::::new(); for &(_, node) in matches { - if node.is_null() || unsafe { (*node).type_ } != pg_sys::NodeTag::T_Const { - return None; - } - let value = unsafe { &*node.cast::() }; - if value.constisnull || value.consttype != pg_sys::TEXTOID { - return None; + if node.is_null() || unsafe { pg_sys::exprType(node) } != pg_sys::TEXTOID { + continue; } - queries.push(unsafe { String::from_datum(value.constvalue, false)? }); + elements.push(unsafe { pg_sys::copyObjectImpl(node.cast()).cast() }); } - let combined = queries - .into_iter() - .map(|query| format!("({query})")) - .collect::>() - .join(" OR "); - let datum = combined.into_datum()?; - Some(unsafe { - pg_sys::makeConst( - pg_sys::TEXTOID, - -1, - pg_sys::DEFAULT_COLLATION_OID, - -1, - datum, - false, - false, - ) - .cast() - }) + if elements.is_empty() { + elements.push(unsafe { pg_sys::copyObjectImpl(first_query.cast()).cast() }); + } + let mut array = unsafe { PgBox::::alloc_node(pg_sys::NodeTag::T_ArrayExpr) }; + array.array_typeid = pg_sys::TEXTARRAYOID; + array.array_collid = pg_sys::DEFAULT_COLLATION_OID; + array.element_typeid = pg_sys::TEXTOID; + array.elements = elements.into_pg(); + array.multidims = false; + array.location = -1; + array.into_pg().cast() } unsafe fn make_int4_const(value: i32) -> *mut pg_sys::Const { @@ -743,7 +790,7 @@ unsafe fn lookup_score_bound() -> pg_sys::Oid { let names = unsafe { pg_sys::stringToQualifiedNameList(name.as_ptr(), std::ptr::null_mut()) }; let types = [ pg_sys::TEXTOID, - pg_sys::TEXTOID, + pg_sys::TEXTARRAYOID, pg_sys::INT4OID, pg_sys::INT4OID, pg_sys::INT4OID, @@ -753,7 +800,19 @@ unsafe fn lookup_score_bound() -> pg_sys::Oid { pg_sys::TEXTARRAYOID, pg_sys::TEXTARRAYOID, ]; - unsafe { pg_sys::LookupFuncName(names, types.len() as i32, types.as_ptr(), false) } + let oid = unsafe { pg_sys::LookupFuncName(names, types.len() as i32, types.as_ptr(), true) }; + if oid == pg_sys::InvalidOid { + pgrx::ereport!( + ERROR, + pg_sys::errcodes::PgSqlErrorCode::ERRCODE_UNDEFINED_FUNCTION, + "tin.score_bound() is missing: the installed tin SQL predates this build", + "Lead ships no extension upgrade scripts, so the extension has to be \ + reinstalled: DROP EXTENSION tin CASCADE; CREATE EXTENSION tin; \ + CASCADE also drops every tin index and anything else that depends on \ + them, so recreate those afterwards." + ); + } + oid } pgrx::extension_sql!( @@ -762,7 +821,7 @@ ALTER FUNCTION @extschema@.full_score(pg_catalog.tid) SUPPORT @extschema@.score_ ALTER FUNCTION @extschema@.full_score(pg_catalog.tid, pg_catalog.float4, pg_catalog.float4) SUPPORT @extschema@.score_support; ALTER FUNCTION @extschema@.score(pg_catalog.tid, pg_catalog.float4, pg_catalog.float4, pg_catalog.float4, pg_catalog.text[], pg_catalog.text[]) SUPPORT @extschema@.score_support; ALTER FUNCTION @extschema@.max_score(pg_catalog.tid) SUPPORT @extschema@.score_support; -REVOKE ALL ON FUNCTION @extschema@.score_bound(pg_catalog.text, pg_catalog.text, pg_catalog.int4, pg_catalog.int4, pg_catalog.int4, pg_catalog.float4, pg_catalog.float4, pg_catalog.float4, pg_catalog.text[], pg_catalog.text[]) FROM PUBLIC; +REVOKE ALL ON FUNCTION @extschema@.score_bound(pg_catalog.text, pg_catalog.text[], pg_catalog.int4, pg_catalog.int4, pg_catalog.int4, pg_catalog.float4, pg_catalog.float4, pg_catalog.float4, pg_catalog.text[], pg_catalog.text[]) FROM PUBLIC; "#, name = "score_support_bindings", requires = [