uni-query 2.0.2

OpenCypher query parser, planner, and vectorized executor for Uni
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

// SPDX-License-Identifier: Apache-2.0
// Copyright 2024-2026 Dragonscale Team

//! `uni.search` — hybrid (vector + FTS) search with RRF fusion.

use std::sync::Arc;
use std::sync::OnceLock;

use arrow_schema::DataType;
use datafusion::execution::SendableRecordBatchStream;
use datafusion::logical_expr::ColumnarValue;
use uni_plugin::traits::procedure::{
    NamedArgType, ProcedureContext, ProcedureMode, ProcedurePlugin, ProcedureSignature,
};
use uni_plugin::traits::scalar::ArgType;
use uni_plugin::{FnError, PluginError, PluginRegistrar, QName, SideEffects};

use crate::procedures_plugin::vector::{hybrid_search_yields, run_search_procedure};
use crate::query::df_graph::search_procedures::run_hybrid_search;

// Rust guideline compliant

fn signature() -> &'static ProcedureSignature {
    static SIG: OnceLock<ProcedureSignature> = OnceLock::new();
    SIG.get_or_init(|| ProcedureSignature {
        args: vec![
            NamedArgType {
                name: smol_str::SmolStr::new("label"),
                ty: ArgType::Primitive(DataType::Utf8),
                default: None,
                doc: "Vertex label to search.".to_owned(),
            },
            NamedArgType {
                name: smol_str::SmolStr::new("properties"),
                ty: ArgType::CypherValue,
                default: None,
                doc: "Either a property name (used for both vector and fts) or a map `{vector: '...', fts: '...'}`."
                    .to_owned(),
            },
            NamedArgType {
                name: smol_str::SmolStr::new("query_text"),
                ty: ArgType::Primitive(DataType::Utf8),
                default: None,
                doc: "Free-text query (used for FTS and, optionally, auto-embedding).".to_owned(),
            },
            NamedArgType {
                name: smol_str::SmolStr::new("query_vector"),
                ty: ArgType::CypherValue,
                default: None,
                doc: "Optional pre-computed query vector (List<Float>); omit to auto-embed.".to_owned(),
            },
            NamedArgType {
                name: smol_str::SmolStr::new("k"),
                ty: ArgType::Primitive(DataType::Int64),
                default: None,
                doc: "Number of fused results to return.".to_owned(),
            },
            NamedArgType {
                name: smol_str::SmolStr::new("filter"),
                ty: ArgType::Primitive(DataType::Utf8),
                default: None,
                doc: "Optional pushdown filter expression.".to_owned(),
            },
            NamedArgType {
                name: smol_str::SmolStr::new("options"),
                ty: ArgType::CypherValue,
                default: None,
                doc: "Optional options map (fusion method, alpha, rrf_k, reranker, …).".to_owned(),
            },
        ],
        yields: hybrid_search_yields(),
        mode: ProcedureMode::Read,
        side_effects: SideEffects::ReadOnly,
        retry_contract: None,
        batch_input: None,
        docs: "Hybrid vector + FTS search with RRF (or weighted) fusion and optional rerank."
            .to_owned(),
    })
}

#[derive(Debug)]
struct HybridSearchProc;

impl ProcedurePlugin for HybridSearchProc {
    fn signature(&self) -> &ProcedureSignature {
        signature()
    }

    fn invoke(
        &self,
        ctx: ProcedureContext<'_>,
        args: &[ColumnarValue],
    ) -> Result<SendableRecordBatchStream, FnError> {
        run_search_procedure(
            "uni.search",
            &ctx,
            args,
            signature(),
            |host, uni_args, yield_items, output_schema| async move {
                let target_properties = host.target_properties().clone();
                run_hybrid_search(
                    &host,
                    &uni_args,
                    &yield_items,
                    &target_properties,
                    &output_schema,
                )
                .await
            },
        )
    }
}

/// Register `uni.search` into `r`.
///
/// # Errors
///
/// Returns [`PluginError::DuplicateRegistration`] if a qname is taken.
pub fn register_into(r: &mut PluginRegistrar<'_>) -> Result<(), PluginError> {
    r.procedure(
        QName::new("uni", "search"),
        signature().clone(),
        Arc::new(HybridSearchProc),
    )?;
    Ok(())
}