use bytes::{BufMut, BytesMut};
#[allow(dead_code)]
const WIRE_VARINT: u8 = 0;
#[allow(dead_code)]
const WIRE_LEN: u8 = 2;
#[allow(dead_code)]
const WIRE_FIXED32: u8 = 5;
const SEARCH_COLLECTION: u8 = 0x0A;
const SEARCH_VECTOR: u8 = 0x12;
const SEARCH_FILTER: u8 = 0x1A;
const SEARCH_LIMIT: u8 = 0x20;
const SEARCH_WITH_PAYLOAD: u8 = 0x32;
const SEARCH_SCORE_THRESHOLD: u8 = 0x45;
const SEARCH_VECTOR_NAME: u8 = 0x52;
const UPSERT_COLLECTION: u8 = 0x0A;
const UPSERT_WAIT: u8 = 0x10;
const UPSERT_POINTS: u8 = 0x1A;
const POINT_ID: u8 = 0x0A;
const POINT_VECTORS: u8 = 0x22;
const POINT_PAYLOAD: u8 = 0x1A;
const POINT_ID_NUM: u8 = 0x08;
const POINT_ID_UUID: u8 = 0x12;
const FILTER_MUST: u8 = 0x0A;
const FILTER_SHOULD: u8 = 0x12;
#[allow(dead_code)]
const FILTER_MUST_NOT: u8 = 0x1A;
#[inline]
pub fn encode_varint(buf: &mut BytesMut, mut value: usize) {
loop {
let byte = (value & 0x7F) as u8;
value >>= 7;
if value == 0 {
buf.put_u8(byte);
break;
} else {
buf.put_u8(byte | 0x80);
}
}
}
#[inline]
pub fn encode_varint_u64(buf: &mut BytesMut, mut value: u64) {
loop {
let byte = (value & 0x7F) as u8;
value >>= 7;
if value == 0 {
buf.put_u8(byte);
break;
} else {
buf.put_u8(byte | 0x80);
}
}
}
pub fn encode_search_proto(
buf: &mut BytesMut,
collection: &str,
vector: &[f32],
limit: u64,
score_threshold: Option<f32>,
vector_name: Option<&str>,
) {
buf.clear();
buf.put_u8(SEARCH_COLLECTION);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
buf.put_u8(SEARCH_VECTOR);
let vector_bytes_len = vector.len() * 4; encode_varint(buf, vector_bytes_len);
let float_bytes: &[u8] = bytemuck::cast_slice(vector);
buf.extend_from_slice(float_bytes);
buf.put_u8(SEARCH_LIMIT);
encode_varint_u64(buf, limit);
encode_with_payload_true(buf);
if let Some(threshold) = score_threshold {
buf.put_u8(SEARCH_SCORE_THRESHOLD);
buf.put_f32_le(threshold);
}
if let Some(name) = vector_name {
buf.put_u8(SEARCH_VECTOR_NAME);
encode_varint(buf, name.len());
buf.extend_from_slice(name.as_bytes());
}
}
#[allow(clippy::too_many_arguments)]
pub fn encode_search_with_filter_proto(
buf: &mut BytesMut,
collection: &str,
vector: &[f32],
limit: u64,
score_threshold: Option<f32>,
vector_name: Option<&str>,
conditions: &[qail_core::ast::Condition],
is_or: bool,
) {
buf.clear();
buf.put_u8(SEARCH_COLLECTION);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
buf.put_u8(SEARCH_VECTOR);
let vector_bytes_len = vector.len() * 4;
encode_varint(buf, vector_bytes_len);
let float_bytes: &[u8] = bytemuck::cast_slice(vector);
buf.extend_from_slice(float_bytes);
if !conditions.is_empty() {
let filter_buf = encode_filter_message(conditions, is_or);
buf.put_u8(SEARCH_FILTER);
encode_varint(buf, filter_buf.len());
buf.extend_from_slice(&filter_buf);
}
buf.put_u8(SEARCH_LIMIT);
encode_varint_u64(buf, limit);
encode_with_payload_true(buf);
if let Some(threshold) = score_threshold {
buf.put_u8(SEARCH_SCORE_THRESHOLD);
buf.put_f32_le(threshold);
}
if let Some(name) = vector_name {
buf.put_u8(SEARCH_VECTOR_NAME);
encode_varint(buf, name.len());
buf.extend_from_slice(name.as_bytes());
}
}
pub fn encode_with_payload_true(buf: &mut BytesMut) {
buf.put_u8(SEARCH_WITH_PAYLOAD);
encode_varint(buf, 2); buf.put_u8(0x08); buf.put_u8(0x01); }
fn encode_filter_message(conditions: &[qail_core::ast::Condition], is_or: bool) -> BytesMut {
use qail_core::ast::{Expr, Operator, Value};
let mut filter_buf = BytesMut::with_capacity(conditions.len() * 32);
let clause_tag = if is_or { FILTER_SHOULD } else { FILTER_MUST };
for cond in conditions {
let key = match &cond.left {
Expr::Named(name) => name.as_str(),
Expr::Aliased { name, .. } => name.as_str(),
_ => continue,
};
let cond_buf = match (&cond.op, &cond.value) {
(Operator::Eq, Value::String(s)) => {
encode_field_condition_match_keyword(key, s)
}
(Operator::Eq, Value::Int(n)) => {
encode_field_condition_match_integer(key, *n)
}
(Operator::Eq, Value::Bool(b)) => {
encode_field_condition_match_bool(key, *b)
}
(Operator::Gt, Value::Int(n)) => encode_field_condition_range(key, None, None, Some(*n as f64), None),
(Operator::Gt, Value::Float(f)) => encode_field_condition_range(key, None, None, Some(*f), None),
(Operator::Gte, Value::Int(n)) => encode_field_condition_range(key, None, None, None, Some(*n as f64)),
(Operator::Gte, Value::Float(f)) => encode_field_condition_range(key, None, None, None, Some(*f)),
(Operator::Lt, Value::Int(n)) => encode_field_condition_range(key, Some(*n as f64), None, None, None),
(Operator::Lt, Value::Float(f)) => encode_field_condition_range(key, Some(*f), None, None, None),
(Operator::Lte, Value::Int(n)) => encode_field_condition_range(key, None, Some(*n as f64), None, None),
(Operator::Lte, Value::Float(f)) => encode_field_condition_range(key, None, Some(*f), None, None),
(Operator::Contains | Operator::Like, Value::String(s)) => {
encode_field_condition_match_text(key, s)
}
_ => continue,
};
filter_buf.put_u8(clause_tag);
encode_varint(&mut filter_buf, cond_buf.len());
filter_buf.extend_from_slice(&cond_buf);
}
filter_buf
}
fn encode_field_condition_match_keyword(key: &str, value: &str) -> BytesMut {
let mut match_buf = BytesMut::with_capacity(value.len() + 8);
match_buf.put_u8(0x0A); encode_varint(&mut match_buf, value.len());
match_buf.extend_from_slice(value.as_bytes());
let mut fc_buf = BytesMut::with_capacity(key.len() + match_buf.len() + 16);
fc_buf.put_u8(0x0A);
encode_varint(&mut fc_buf, key.len());
fc_buf.extend_from_slice(key.as_bytes());
fc_buf.put_u8(0x12);
encode_varint(&mut fc_buf, match_buf.len());
fc_buf.extend_from_slice(&match_buf);
let mut cond_buf = BytesMut::with_capacity(fc_buf.len() + 4);
cond_buf.put_u8(0x0A); encode_varint(&mut cond_buf, fc_buf.len());
cond_buf.extend_from_slice(&fc_buf);
cond_buf
}
fn encode_field_condition_match_integer(key: &str, value: i64) -> BytesMut {
let mut match_buf = BytesMut::with_capacity(16);
match_buf.put_u8(0x10); encode_varint_u64(&mut match_buf, value as u64);
let mut fc_buf = BytesMut::with_capacity(key.len() + match_buf.len() + 16);
fc_buf.put_u8(0x0A);
encode_varint(&mut fc_buf, key.len());
fc_buf.extend_from_slice(key.as_bytes());
fc_buf.put_u8(0x12);
encode_varint(&mut fc_buf, match_buf.len());
fc_buf.extend_from_slice(&match_buf);
let mut cond_buf = BytesMut::with_capacity(fc_buf.len() + 4);
cond_buf.put_u8(0x0A);
encode_varint(&mut cond_buf, fc_buf.len());
cond_buf.extend_from_slice(&fc_buf);
cond_buf
}
fn encode_field_condition_match_bool(key: &str, value: bool) -> BytesMut {
let mut match_buf = BytesMut::with_capacity(4);
match_buf.put_u8(0x20); match_buf.put_u8(if value { 1 } else { 0 });
let mut fc_buf = BytesMut::with_capacity(key.len() + match_buf.len() + 16);
fc_buf.put_u8(0x0A);
encode_varint(&mut fc_buf, key.len());
fc_buf.extend_from_slice(key.as_bytes());
fc_buf.put_u8(0x12);
encode_varint(&mut fc_buf, match_buf.len());
fc_buf.extend_from_slice(&match_buf);
let mut cond_buf = BytesMut::with_capacity(fc_buf.len() + 4);
cond_buf.put_u8(0x0A);
encode_varint(&mut cond_buf, fc_buf.len());
cond_buf.extend_from_slice(&fc_buf);
cond_buf
}
fn encode_field_condition_match_text(key: &str, value: &str) -> BytesMut {
let mut match_buf = BytesMut::with_capacity(value.len() + 8);
match_buf.put_u8(0x1A); encode_varint(&mut match_buf, value.len());
match_buf.extend_from_slice(value.as_bytes());
let mut fc_buf = BytesMut::with_capacity(key.len() + match_buf.len() + 16);
fc_buf.put_u8(0x0A);
encode_varint(&mut fc_buf, key.len());
fc_buf.extend_from_slice(key.as_bytes());
fc_buf.put_u8(0x12);
encode_varint(&mut fc_buf, match_buf.len());
fc_buf.extend_from_slice(&match_buf);
let mut cond_buf = BytesMut::with_capacity(fc_buf.len() + 4);
cond_buf.put_u8(0x0A);
encode_varint(&mut cond_buf, fc_buf.len());
cond_buf.extend_from_slice(&fc_buf);
cond_buf
}
fn encode_field_condition_range(
key: &str,
lt: Option<f64>,
lte: Option<f64>,
gt: Option<f64>,
gte: Option<f64>,
) -> BytesMut {
let mut range_buf = BytesMut::with_capacity(40);
if let Some(v) = lt {
range_buf.put_u8(0x09); range_buf.put_f64_le(v);
}
if let Some(v) = gt {
range_buf.put_u8(0x11); range_buf.put_f64_le(v);
}
if let Some(v) = gte {
range_buf.put_u8(0x19); range_buf.put_f64_le(v);
}
if let Some(v) = lte {
range_buf.put_u8(0x21); range_buf.put_f64_le(v);
}
let mut fc_buf = BytesMut::with_capacity(key.len() + range_buf.len() + 16);
fc_buf.put_u8(0x0A); encode_varint(&mut fc_buf, key.len());
fc_buf.extend_from_slice(key.as_bytes());
fc_buf.put_u8(0x1A); encode_varint(&mut fc_buf, range_buf.len());
fc_buf.extend_from_slice(&range_buf);
let mut cond_buf = BytesMut::with_capacity(fc_buf.len() + 4);
cond_buf.put_u8(0x0A);
encode_varint(&mut cond_buf, fc_buf.len());
cond_buf.extend_from_slice(&fc_buf);
cond_buf
}
pub fn encode_upsert_proto(
buf: &mut BytesMut,
collection: &str,
points: &[crate::Point],
wait: bool,
) {
buf.clear();
buf.put_u8(UPSERT_COLLECTION);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
if wait {
buf.put_u8(UPSERT_WAIT);
buf.put_u8(0x01);
}
for point in points {
encode_point_struct(buf, point);
}
}
fn encode_point_struct(buf: &mut BytesMut, point: &crate::Point) {
let mut point_buf = BytesMut::with_capacity(point.vector.len() * 4 + 64);
encode_point_id_field(&mut point_buf, &point.id);
if !point.payload.is_empty() {
encode_payload_map(&mut point_buf, &point.payload);
}
let vector_bytes_len = point.vector.len() * 4;
let vector_inner_len = 1 + varint_len(vector_bytes_len as u64) + vector_bytes_len;
let vectors_len = 1 + varint_len(vector_inner_len as u64) + vector_inner_len;
point_buf.put_u8(POINT_VECTORS);
encode_varint(&mut point_buf, vectors_len);
point_buf.put_u8(0x0A); encode_varint(&mut point_buf, vector_inner_len);
point_buf.put_u8(0x0A); encode_varint(&mut point_buf, vector_bytes_len);
let float_bytes: &[u8] = bytemuck::cast_slice(&point.vector);
point_buf.extend_from_slice(float_bytes);
buf.put_u8(UPSERT_POINTS);
encode_varint(buf, point_buf.len());
buf.extend_from_slice(&point_buf);
}
fn encode_point_id_field(buf: &mut BytesMut, id: &crate::PointId) {
match id {
crate::PointId::Num(n) => {
buf.put_u8(POINT_ID);
let id_len = 1 + varint_len(*n);
encode_varint(buf, id_len);
buf.put_u8(POINT_ID_NUM);
encode_varint_u64(buf, *n);
}
crate::PointId::Uuid(s) => {
buf.put_u8(POINT_ID);
let id_len = 1 + varint_len(s.len() as u64) + s.len();
encode_varint(buf, id_len);
buf.put_u8(POINT_ID_UUID);
encode_varint(buf, s.len());
buf.extend_from_slice(s.as_bytes());
}
}
}
fn encode_payload_map(buf: &mut BytesMut, payload: &crate::point::Payload) {
for (key, value) in payload {
let mut entry_buf = BytesMut::with_capacity(key.len() + 32);
entry_buf.put_u8(0x0A);
encode_varint(&mut entry_buf, key.len());
entry_buf.extend_from_slice(key.as_bytes());
let value_buf = encode_payload_value(value);
entry_buf.put_u8(0x12);
encode_varint(&mut entry_buf, value_buf.len());
entry_buf.extend_from_slice(&value_buf);
buf.put_u8(POINT_PAYLOAD);
encode_varint(buf, entry_buf.len());
buf.extend_from_slice(&entry_buf);
}
}
fn encode_payload_value(value: &crate::point::PayloadValue) -> BytesMut {
use crate::point::PayloadValue;
let mut buf = BytesMut::with_capacity(32);
match value {
PayloadValue::Null => {
buf.put_u8(0x08);
buf.put_u8(0x00);
}
PayloadValue::Float(f) => {
buf.put_u8(0x11); buf.put_f64_le(*f);
}
PayloadValue::Integer(n) => {
buf.put_u8(0x18); encode_varint_u64(&mut buf, *n as u64);
}
PayloadValue::String(s) => {
buf.put_u8(0x22); encode_varint(&mut buf, s.len());
buf.extend_from_slice(s.as_bytes());
}
PayloadValue::Bool(b) => {
buf.put_u8(0x28); buf.put_u8(if *b { 1 } else { 0 });
}
PayloadValue::List(items) => {
let mut list_buf = BytesMut::with_capacity(items.len() * 16);
for item in items {
let val_buf = encode_payload_value(item);
list_buf.put_u8(0x0A);
encode_varint(&mut list_buf, val_buf.len());
list_buf.extend_from_slice(&val_buf);
}
buf.put_u8(0x3A); encode_varint(&mut buf, list_buf.len());
buf.extend_from_slice(&list_buf);
}
PayloadValue::Object(map) => {
let mut struct_buf = BytesMut::with_capacity(map.len() * 32);
for (k, v) in map {
let val_buf = encode_payload_value(v);
let mut entry_buf = BytesMut::with_capacity(k.len() + val_buf.len() + 8);
entry_buf.put_u8(0x0A);
encode_varint(&mut entry_buf, k.len());
entry_buf.extend_from_slice(k.as_bytes());
entry_buf.put_u8(0x12);
encode_varint(&mut entry_buf, val_buf.len());
entry_buf.extend_from_slice(&val_buf);
struct_buf.put_u8(0x0A);
encode_varint(&mut struct_buf, entry_buf.len());
struct_buf.extend_from_slice(&entry_buf);
}
buf.put_u8(0x32); encode_varint(&mut buf, struct_buf.len());
buf.extend_from_slice(&struct_buf);
}
}
buf
}
pub fn encode_get_points_proto(
buf: &mut BytesMut,
collection: &str,
ids: &[crate::PointId],
with_vectors: bool,
) {
buf.clear();
buf.put_u8(0x0A);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
for id in ids {
let mut id_buf = BytesMut::with_capacity(40);
match id {
crate::PointId::Num(n) => {
id_buf.put_u8(POINT_ID_NUM);
encode_varint_u64(&mut id_buf, *n);
}
crate::PointId::Uuid(s) => {
id_buf.put_u8(POINT_ID_UUID);
encode_varint(&mut id_buf, s.len());
id_buf.extend_from_slice(s.as_bytes());
}
}
buf.put_u8(0x12); encode_varint(buf, id_buf.len());
buf.extend_from_slice(&id_buf);
}
buf.put_u8(0x22); encode_varint(buf, 2);
buf.put_u8(0x08); buf.put_u8(0x01);
if with_vectors {
buf.put_u8(0x2A); encode_varint(buf, 2);
buf.put_u8(0x08); buf.put_u8(0x01);
}
}
pub fn encode_scroll_points_proto(
buf: &mut BytesMut,
collection: &str,
limit: u32,
offset: Option<&crate::PointId>,
with_vectors: bool,
) {
buf.clear();
buf.put_u8(0x0A);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
if let Some(id) = offset {
let mut id_buf = BytesMut::with_capacity(40);
match id {
crate::PointId::Num(n) => {
id_buf.put_u8(POINT_ID_NUM);
encode_varint_u64(&mut id_buf, *n);
}
crate::PointId::Uuid(s) => {
id_buf.put_u8(POINT_ID_UUID);
encode_varint(&mut id_buf, s.len());
id_buf.extend_from_slice(s.as_bytes());
}
}
buf.put_u8(0x1A); encode_varint(buf, id_buf.len());
buf.extend_from_slice(&id_buf);
}
buf.put_u8(0x20); encode_varint(buf, limit as usize);
buf.put_u8(0x32); encode_varint(buf, 2);
buf.put_u8(0x08);
buf.put_u8(0x01);
if with_vectors {
buf.put_u8(0x3A); encode_varint(buf, 2);
buf.put_u8(0x08);
buf.put_u8(0x01);
}
}
pub fn encode_set_payload_proto(
buf: &mut BytesMut,
collection: &str,
point_ids: &[crate::PointId],
payload: &crate::point::Payload,
wait: bool,
) {
buf.clear();
buf.put_u8(0x0A);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
if wait {
buf.put_u8(0x10);
buf.put_u8(0x01);
}
for (key, value) in payload {
let mut entry_buf = BytesMut::with_capacity(key.len() + 32);
entry_buf.put_u8(0x0A); encode_varint(&mut entry_buf, key.len());
entry_buf.extend_from_slice(key.as_bytes());
let val_buf = encode_payload_value(value);
entry_buf.put_u8(0x12); encode_varint(&mut entry_buf, val_buf.len());
entry_buf.extend_from_slice(&val_buf);
buf.put_u8(0x1A); encode_varint(buf, entry_buf.len());
buf.extend_from_slice(&entry_buf);
}
let selector_buf = encode_points_selector(point_ids);
buf.put_u8(0x2A); encode_varint(buf, selector_buf.len());
buf.extend_from_slice(&selector_buf);
}
pub fn encode_create_field_index_proto(
buf: &mut BytesMut,
collection: &str,
field_name: &str,
field_type: FieldType,
wait: bool,
) {
buf.clear();
buf.put_u8(0x0A);
encode_varint(buf, collection.len());
buf.extend_from_slice(collection.as_bytes());
if wait {
buf.put_u8(0x10);
buf.put_u8(0x01);
}
buf.put_u8(0x1A);
encode_varint(buf, field_name.len());
buf.extend_from_slice(field_name.as_bytes());
buf.put_u8(0x20); encode_varint(buf, field_type as usize);
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum FieldType {
Keyword = 0,
Integer = 1,
Float = 2,
Geo = 3,
Text = 4,
Bool = 5,
Datetime = 6,
}
const CREATE_COLLECTION_NAME: u8 = 0x0A;
const CREATE_VECTORS_CONFIG: u8 = 0x52;
const CREATE_ON_DISK: u8 = 0x40;
const DELETE_COLLECTION_NAME: u8 = 0x0A;
pub fn encode_create_collection_proto(
buf: &mut BytesMut,
collection_name: &str,
vector_size: u64,
distance: qail_core::ast::Distance,
on_disk: bool,
) {
buf.clear();
buf.put_u8(CREATE_COLLECTION_NAME);
encode_varint(buf, collection_name.len());
buf.extend_from_slice(collection_name.as_bytes());
let mut params_buf = BytesMut::with_capacity(32);
params_buf.put_u8(0x08);
encode_varint_u64(&mut params_buf, vector_size);
params_buf.put_u8(0x10);
let distance_val = match distance {
qail_core::ast::Distance::Cosine => 1,
qail_core::ast::Distance::Euclid => 2,
qail_core::ast::Distance::Dot => 3,
};
encode_varint(&mut params_buf, distance_val);
if on_disk {
params_buf.put_u8(0x28);
params_buf.put_u8(0x01);
}
let mut config_buf = BytesMut::with_capacity(params_buf.len() + 4);
config_buf.put_u8(0x0A);
encode_varint(&mut config_buf, params_buf.len());
config_buf.extend_from_slice(¶ms_buf);
buf.put_u8(CREATE_VECTORS_CONFIG);
encode_varint(buf, config_buf.len());
buf.extend_from_slice(&config_buf);
if on_disk {
buf.put_u8(CREATE_ON_DISK);
buf.put_u8(0x01);
}
}
pub fn encode_delete_collection_proto(buf: &mut BytesMut, collection_name: &str) {
buf.clear();
buf.put_u8(DELETE_COLLECTION_NAME);
encode_varint(buf, collection_name.len());
buf.extend_from_slice(collection_name.as_bytes());
}
pub fn encode_delete_points_mixed_proto(
buf: &mut BytesMut,
collection_name: &str,
point_ids: &[crate::PointId],
) {
buf.clear();
buf.put_u8(0x0A);
encode_varint(buf, collection_name.len());
buf.put_slice(collection_name.as_bytes());
buf.put_u8(0x10);
buf.put_u8(1);
let selector_buf = encode_points_selector(point_ids);
buf.put_u8(0x22);
encode_varint(buf, selector_buf.len());
buf.extend_from_slice(&selector_buf);
}
pub fn encode_delete_points_proto(
buf: &mut BytesMut,
collection_name: &str,
point_ids: &[u64],
) {
let ids: Vec<crate::PointId> = point_ids.iter().map(|&id| crate::PointId::Num(id)).collect();
encode_delete_points_mixed_proto(buf, collection_name, &ids);
}
fn encode_points_selector(ids: &[crate::PointId]) -> BytesMut {
let mut ids_list = BytesMut::with_capacity(ids.len() * 40);
for id in ids {
let mut id_buf = BytesMut::with_capacity(40);
match id {
crate::PointId::Num(n) => {
id_buf.put_u8(POINT_ID_NUM);
encode_varint_u64(&mut id_buf, *n);
}
crate::PointId::Uuid(s) => {
id_buf.put_u8(POINT_ID_UUID);
encode_varint(&mut id_buf, s.len());
id_buf.extend_from_slice(s.as_bytes());
}
}
ids_list.put_u8(0x0A); encode_varint(&mut ids_list, id_buf.len());
ids_list.extend_from_slice(&id_buf);
}
let mut selector = BytesMut::with_capacity(ids_list.len() + 8);
selector.put_u8(0x0A); encode_varint(&mut selector, ids_list.len());
selector.extend_from_slice(&ids_list);
selector
}
pub fn encode_list_collections_proto(buf: &mut BytesMut) {
buf.clear();
}
pub fn encode_collection_info_proto(buf: &mut BytesMut, collection_name: &str) {
buf.clear();
buf.put_u8(0x0A);
encode_varint(buf, collection_name.len());
buf.extend_from_slice(collection_name.as_bytes());
}
#[inline]
pub fn varint_len(value: u64) -> usize {
if value == 0 {
1
} else {
let bits = 64 - value.leading_zeros() as usize;
bits.div_ceil(7)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_varint_encoding() {
let mut buf = BytesMut::new();
encode_varint(&mut buf, 1);
assert_eq!(&buf[..], &[0x01]);
buf.clear();
encode_varint(&mut buf, 127);
assert_eq!(&buf[..], &[0x7F]);
buf.clear();
encode_varint(&mut buf, 128);
assert_eq!(&buf[..], &[0x80, 0x01]);
buf.clear();
encode_varint(&mut buf, 300);
assert_eq!(&buf[..], &[0xAC, 0x02]);
}
#[test]
fn test_encode_search_basic() {
let mut buf = BytesMut::with_capacity(1024);
let vector = vec![0.1f32, 0.2, 0.3, 0.4];
encode_search_proto(&mut buf, "test_collection", &vector, 10, None, None);
assert_eq!(buf[0], SEARCH_COLLECTION);
assert!(buf.len() > 20);
}
#[test]
fn test_zero_copy_vector() {
let mut buf = BytesMut::with_capacity(1024);
let vector = vec![1.0f32, 2.0, 3.0, 4.0];
encode_search_proto(&mut buf, "test", &vector, 5, None, None);
let vector_start = 8;
let vector_bytes = &buf[vector_start..vector_start + 16];
let float_bytes: [u8; 4] = 1.0f32.to_le_bytes();
assert_eq!(&vector_bytes[0..4], &float_bytes);
}
#[test]
fn test_varint_len() {
assert_eq!(varint_len(0), 1);
assert_eq!(varint_len(1), 1);
assert_eq!(varint_len(127), 1);
assert_eq!(varint_len(128), 2);
assert_eq!(varint_len(16383), 2);
assert_eq!(varint_len(16384), 3);
}
#[test]
fn test_encode_search_with_filter() {
use qail_core::ast::{Condition, Expr, Operator, Value};
let mut buf = BytesMut::with_capacity(1024);
let vector = vec![0.1f32, 0.2, 0.3];
let conditions = vec![
Condition {
left: Expr::Named("category".to_string()),
op: Operator::Eq,
value: Value::String("electronics".to_string()),
is_array_unnest: false,
},
Condition {
left: Expr::Named("price".to_string()),
op: Operator::Lt,
value: Value::Int(1000),
is_array_unnest: false,
},
];
encode_search_with_filter_proto(
&mut buf, "products", &vector, 10, None, None, &conditions, false,
);
assert!(buf.len() > 50);
assert_eq!(buf[0], SEARCH_COLLECTION);
assert!(buf.iter().any(|&b| b == SEARCH_FILTER));
}
#[test]
fn test_encode_get_points() {
let mut buf = BytesMut::with_capacity(1024);
let ids = vec![
crate::PointId::Num(42),
crate::PointId::Uuid("abc-123".to_string()),
];
encode_get_points_proto(&mut buf, "my_collection", &ids, true);
assert_eq!(buf[0], 0x0A); assert!(buf.len() > 20);
}
#[test]
fn test_encode_scroll_points() {
let mut buf = BytesMut::with_capacity(1024);
encode_scroll_points_proto(&mut buf, "my_collection", 100, None, false);
assert_eq!(buf[0], 0x0A);
assert!(buf.len() > 10);
}
#[test]
fn test_encode_delete_points_uuid() {
let mut buf = BytesMut::with_capacity(1024);
let ids = vec![
crate::PointId::Uuid("test-uuid-1".to_string()),
crate::PointId::Num(99),
];
encode_delete_points_mixed_proto(&mut buf, "products", &ids);
assert_eq!(buf[0], 0x0A); assert!(buf.len() > 20);
}
#[test]
fn test_encode_set_payload() {
let mut buf = BytesMut::with_capacity(1024);
let ids = vec![crate::PointId::Num(1)];
let mut payload = crate::point::Payload::new();
payload.insert("name".to_string(), crate::point::PayloadValue::String("updated".to_string()));
encode_set_payload_proto(&mut buf, "my_col", &ids, &payload, true);
assert_eq!(buf[0], 0x0A);
assert!(buf.len() > 15);
}
#[test]
fn test_encode_create_field_index() {
let mut buf = BytesMut::with_capacity(256);
encode_create_field_index_proto(&mut buf, "products", "category", FieldType::Keyword, true);
assert_eq!(buf[0], 0x0A);
assert!(buf.len() > 10);
}
#[test]
fn test_encode_payload_value_string() {
let val = crate::point::PayloadValue::String("hello".to_string());
let buf = encode_payload_value(&val);
assert_eq!(buf[0], 0x22);
assert!(buf.len() > 5);
}
#[test]
fn test_encode_payload_value_integer() {
let val = crate::point::PayloadValue::Integer(42);
let buf = encode_payload_value(&val);
assert_eq!(buf[0], 0x18);
}
}