Skip to main content

hirn_engine/operators/
recall.rs

1//! Source operators: VectorRecall, HybridRecall, MultivectorRecall.
2//!
3//! These are data-source operators — they produce `RecordBatch`es from the
4//! store and ignore any input batches.
5
6use arrow_array::RecordBatch;
7use async_trait::async_trait;
8
9use hirn_core::error::HirnResult;
10use hirn_storage::store::{HybridSearchOptions, MultivectorSearchOptions, VectorSearchOptions};
11
12use super::{OpContext, Operator};
13
14// ── VectorRecall ────────────────────────────────────────────────────────
15
16/// Source operator that performs a vector similarity search.
17pub struct VectorRecall {
18    pub dataset: String,
19    pub opts: VectorSearchOptions,
20}
21
22#[async_trait]
23impl Operator for VectorRecall {
24    async fn execute(
25        &self,
26        _input: Vec<RecordBatch>,
27        ctx: &OpContext,
28    ) -> HirnResult<Vec<RecordBatch>> {
29        let batches = ctx
30            .store
31            .vector_search(&self.dataset, self.opts.clone())
32            .await
33            .map_err(|e| hirn_core::error::HirnError::storage(e))?;
34        Ok(batches)
35    }
36}
37
38// ── HybridRecall ────────────────────────────────────────────────────────
39
40/// Source operator that performs a hybrid (vector + FTS) search.
41pub struct HybridRecall {
42    pub dataset: String,
43    pub opts: HybridSearchOptions,
44}
45
46#[async_trait]
47impl Operator for HybridRecall {
48    async fn execute(
49        &self,
50        _input: Vec<RecordBatch>,
51        ctx: &OpContext,
52    ) -> HirnResult<Vec<RecordBatch>> {
53        let batches = ctx
54            .store
55            .hybrid_search(&self.dataset, self.opts.clone())
56            .await
57            .map_err(|e| hirn_core::error::HirnError::storage(e))?;
58        Ok(batches)
59    }
60}
61
62// ── MultivectorRecall ───────────────────────────────────────────────────
63
64/// Source operator that performs a multivector search (e.g., ColBERT MaxSim).
65pub struct MultivectorRecall {
66    pub dataset: String,
67    pub opts: MultivectorSearchOptions,
68}
69
70#[async_trait]
71impl Operator for MultivectorRecall {
72    async fn execute(
73        &self,
74        _input: Vec<RecordBatch>,
75        ctx: &OpContext,
76    ) -> HirnResult<Vec<RecordBatch>> {
77        let batches = ctx
78            .store
79            .multivector_search(&self.dataset, self.opts.clone())
80            .await
81            .map_err(|e| hirn_core::error::HirnError::storage(e))?;
82        Ok(batches)
83    }
84}