ailake-query 0.0.11

Query planner and executor for AI-Lake — vector search, compaction, ContextAssembler
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

// SPDX-License-Identifier: MIT OR Apache-2.0
use ailake_catalog::{decode_centroid, DataFileEntry};
use ailake_core::VectorMetric;
use ailake_vec::{cosine_distance, dot_product, euclidean_distance};
use tracing::debug;

pub struct VectorPruner;

impl VectorPruner {
    /// Remove files whose centroid is geometrically guaranteed to contain no vectors
    /// within `threshold` distance of `query`.
    ///
    /// Pruning condition: `distance(query, centroid) - radius > threshold`
    /// Files without centroid metadata are kept (conservative fallback).
    pub fn prune(
        files: Vec<DataFileEntry>,
        query: &[f32],
        metric: VectorMetric,
        threshold: f32,
    ) -> Vec<DataFileEntry> {
        files
            .into_iter()
            .filter(|entry| {
                match decode_centroid(entry, metric) {
                    Some(centroid) => {
                        let dist = compute_distance(query, &centroid.values, metric);
                        let keep = dist - centroid.radius <= threshold;
                        debug!(
                            "ailake: pruner {} — dist={:.4} radius={:.4} edge={:.4} threshold={:.4} → {}",
                            entry.path,
                            dist,
                            centroid.radius,
                            dist - centroid.radius,
                            threshold,
                            if keep { "KEEP" } else { "PRUNE" }
                        );
                        keep
                    }
                    None => {
                        debug!(
                            "ailake: pruner {} — no centroid metadata, keeping (conservative fallback)",
                            entry.path
                        );
                        true // no centroid → keep (safe fallback)
                    }
                }
            })
            .collect()
    }
}

fn compute_distance(a: &[f32], b: &[f32], metric: VectorMetric) -> f32 {
    match metric {
        VectorMetric::Cosine | VectorMetric::NormalizedCosine => cosine_distance(a, b),
        VectorMetric::Euclidean => euclidean_distance(a, b),
        VectorMetric::DotProduct => -dot_product(a, b),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ailake_catalog::{make_data_file_entry, VectorIndexInfo};
    use ailake_core::VectorMetric;
    use ailake_vec::compute_centroid_and_radius;

    fn make_entry(path: &str, vecs: &[Vec<f32>], metric: VectorMetric) -> DataFileEntry {
        let centroid = compute_centroid_and_radius(vecs, metric);
        make_data_file_entry(
            path,
            vecs.len() as u64,
            1024,
            &centroid,
            VectorIndexInfo {
                column: "embedding",
                dim: vecs[0].len() as u32,
                hnsw_offset: 0,
                hnsw_len: 0,
            },
        )
    }

    #[test]
    fn prunes_far_file() {
        // File centroid near [1,0,0], query near [0,0,1] — orthogonal → prune
        let vecs = vec![vec![1.0f32, 0.0, 0.0], vec![0.9, 0.1, 0.0]];
        let entry = make_entry("far.parquet", &vecs, VectorMetric::Cosine);
        let query = vec![0.0f32, 0.0, 1.0];
        let pruned = VectorPruner::prune(vec![entry], &query, VectorMetric::Cosine, 0.1);
        assert!(pruned.is_empty(), "far file should be pruned");
    }

    #[test]
    fn keeps_nearby_file() {
        let vecs = vec![vec![1.0f32, 0.0, 0.0], vec![0.99, 0.1, 0.0]];
        let entry = make_entry("near.parquet", &vecs, VectorMetric::Cosine);
        let query = vec![1.0f32, 0.0, 0.0];
        let kept = VectorPruner::prune(vec![entry], &query, VectorMetric::Cosine, 0.5);
        assert_eq!(kept.len(), 1, "nearby file should be kept");
    }

    #[test]
    fn no_centroid_always_kept() {
        let entry = DataFileEntry {
            path: "unknown.parquet".into(),
            record_count: 10,
            file_size_bytes: 512,
            centroid_b64: None,
            radius: None,
            hnsw_offset: None,
            hnsw_len: None,
            vector_column: None,
            vector_dim: None,
            extra_vector_indexes: vec![],
            index_status: ailake_catalog::IndexStatus::Ready,
            batch_id: None,
        };
        let query = vec![0.0f32, 0.0, 1.0];
        let kept = VectorPruner::prune(vec![entry], &query, VectorMetric::Cosine, 0.0);
        assert_eq!(kept.len(), 1);
    }
}