qail_qdrant/
decoder.rs

1//! Zero-copy Protobuf decoder for Qdrant gRPC responses.
2//!
3//! Decodes protobuf wire format directly without intermediate allocations.
4//! Matches the zero-copy pattern of proto_encoder.rs.
5
6use crate::error::{QdrantError, QdrantResult};
7use crate::point::{PointId, ScoredPoint, Payload};
8
9// ============================================================================
10// Wire Type Constants
11// ============================================================================
12
13const WIRE_VARINT: u8 = 0;
14const WIRE_FIXED64: u8 = 1;
15const WIRE_LEN: u8 = 2;
16const WIRE_FIXED32: u8 = 5;
17
18// ============================================================================
19// SearchResponse Field Numbers
20// ============================================================================
21// message SearchResponse {
22//   repeated ScoredPoint result = 1;
23//   double time = 2;
24// }
25
26const SEARCH_RESULT: u32 = 1;
27#[allow(dead_code)]
28const SEARCH_TIME: u32 = 2;
29
30// ============================================================================
31// ScoredPoint Field Numbers
32// ============================================================================
33// message ScoredPoint {
34//   PointId id = 1;
35//   map<string, Value> payload = 2;
36//   float score = 3;
37//   uint64 version = 5;
38// }
39
40const SCORED_POINT_ID: u32 = 1;
41#[allow(dead_code)]
42const SCORED_POINT_PAYLOAD: u32 = 2;
43const SCORED_POINT_SCORE: u32 = 3;
44#[allow(dead_code)]
45const SCORED_POINT_VERSION: u32 = 5;
46
47// ============================================================================
48// PointId Field Numbers
49// ============================================================================
50// message PointId {
51//   oneof point_id_options {
52//     uint64 num = 1;
53//     string uuid = 2;
54//   }
55// }
56
57const POINT_ID_NUM: u32 = 1;
58const POINT_ID_UUID: u32 = 2;
59
60// ============================================================================
61// Varint Decoding
62// ============================================================================
63
64/// Decode a varint from the buffer, advancing the cursor.
65#[inline]
66fn decode_varint(buf: &mut &[u8]) -> QdrantResult<u64> {
67    let mut result: u64 = 0;
68    let mut shift = 0;
69    
70    loop {
71        if buf.is_empty() {
72            return Err(QdrantError::Decode("Unexpected end of data in varint".to_string()));
73        }
74        
75        let byte = buf[0];
76        *buf = &buf[1..];
77        
78        result |= ((byte & 0x7F) as u64) << shift;
79        
80        if byte & 0x80 == 0 {
81            return Ok(result);
82        }
83        
84        shift += 7;
85        if shift >= 64 {
86            return Err(QdrantError::Decode("Varint too long".to_string()));
87        }
88    }
89}
90
91/// Decode a field tag (field_number << 3 | wire_type).
92#[inline]
93fn decode_tag(buf: &mut &[u8]) -> QdrantResult<(u32, u8)> {
94    let tag = decode_varint(buf)?;
95    let field_number = (tag >> 3) as u32;
96    let wire_type = (tag & 0x07) as u8;
97    Ok((field_number, wire_type))
98}
99
100/// Skip a field value based on wire type.
101#[inline]
102fn skip_field(buf: &mut &[u8], wire_type: u8) -> QdrantResult<()> {
103    match wire_type {
104        WIRE_VARINT => {
105            decode_varint(buf)?;
106        }
107        WIRE_FIXED64 => {
108            if buf.len() < 8 {
109                return Err(QdrantError::Decode("Unexpected end of data".to_string()));
110            }
111            *buf = &buf[8..];
112        }
113        WIRE_LEN => {
114            let len = decode_varint(buf)? as usize;
115            if buf.len() < len {
116                return Err(QdrantError::Decode("Unexpected end of data".to_string()));
117            }
118            *buf = &buf[len..];
119        }
120        WIRE_FIXED32 => {
121            if buf.len() < 4 {
122                return Err(QdrantError::Decode("Unexpected end of data".to_string()));
123            }
124            *buf = &buf[4..];
125        }
126        _ => {
127            return Err(QdrantError::Decode(format!("Unknown wire type: {}", wire_type)));
128        }
129    }
130    Ok(())
131}
132
133// ============================================================================
134// SearchResponse Decoder
135// ============================================================================
136
137/// Decode a SearchResponse protobuf message.
138///
139/// # Zero-Copy Pattern
140/// - Parses in a single pass through the buffer
141/// - Minimal allocations (only for result Vec and PointId strings)
142/// - No intermediate struct copies
143pub fn decode_search_response(data: &[u8]) -> QdrantResult<Vec<ScoredPoint>> {
144    let mut results = Vec::new();
145    let mut buf = data;
146    
147    while !buf.is_empty() {
148        let (field_number, wire_type) = decode_tag(&mut buf)?;
149        
150        match field_number {
151            SEARCH_RESULT => {
152                // repeated ScoredPoint - length-delimited
153                if wire_type != WIRE_LEN {
154                    return Err(QdrantError::Decode("Expected length-delimited for ScoredPoint".to_string()));
155                }
156                
157                let len = decode_varint(&mut buf)? as usize;
158                if buf.len() < len {
159                    return Err(QdrantError::Decode("Truncated ScoredPoint".to_string()));
160                }
161                
162                let point_data = &buf[..len];
163                buf = &buf[len..];
164                
165                let point = decode_scored_point(point_data)?;
166                results.push(point);
167            }
168            _ => {
169                // Skip unknown fields (including time, usage)
170                skip_field(&mut buf, wire_type)?;
171            }
172        }
173    }
174    
175    Ok(results)
176}
177
178/// Decode a single ScoredPoint message.
179fn decode_scored_point(data: &[u8]) -> QdrantResult<ScoredPoint> {
180    let mut id = PointId::Num(0);
181    let mut score = 0.0f32;
182    let mut buf = data;
183    
184    while !buf.is_empty() {
185        let (field_number, wire_type) = decode_tag(&mut buf)?;
186        
187        match field_number {
188            SCORED_POINT_ID => {
189                // PointId - length-delimited submessage
190                if wire_type != WIRE_LEN {
191                    skip_field(&mut buf, wire_type)?;
192                    continue;
193                }
194                
195                let len = decode_varint(&mut buf)? as usize;
196                if buf.len() < len {
197                    return Err(QdrantError::Decode("Truncated PointId".to_string()));
198                }
199                
200                let id_data = &buf[..len];
201                buf = &buf[len..];
202                
203                id = decode_point_id(id_data)?;
204            }
205            SCORED_POINT_SCORE => {
206                // float - fixed32
207                if wire_type != WIRE_FIXED32 {
208                    skip_field(&mut buf, wire_type)?;
209                    continue;
210                }
211                
212                if buf.len() < 4 {
213                    return Err(QdrantError::Decode("Truncated score".to_string()));
214                }
215                
216                let bytes = [buf[0], buf[1], buf[2], buf[3]];
217                score = f32::from_le_bytes(bytes);
218                buf = &buf[4..];
219            }
220            _ => {
221                // Skip payload, version, vectors, etc.
222                skip_field(&mut buf, wire_type)?;
223            }
224        }
225    }
226    
227    Ok(ScoredPoint {
228        id,
229        score,
230        payload: Payload::new(),
231        vector: None,
232    })
233}
234
235/// Decode a PointId message.
236fn decode_point_id(data: &[u8]) -> QdrantResult<PointId> {
237    let mut buf = data;
238    
239    while !buf.is_empty() {
240        let (field_number, wire_type) = decode_tag(&mut buf)?;
241        
242        match field_number {
243            POINT_ID_NUM => {
244                // uint64
245                if wire_type != WIRE_VARINT {
246                    skip_field(&mut buf, wire_type)?;
247                    continue;
248                }
249                let num = decode_varint(&mut buf)?;
250                return Ok(PointId::Num(num));
251            }
252            POINT_ID_UUID => {
253                // string
254                if wire_type != WIRE_LEN {
255                    skip_field(&mut buf, wire_type)?;
256                    continue;
257                }
258                let len = decode_varint(&mut buf)? as usize;
259                if buf.len() < len {
260                    return Err(QdrantError::Decode("Truncated UUID".to_string()));
261                }
262                
263                let uuid_str = std::str::from_utf8(&buf[..len])
264                    .map_err(|e| QdrantError::Decode(format!("Invalid UTF-8: {}", e)))?;
265                return Ok(PointId::Uuid(uuid_str.to_string()));
266            }
267            _ => {
268                skip_field(&mut buf, wire_type)?;
269            }
270        }
271    }
272    
273    // Default to Num(0) if empty
274    Ok(PointId::Num(0))
275}
276
277// ============================================================================
278// Tests
279// ============================================================================
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    #[test]
286    fn test_decode_varint() {
287        // Single byte
288        let mut buf: &[u8] = &[0x01];
289        assert_eq!(decode_varint(&mut buf).unwrap(), 1);
290        assert!(buf.is_empty());
291        
292        // Two bytes (300 = 0xAC 0x02)
293        let mut buf: &[u8] = &[0xAC, 0x02];
294        assert_eq!(decode_varint(&mut buf).unwrap(), 300);
295        assert!(buf.is_empty());
296    }
297
298    #[test]
299    fn test_decode_tag() {
300        // Field 1, wire type 2 = 0x0A
301        let mut buf: &[u8] = &[0x0A];
302        let (field, wire) = decode_tag(&mut buf).unwrap();
303        assert_eq!(field, 1);
304        assert_eq!(wire, WIRE_LEN);
305        
306        // Field 3, wire type 5 = (3 << 3) | 5 = 0x1D
307        let mut buf: &[u8] = &[0x1D];
308        let (field, wire) = decode_tag(&mut buf).unwrap();
309        assert_eq!(field, 3);
310        assert_eq!(wire, WIRE_FIXED32);
311    }
312
313    #[test]
314    fn test_decode_point_id_num() {
315        // PointId { num = 42 }
316        // Field 1 (num), varint = 0x08, value 42 = 0x2A
317        let data = &[0x08, 0x2A];
318        let id = decode_point_id(data).unwrap();
319        assert_eq!(id, PointId::Num(42));
320    }
321
322    #[test]
323    fn test_decode_point_id_uuid() {
324        // PointId { uuid = "abc" }
325        // Field 2 (uuid), len-delimited = 0x12, len 3, "abc"
326        let data = &[0x12, 0x03, b'a', b'b', b'c'];
327        let id = decode_point_id(data).unwrap();
328        assert_eq!(id, PointId::Uuid("abc".to_string()));
329    }
330
331    #[test]
332    fn test_decode_scored_point() {
333        // ScoredPoint { id: PointId { num: 1 }, score: 0.5 }
334        // Field 1 (id): 0x0A, len 2, [0x08, 0x01] (num = 1)
335        // Field 3 (score): 0x1D, f32 bytes for 0.5
336        let score_bytes = 0.5f32.to_le_bytes();
337        let data = &[
338            0x0A, 0x02, 0x08, 0x01,  // id = PointId { num = 1 }
339            0x1D, score_bytes[0], score_bytes[1], score_bytes[2], score_bytes[3],  // score = 0.5
340        ];
341        
342        let point = decode_scored_point(data).unwrap();
343        assert_eq!(point.id, PointId::Num(1));
344        assert!((point.score - 0.5).abs() < 0.0001);
345    }
346
347    #[test]
348    fn test_decode_search_response_empty() {
349        let data: &[u8] = &[];
350        let results = decode_search_response(data).unwrap();
351        assert!(results.is_empty());
352    }
353}