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 read_entity_from_reader(&mut NativeReader::new(bytes, position), format_version)
267}
268
269pub fn read_entity_binary_with_budget(
270 bytes: &[u8],
271 position: &mut usize,
272 format_version: u32,
273 budget: &mut reddb_file::native_value::DecodeBudget,
274) -> RdbFileResult<UnifiedEntity> {
275 budget.elements::<UnifiedEntity>(1)?;
276 read_entity_from_reader(
277 &mut NativeReader::with_budget(bytes, position, budget),
278 format_version,
279 )
280}
281
282fn read_entity_from_reader(
283 reader: &mut NativeReader<'_, '_>,
284 format_version: u32,
285) -> RdbFileResult<UnifiedEntity> {
286 let id = EntityId::new(reader.variable_u64()?);
287 let kind = match reader.byte()? {
288 0 => {
289 let table = reader.string()?;
290 reader.reserve_bytes(table.len().saturating_add(2 * std::mem::size_of::<usize>()))?;
291 EntityKind::TableRow {
292 table: table.into(),
293 row_id: reader.variable_u64()?,
294 }
295 }
296 1 => {
297 reader.reserve_elements::<GraphNodeKind>(1)?;
298 EntityKind::GraphNode(Box::new(GraphNodeKind {
299 label: reader.string()?,
300 node_type: reader.string()?,
301 }))
302 }
303 2 => {
304 reader.reserve_elements::<GraphEdgeKind>(1)?;
305 EntityKind::GraphEdge(Box::new(GraphEdgeKind {
306 label: reader.string()?,
307 from_node: reader.string()?,
308 to_node: reader.string()?,
309 weight: u32::from_le_bytes(reader.fixed()?),
310 }))
311 }
312 3 => EntityKind::Vector {
313 collection: reader.string()?,
314 },
315 4 => {
316 reader.reserve_elements::<TimeSeriesPointKind>(1)?;
317 EntityKind::TimeSeriesPoint(Box::new(TimeSeriesPointKind {
318 series: reader.string()?,
319 metric: reader.string()?,
320 }))
321 }
322 5 => EntityKind::QueueMessage {
323 queue: reader.string()?,
324 position: reader.variable_u64()?,
325 },
326 tag => return Err(invalid(format!("Unknown EntityKind type: {tag}"))),
327 };
328 let mut data = match reader.byte()? {
329 0 => {
330 let count = reader.count(1)?;
331 reader.reserve_elements::<reddb_types::Value>(count)?;
332 let mut columns = Vec::with_capacity(count);
333 for _ in 0..count {
334 columns.push(reader.value()?);
335 }
336 EntityData::Row(RowData::new(columns))
337 }
338 1 => EntityData::Node(NodeData::with_properties(read_properties(reader)?)),
339 2 => {
340 let mut edge = EdgeData::new(f32::from_le_bytes(reader.fixed()?));
341 edge.properties = read_properties(reader)?;
342 EntityData::Edge(edge)
343 }
344 3 => {
345 let mut vector = VectorData::new(read_vector(reader)?);
346 if format_version >= STORE_VERSION_V6 && reader.byte()? != 0 {
347 vector.content = Some(reader.string()?);
348 }
349 EntityData::Vector(vector)
350 }
351 4 => {
352 let metric = reader.string()?;
353 let timestamp_ns = reader.variable_u64()?;
354 let value = f64::from_le_bytes(reader.fixed()?);
355 let mut tags = HashMap::new();
356 let series_id = if format_version >= STORE_VERSION_V11 && reader.variable_u32()? == 1 {
357 Some(reader.variable_u64()?)
358 } else {
359 if format_version >= STORE_VERSION_V5 {
360 let count = reader.count(2)?;
361 reader.reserve_map::<String, String>(count)?;
362 tags = HashMap::with_capacity(count);
363 for _ in 0..count {
364 tags.insert(reader.string()?, reader.string()?);
365 }
366 }
367 None
368 };
369 let fields = if format_version >= STORE_VERSION_V12 {
370 read_properties(reader)?
371 } else {
372 HashMap::new()
373 };
374 EntityData::TimeSeries(TimeSeriesData {
375 metric,
376 timestamp_ns,
377 value,
378 series_id,
379 tags,
380 fields,
381 })
382 }
383 5 => EntityData::QueueMessage(QueueMessageData {
384 payload: reader.value()?,
385 priority: None,
386 enqueued_at_ns: reader.variable_u64()?,
387 attempts: reader.variable_u32()?,
388 max_attempts: 3,
389 acked: false,
390 }),
391 6 => EntityData::Row(RowData {
392 columns: Vec::new(),
393 named: Some(read_properties(reader)?),
394 schema: None,
395 }),
396 tag => return Err(invalid(format!("Unknown EntityData type: {tag}"))),
397 };
398 let created_at = reader.variable_u64()?;
399 let updated_at = reader.variable_u64()?;
400 let count = reader.count(3)?;
401 reader.reserve_elements::<EmbeddingSlot>(count)?;
402 reader.reserve_elements::<EmbeddingSlot>(count)?;
403 let mut embeddings = Vec::with_capacity(count);
404 for _ in 0..count {
405 let name = reader.string()?;
406 let vector = read_vector(reader)?;
407 let model = reader.string()?;
408 embeddings.push(EmbeddingSlot::new(name, vector, model));
409 }
410 let count = reader.count(3)?;
411 reader.reserve_elements::<CrossRef>(count)?;
412 reader.reserve_elements::<CrossRef>(count)?;
413 let mut cross_refs = Vec::with_capacity(count);
414 for _ in 0..count {
415 let source = EntityId::new(reader.variable_u64()?);
416 let target = EntityId::new(reader.variable_u64()?);
417 let ref_type = RefType::from_byte(reader.byte()?);
418 let (collection, weight, created_at) = if format_version >= STORE_VERSION_V2 {
419 (
420 reader.string()?,
421 f32::from_le_bytes(reader.fixed()?),
422 reader.variable_u64()?,
423 )
424 } else {
425 (String::new(), 1.0, 0)
426 };
427 let mut cross_ref = CrossRef::new(source, target, collection, ref_type);
428 cross_ref.weight = weight;
429 cross_ref.created_at = created_at;
430 cross_refs.push(cross_ref);
431 }
432 let sequence_id = reader.variable_u64()?;
433 if format_version >= STORE_VERSION_V4 {
434 if let EntityData::QueueMessage(message) = &mut data {
435 if reader.remaining() != 0 {
436 message.priority = if reader.byte()? != 0 {
437 Some(i32::from_le_bytes(reader.fixed()?))
438 } else {
439 None
440 };
441 message.max_attempts = reader.variable_u32()?;
442 message.acked = reader.byte()? != 0;
443 }
444 }
445 }
446 let mut entity = UnifiedEntity::new(id, kind, data);
447 entity.created_at = created_at;
448 entity.updated_at = updated_at;
449 entity.sequence_id = sequence_id;
450 if format_version >= STORE_VERSION_V8 && reader.remaining() != 0 && reader.byte()? != 0 {
451 entity.set_logical_id(EntityId::new(reader.variable_u64()?));
452 }
453 if format_version >= STORE_VERSION_V9 && reader.remaining() != 0 {
454 entity.set_xmin(reader.variable_u64()?);
455 entity.set_xmax(reader.variable_u64()?);
456 }
457 if !embeddings.is_empty() || !cross_refs.is_empty() {
458 reader.reserve_elements::<EntityAux>(1)?;
459 entity.embeddings_mut().extend(embeddings);
460 entity.cross_refs_mut().extend(cross_refs);
461 }
462 Ok(entity)
463}
464
465fn read_properties(
466 reader: &mut NativeReader<'_, '_>,
467) -> RdbFileResult<HashMap<String, reddb_types::Value>> {
468 let count = reader.count(2)?;
469 reader.reserve_map::<String, reddb_types::Value>(count)?;
470 let mut properties = HashMap::with_capacity(count);
471 for _ in 0..count {
472 properties.insert(reader.string()?, reader.value()?);
473 }
474 Ok(properties)
475}
476
477fn read_vector(reader: &mut NativeReader<'_, '_>) -> RdbFileResult<Vec<f32>> {
478 let count = reader.count(4)?;
479 reader.reserve_elements::<f32>(count)?;
480 let mut vector = Vec::with_capacity(count);
481 for _ in 0..count {
482 vector.push(f32::from_le_bytes(reader.fixed()?));
483 }
484 Ok(vector)
485}
486
487fn invalid(error: impl std::fmt::Display) -> RdbFileError {
488 RdbFileError::InvalidOperation(error.to_string())
489}