1use super::entity::*;
4use reddb_file::native_store::{
5 write_native_varu32 as write_varu32, write_native_varu64 as write_varu64,
6};
7use reddb_file::{
8 embedded::{RdbFileError, RdbFileResult},
9 native_store::*,
10 native_value::{write_native_value, NativeReader},
11};
12use std::collections::HashMap;
13
14pub fn write_entity_binary(buf: &mut Vec<u8>, entity: &UnifiedEntity, format_version: u32) {
16 write_varu64(buf, entity.id.raw());
18
19 match &entity.kind {
21 EntityKind::TableRow { table, row_id } => {
22 buf.push(0);
23 write_varu32(buf, table.len() as u32);
24 buf.extend_from_slice(table.as_bytes());
25 write_varu64(buf, *row_id);
26 }
27 EntityKind::GraphNode(ref node) => {
28 buf.push(1);
29 write_varu32(buf, node.label.len() as u32);
30 buf.extend_from_slice(node.label.as_bytes());
31 write_varu32(buf, node.node_type.len() as u32);
32 buf.extend_from_slice(node.node_type.as_bytes());
33 }
34 EntityKind::GraphEdge(ref edge) => {
35 buf.push(2);
36 write_varu32(buf, edge.label.len() as u32);
37 buf.extend_from_slice(edge.label.as_bytes());
38 write_varu32(buf, edge.from_node.len() as u32);
39 buf.extend_from_slice(edge.from_node.as_bytes());
40 write_varu32(buf, edge.to_node.len() as u32);
41 buf.extend_from_slice(edge.to_node.as_bytes());
42 buf.extend_from_slice(&edge.weight.to_le_bytes());
43 }
44 EntityKind::Vector { collection } => {
45 buf.push(3);
46 write_varu32(buf, collection.len() as u32);
47 buf.extend_from_slice(collection.as_bytes());
48 }
49 EntityKind::TimeSeriesPoint(ref ts) => {
50 buf.push(4);
51 write_varu32(buf, ts.series.len() as u32);
52 buf.extend_from_slice(ts.series.as_bytes());
53 write_varu32(buf, ts.metric.len() as u32);
54 buf.extend_from_slice(ts.metric.as_bytes());
55 }
56 EntityKind::QueueMessage { queue, position } => {
57 buf.push(5);
58 write_varu32(buf, queue.len() as u32);
59 buf.extend_from_slice(queue.as_bytes());
60 write_varu64(buf, *position);
61 }
62 }
63
64 match &entity.data {
66 EntityData::Row(row) => {
67 if row.named.is_some()
68 || row
69 .schema
70 .as_ref()
71 .is_some_and(|schema| schema.len() == row.columns.len())
72 {
73 buf.push(6);
76 let mut entries: Vec<_> = if let Some(named) = &row.named {
77 named
78 .iter()
79 .map(|(name, value)| (name.as_str(), value))
80 .collect()
81 } else {
82 row.schema
83 .as_ref()
84 .expect("validated schema")
85 .iter()
86 .map(String::as_str)
87 .zip(&row.columns)
88 .collect()
89 };
90 write_varu32(buf, entries.len() as u32);
91 entries.sort_by(|a, b| a.0.cmp(b.0));
92 for (key, value) in entries {
93 write_varu32(buf, key.len() as u32);
94 buf.extend_from_slice(key.as_bytes());
95 write_native_value(buf, value);
96 }
97 } else {
98 buf.push(0);
99 write_varu32(buf, row.columns.len() as u32);
100 for col in &row.columns {
101 write_native_value(buf, col);
102 }
103 }
104 }
105 EntityData::Node(node) => {
106 buf.push(1);
107 write_varu32(buf, node.properties.len() as u32);
108 let mut entries: Vec<_> = node.properties.iter().collect();
109 entries.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str()));
110 for (key, value) in entries {
111 write_varu32(buf, key.len() as u32);
112 buf.extend_from_slice(key.as_bytes());
113 write_native_value(buf, value);
114 }
115 }
116 EntityData::Edge(edge) => {
117 buf.push(2);
118 buf.extend_from_slice(&edge.weight.to_le_bytes());
119 write_varu32(buf, edge.properties.len() as u32);
120 let mut entries: Vec<_> = edge.properties.iter().collect();
121 entries.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str()));
122 for (key, value) in entries {
123 write_varu32(buf, key.len() as u32);
124 buf.extend_from_slice(key.as_bytes());
125 write_native_value(buf, value);
126 }
127 }
128 EntityData::Vector(vec) => {
129 buf.push(3);
130 write_varu32(buf, vec.dense.len() as u32);
131 for f in &vec.dense {
132 buf.extend_from_slice(&f.to_le_bytes());
133 }
134 if format_version >= STORE_VERSION_V6 {
135 buf.push(u8::from(vec.content.is_some()));
136 if let Some(content) = &vec.content {
137 write_varu32(buf, content.len() as u32);
138 buf.extend_from_slice(content.as_bytes());
139 }
140 }
141 }
142 EntityData::TimeSeries(ts) => {
143 buf.push(4);
144 write_varu32(buf, ts.metric.len() as u32);
145 buf.extend_from_slice(ts.metric.as_bytes());
146 write_varu64(buf, ts.timestamp_ns);
147 buf.extend_from_slice(&ts.value.to_le_bytes());
148 if format_version >= STORE_VERSION_V11 {
149 match ts.series_id {
155 Some(series_id) => {
156 write_varu32(buf, 1);
157 write_varu64(buf, series_id);
158 }
159 None => {
160 write_varu32(buf, 0);
161 write_varu32(buf, ts.tags.len() as u32);
162 let mut tag_entries: Vec<_> = ts.tags.iter().collect();
163 tag_entries.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str()));
164 for (key, value) in tag_entries {
165 write_varu32(buf, key.len() as u32);
166 buf.extend_from_slice(key.as_bytes());
167 write_varu32(buf, value.len() as u32);
168 buf.extend_from_slice(value.as_bytes());
169 }
170 }
171 }
172 } else if format_version >= STORE_VERSION_V5 {
173 write_varu32(buf, ts.tags.len() as u32);
174 let mut tag_entries: Vec<_> = ts.tags.iter().collect();
175 tag_entries.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str()));
176 for (key, value) in tag_entries {
177 write_varu32(buf, key.len() as u32);
178 buf.extend_from_slice(key.as_bytes());
179 write_varu32(buf, value.len() as u32);
180 buf.extend_from_slice(value.as_bytes());
181 }
182 }
183 if format_version >= STORE_VERSION_V12 {
184 write_varu32(buf, ts.fields.len() as u32);
185 let mut field_entries: Vec<_> = ts.fields.iter().collect();
186 field_entries.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str()));
187 for (key, value) in field_entries {
188 write_varu32(buf, key.len() as u32);
189 buf.extend_from_slice(key.as_bytes());
190 write_native_value(buf, value);
191 }
192 }
193 }
194 EntityData::QueueMessage(msg) => {
195 buf.push(5);
196 write_native_value(buf, &msg.payload);
197 write_varu64(buf, msg.enqueued_at_ns);
198 write_varu32(buf, msg.attempts);
199 }
200 }
201
202 write_varu64(buf, entity.created_at);
204 write_varu64(buf, entity.updated_at);
205
206 write_varu32(buf, entity.embeddings().len() as u32);
208 for emb in entity.embeddings() {
209 write_varu32(buf, emb.name.len() as u32);
210 buf.extend_from_slice(emb.name.as_bytes());
211 write_varu32(buf, emb.vector.len() as u32);
212 for f in &emb.vector {
213 buf.extend_from_slice(&f.to_le_bytes());
214 }
215 write_varu32(buf, emb.model.len() as u32);
216 buf.extend_from_slice(emb.model.as_bytes());
217 }
218
219 write_varu32(buf, entity.cross_refs().len() as u32);
221 for cross_ref in entity.cross_refs() {
222 write_varu64(buf, cross_ref.source.raw());
223 write_varu64(buf, cross_ref.target.raw());
224 buf.push(cross_ref.ref_type.to_byte());
225 if format_version >= STORE_VERSION_V2 {
226 write_varu32(buf, cross_ref.target_collection.len() as u32);
227 buf.extend_from_slice(cross_ref.target_collection.as_bytes());
228 buf.extend_from_slice(&cross_ref.weight.to_le_bytes());
229 write_varu64(buf, cross_ref.created_at);
230 }
231 }
232
233 write_varu64(buf, entity.sequence_id);
235
236 if format_version >= STORE_VERSION_V4 {
237 if let EntityData::QueueMessage(message) = &entity.data {
238 buf.push(u8::from(message.priority.is_some()));
239 if let Some(priority) = message.priority {
240 buf.extend_from_slice(&priority.to_le_bytes());
241 }
242 write_varu32(buf, message.max_attempts);
243 buf.push(u8::from(message.acked));
244 }
245 }
246
247 if format_version >= STORE_VERSION_V8 {
248 buf.push(u8::from(entity.has_explicit_logical_id()));
249 if entity.has_explicit_logical_id() {
250 write_varu64(buf, entity.logical_id().raw());
251 }
252 }
253 if format_version >= STORE_VERSION_V9 {
254 write_varu64(buf, entity.xmin);
255 write_varu64(buf, entity.xmax);
256 }
257}
258
259pub fn read_entity_binary(
262 bytes: &[u8],
263 position: &mut usize,
264 format_version: u32,
265) -> RdbFileResult<UnifiedEntity> {
266 let mut reader = NativeReader::new(bytes, position);
267 let id = EntityId::new(reader.variable_u64()?);
268 let kind = match reader.byte()? {
269 0 => EntityKind::TableRow {
270 table: reader.string()?.into(),
271 row_id: reader.variable_u64()?,
272 },
273 1 => EntityKind::GraphNode(Box::new(GraphNodeKind {
274 label: reader.string()?,
275 node_type: reader.string()?,
276 })),
277 2 => EntityKind::GraphEdge(Box::new(GraphEdgeKind {
278 label: reader.string()?,
279 from_node: reader.string()?,
280 to_node: reader.string()?,
281 weight: u32::from_le_bytes(reader.fixed()?),
282 })),
283 3 => EntityKind::Vector {
284 collection: reader.string()?,
285 },
286 4 => EntityKind::TimeSeriesPoint(Box::new(TimeSeriesPointKind {
287 series: reader.string()?,
288 metric: reader.string()?,
289 })),
290 5 => EntityKind::QueueMessage {
291 queue: reader.string()?,
292 position: reader.variable_u64()?,
293 },
294 tag => return Err(invalid(format!("Unknown EntityKind type: {tag}"))),
295 };
296 let mut data = match reader.byte()? {
297 0 => {
298 let count = reader.count(1)?;
299 let mut columns = Vec::with_capacity(count);
300 for _ in 0..count {
301 columns.push(reader.value()?);
302 }
303 EntityData::Row(RowData::new(columns))
304 }
305 1 => EntityData::Node(NodeData::with_properties(read_properties(&mut reader)?)),
306 2 => {
307 let mut edge = EdgeData::new(f32::from_le_bytes(reader.fixed()?));
308 edge.properties = read_properties(&mut reader)?;
309 EntityData::Edge(edge)
310 }
311 3 => {
312 let mut vector = VectorData::new(read_vector(&mut reader)?);
313 if format_version >= STORE_VERSION_V6 && reader.byte()? != 0 {
314 vector.content = Some(reader.string()?);
315 }
316 EntityData::Vector(vector)
317 }
318 4 => {
319 let metric = reader.string()?;
320 let timestamp_ns = reader.variable_u64()?;
321 let value = f64::from_le_bytes(reader.fixed()?);
322 let mut tags = HashMap::new();
323 let series_id = if format_version >= STORE_VERSION_V11 && reader.variable_u32()? == 1 {
324 Some(reader.variable_u64()?)
325 } else {
326 if format_version >= STORE_VERSION_V5 {
327 let count = reader.count(2)?;
328 tags = HashMap::with_capacity(count);
329 for _ in 0..count {
330 tags.insert(reader.string()?, reader.string()?);
331 }
332 }
333 None
334 };
335 let fields = if format_version >= STORE_VERSION_V12 {
336 read_properties(&mut reader)?
337 } else {
338 HashMap::new()
339 };
340 EntityData::TimeSeries(TimeSeriesData {
341 metric,
342 timestamp_ns,
343 value,
344 series_id,
345 tags,
346 fields,
347 })
348 }
349 5 => EntityData::QueueMessage(QueueMessageData {
350 payload: reader.value()?,
351 priority: None,
352 enqueued_at_ns: reader.variable_u64()?,
353 attempts: reader.variable_u32()?,
354 max_attempts: 3,
355 acked: false,
356 }),
357 6 => EntityData::Row(RowData {
358 columns: Vec::new(),
359 named: Some(read_properties(&mut reader)?),
360 schema: None,
361 }),
362 tag => return Err(invalid(format!("Unknown EntityData type: {tag}"))),
363 };
364 let created_at = reader.variable_u64()?;
365 let updated_at = reader.variable_u64()?;
366 let count = reader.count(3)?;
367 let mut embeddings = Vec::with_capacity(count);
368 for _ in 0..count {
369 let name = reader.string()?;
370 let vector = read_vector(&mut reader)?;
371 let model = reader.string()?;
372 embeddings.push(EmbeddingSlot::new(name, vector, model));
373 }
374 let count = reader.count(3)?;
375 let mut cross_refs = Vec::with_capacity(count);
376 for _ in 0..count {
377 let source = EntityId::new(reader.variable_u64()?);
378 let target = EntityId::new(reader.variable_u64()?);
379 let ref_type = RefType::from_byte(reader.byte()?);
380 let (collection, weight, created_at) = if format_version >= STORE_VERSION_V2 {
381 (
382 reader.string()?,
383 f32::from_le_bytes(reader.fixed()?),
384 reader.variable_u64()?,
385 )
386 } else {
387 (String::new(), 1.0, 0)
388 };
389 let mut cross_ref = CrossRef::new(source, target, collection, ref_type);
390 cross_ref.weight = weight;
391 cross_ref.created_at = created_at;
392 cross_refs.push(cross_ref);
393 }
394 let sequence_id = reader.variable_u64()?;
395 if format_version >= STORE_VERSION_V4 {
396 if let EntityData::QueueMessage(message) = &mut data {
397 if reader.remaining() != 0 {
398 message.priority = if reader.byte()? != 0 {
399 Some(i32::from_le_bytes(reader.fixed()?))
400 } else {
401 None
402 };
403 message.max_attempts = reader.variable_u32()?;
404 message.acked = reader.byte()? != 0;
405 }
406 }
407 }
408 let mut entity = UnifiedEntity::new(id, kind, data);
409 entity.created_at = created_at;
410 entity.updated_at = updated_at;
411 entity.sequence_id = sequence_id;
412 if format_version >= STORE_VERSION_V8 && reader.remaining() != 0 && reader.byte()? != 0 {
413 entity.set_logical_id(EntityId::new(reader.variable_u64()?));
414 }
415 if format_version >= STORE_VERSION_V9 && reader.remaining() != 0 {
416 entity.set_xmin(reader.variable_u64()?);
417 entity.set_xmax(reader.variable_u64()?);
418 }
419 if !embeddings.is_empty() || !cross_refs.is_empty() {
420 entity.embeddings_mut().extend(embeddings);
421 entity.cross_refs_mut().extend(cross_refs);
422 }
423 Ok(entity)
424}
425
426fn read_properties(
427 reader: &mut NativeReader<'_, '_>,
428) -> RdbFileResult<HashMap<String, reddb_types::Value>> {
429 let count = reader.count(2)?;
430 let mut properties = HashMap::with_capacity(count);
431 for _ in 0..count {
432 properties.insert(reader.string()?, reader.value()?);
433 }
434 Ok(properties)
435}
436
437fn read_vector(reader: &mut NativeReader<'_, '_>) -> RdbFileResult<Vec<f32>> {
438 let count = reader.count(4)?;
439 let mut vector = Vec::with_capacity(count);
440 for _ in 0..count {
441 vector.push(f32::from_le_bytes(reader.fixed()?));
442 }
443 Ok(vector)
444}
445
446fn invalid(error: impl std::fmt::Display) -> RdbFileError {
447 RdbFileError::InvalidOperation(error.to_string())
448}