fix: max_score when used with full_score

.github/workflows/ci.yml

@@ -18,9 +18,9 @@ jobs:
       matrix:
         arch:
           - name: x86-64
-            runner: depot-ubuntu-24.04-8
+            runner: ubuntu-24.04
           - name: arm64
-            runner: depot-ubuntu-24.04-arm-8
+            runner: ubuntu-24.04-arm
         postgres: [17, 18]
 
     env:

postgres/src/lib.rs

@@ -250,6 +250,38 @@ mod tests {
         assert_eq!(inspected, Some(vec!["rare".to_owned()]));
     }
 
+    #[pg_test]
+    fn full_score_normalization_matches_tin() {
+        Spi::run(
+            "CREATE TABLE lite_normalization (id int, body text);
+             INSERT INTO lite_normalization VALUES
+               (1, 'I love fuji apples and juicy mangoes'),
+               (2, 'Grape tasting notes from the orchard'),
+               (3, 'The best juicy fuji apple in town');
+             CREATE INDEX lite_normalization_idx ON lite_normalization USING tin (body)",
+        )
+        .unwrap();
+        for expression in [
+            "tin.full_score(ctid) / tin.max_score(ctid)",
+            "1::real / tin.max_score(ctid) * tin.full_score(ctid)",
+        ] {
+            let sql = format!(
+                "SELECT {expression} FROM lite_normalization
+                 WHERE body ==> 'apple OR grape' AND tin.max_score(ctid) > 0 ORDER BY id"
+            );
+            let scores = Spi::connect(|client| {
+                client
+                    .select(&sql, None, &[])
+                    .unwrap()
+                    .map(|row| row.get::<f32>(1).unwrap().unwrap())
+                    .collect::<Vec<_>>()
+            });
+            assert_eq!(scores.len(), 2);
+            assert!((scores[0] - 1.0).abs() < 0.000001);
+            assert!((scores[1] - 0.9398665).abs() < 0.000001);
+        }
+    }
+
     #[pg_test]
     fn scoring_binds_to_expression_indexes() {
         Spi::run(

postgres/src/score.rs

@@ -104,7 +104,7 @@ fn score_bound(
         heap_oid: heap_oid as u32,
         index_oid: index_oid as u32,
         query: query.to_owned(),
-        full: mode == 1,
+        full: mode == 1 || mode == 3,
         dense: dense_ratio.unwrap_or(DenseRatio::DEFAULT).to_bits(),
         k1: bits(k1),
         b: bits(b),
@@ -116,7 +116,7 @@ fn score_bound(
             *slot = Some(build_corpus(key.clone(), k1, b, term_add, term_replace));
         }
         let corpus = slot.as_ref().expect("score corpus was just populated");
-        if mode == 2 {
+        if mode == 2 || mode == 3 {
             corpus.max
         } else {
             corpus.by_document.get(document).copied().unwrap_or(0.0)
@@ -475,6 +475,57 @@ pub(crate) unsafe fn find_matching_tin_index(
     matched
 }
 
+struct FullScoreBinding {
+    ctid: *const pg_sys::Var,
+    document: *mut pg_sys::Node,
+    support: pg_sys::Oid,
+    bound: pg_sys::Oid,
+}
+
+#[pg_guard]
+unsafe extern "C-unwind" fn has_full_score(node: *mut pg_sys::Node, context: *mut c_void) -> bool {
+    unsafe {
+        if node.is_null() || (*node).type_ == pg_sys::NodeTag::T_Query {
+            return false;
+        }
+        let binding = &*context.cast::<FullScoreBinding>();
+        if (*node).type_ == pg_sys::NodeTag::T_FuncExpr {
+            let function = &*node.cast::<pg_sys::FuncExpr>();
+            // Earlier query clauses may already contain the rewritten scorer.
+            if function.funcid == binding.bound {
+                let mode = pg_sys::list_nth(function.args, 4).cast::<pg_sys::Const>();
+                if (*mode).xpr.type_ == pg_sys::NodeTag::T_Const
+                    && (*mode).constvalue.value() == 1
+                    && pg_sys::equal(pg_sys::list_nth(function.args, 0), binding.document.cast())
+                {
+                    return true;
+                }
+            } else if pg_sys::get_func_support(function.funcid) == binding.support
+                && CStr::from_ptr(pg_sys::get_func_name(function.funcid)).to_bytes()
+                    == b"full_score"
+            {
+                for position in 0..pg_sys::list_length(function.args) {
+                    let mut argument =
+                        pg_sys::list_nth(function.args, position).cast::<pg_sys::Node>();
+                    if (*argument).type_ == pg_sys::NodeTag::T_NamedArgExpr {
+                        let named = &*argument.cast::<pg_sys::NamedArgExpr>();
+                        if named.argnumber != 0 {
+                            continue;
+                        }
+                        argument = named.arg.cast();
+                    } else if position != 0 {
+                        continue;
+                    }
+                    if pg_sys::equal(argument.cast(), binding.ctid.cast()) {
+                        return true;
+                    }
+                }
+            }
+        }
+        pg_sys::expression_tree_walker(node, Some(has_full_score), context)
+    }
+}
+
 #[pg_extern(immutable, parallel_unsafe)]
 fn score_support(request: Internal) -> Internal {
     let unhandled = || Internal::from(Some(pg_sys::Datum::from(0_usize)));
@@ -523,7 +574,19 @@ fn score_support(request: Internal) -> Internal {
         let mode = if fname.as_ref() == "full_score" {
             1
         } else if fname.as_ref() == "max_score" {
-            2
+            let mut binding = FullScoreBinding {
+                ctid,
+                document,
+                support: pg_sys::get_func_support((*request.fcall).funcid),
+                bound: lookup_score_bound(),
+            };
+            let full = pg_sys::query_tree_walker(
+                parse,
+                Some(has_full_score),
+                (&mut binding as *mut FullScoreBinding).cast(),
+                pg_sys::QTW_IGNORE_RC_SUBQUERIES as i32,
+            );
+            if full { 3 } else { 2 }
         } else {
             0
         };