1use super::*;
4use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
5
6const PROB_HLL_STATE_PREFIX: &str = "probabilistic.hll.";
7const PROB_SKETCH_STATE_PREFIX: &str = "probabilistic.sketch.";
8const PROB_FILTER_STATE_PREFIX: &str = "probabilistic.filter.";
9const PROB_ENCODING_MARKER_KEY: &str = "probabilistic.state_encoding";
10const PROB_ENCODING_RAW_V1: &str = "raw-v1";
11const PROB_WAL_KIND_HLL: u8 = 1;
12const PROB_WAL_KIND_SKETCH: u8 = 2;
13const PROB_WAL_KIND_FILTER: u8 = 3;
14const PROB_WAL_OP_ADD: u8 = 1;
15const PROB_WAL_OP_DELETE: u8 = 2;
16const PROB_WAL_OP_SNAPSHOT: u8 = 3;
17const PROB_WAL_OP_DROP: u8 = 4;
18
19fn probabilistic_read<'a, T>(lock: &'a RwLock<T>, _name: &str) -> RwLockReadGuard<'a, T> {
20 lock.read()
21}
22
23fn probabilistic_write<'a, T>(lock: &'a RwLock<T>, _name: &str) -> RwLockWriteGuard<'a, T> {
24 lock.write()
25}
26
27fn probabilistic_collection_contract(
28 name: &str,
29 model: crate::catalog::CollectionModel,
30) -> crate::physical::CollectionContract {
31 let now = crate::utils::now_unix_millis() as u128;
32 crate::physical::CollectionContract {
33 name: name.to_string(),
34 description: None,
35 description_generated: false,
36 declared_model: model,
37 schema_mode: crate::catalog::SchemaMode::Dynamic,
38 origin: crate::physical::ContractOrigin::Explicit,
39 version: 1,
40 created_at_unix_ms: now,
41 updated_at_unix_ms: now,
42 default_ttl_ms: None,
43 vector_dimension: None,
44 vector_metric: None,
45 context_index_fields: Vec::new(),
46 declared_columns: Vec::new(),
47 table_def: None,
48 timestamps_enabled: false,
49 context_index_enabled: false,
50 metrics_raw_retention_ms: None,
51 metrics_rollup_policies: Vec::new(),
52 metrics_tenant_identity: None,
53 metrics_namespace: None,
54 append_only: false,
55 subscriptions: Vec::new(),
56 analytics_config: Vec::new(),
57 session_key: None,
58 session_gap_ms: None,
59 retention_duration_ms: None,
60 analytical_storage: None,
61
62 ai_policy: None,
63 }
64}
65
66fn hll_precision_mismatch_error(
67 operation: &str,
68 expected_name: &str,
69 expected_precision: u8,
70 actual_name: &str,
71 actual_precision: u8,
72) -> RedDBError {
73 RedDBError::Query(format!(
74 "HLL {operation} requires matching precision; '{expected_name}' uses precision {expected_precision}, but '{actual_name}' uses precision {actual_precision}"
75 ))
76}
77
78enum ProbabilisticReadProjection {
79 Cardinality { label: String },
80 Freq { element: String, label: String },
81 Contains { element: String, label: String },
82}
83
84struct ProbabilisticStateEntry {
85 name: String,
86 bytes: Vec<u8>,
87 migrated_from_legacy_hex: bool,
88}
89
90impl RedDBRuntime {
91 pub(crate) fn load_probabilistic_state(&self) -> RedDBResult<()> {
92 let reject_legacy_hex = self.probabilistic_raw_encoding_marker()?;
93 let mut migrated_legacy_hex = false;
94 {
95 let entries = self.latest_probabilistic_state_entries(
96 PROB_HLL_STATE_PREFIX,
97 "HLL",
98 reject_legacy_hex,
99 )?;
100 let mut hlls =
101 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
102 for entry in entries {
103 let Some(hll) = crate::storage::primitives::hyperloglog::HyperLogLog::from_bytes(
104 entry.bytes.clone(),
105 ) else {
106 return Err(RedDBError::Internal(format!(
107 "invalid persisted HLL state for '{}'",
108 entry.name
109 )));
110 };
111 if entry.migrated_from_legacy_hex {
112 self.persist_probabilistic_blob(
113 PROB_HLL_STATE_PREFIX,
114 &entry.name,
115 &entry.bytes,
116 )?;
117 migrated_legacy_hex = true;
118 }
119 hlls.insert(entry.name, hll);
120 }
121 }
122
123 {
124 let entries = self.latest_probabilistic_state_entries(
125 PROB_SKETCH_STATE_PREFIX,
126 "SKETCH",
127 reject_legacy_hex,
128 )?;
129 let mut sketches = probabilistic_write(
130 &self.inner.probabilistic.sketches,
131 "probabilistic sketch store",
132 );
133 for entry in entries {
134 let sketch =
135 crate::storage::primitives::count_min_sketch::CountMinSketch::from_bytes(
136 &entry.bytes,
137 )
138 .ok_or_else(|| {
139 RedDBError::Internal(format!(
140 "invalid persisted SKETCH state for '{}'",
141 entry.name
142 ))
143 })?;
144 if entry.migrated_from_legacy_hex {
145 self.persist_probabilistic_blob(
146 PROB_SKETCH_STATE_PREFIX,
147 &entry.name,
148 &entry.bytes,
149 )?;
150 migrated_legacy_hex = true;
151 }
152 sketches.insert(entry.name, sketch);
153 }
154 }
155
156 {
157 let entries = self.latest_probabilistic_state_entries(
158 PROB_FILTER_STATE_PREFIX,
159 "FILTER",
160 reject_legacy_hex,
161 )?;
162 let mut filters = probabilistic_write(
163 &self.inner.probabilistic.filters,
164 "probabilistic filter store",
165 );
166 for entry in entries {
167 let filter = crate::storage::primitives::cuckoo_filter::CuckooFilter::from_bytes(
168 &entry.bytes,
169 )
170 .ok_or_else(|| {
171 RedDBError::Internal(format!(
172 "invalid persisted FILTER state for '{}'",
173 entry.name
174 ))
175 })?;
176 if entry.migrated_from_legacy_hex {
177 self.persist_probabilistic_blob(
178 PROB_FILTER_STATE_PREFIX,
179 &entry.name,
180 &entry.bytes,
181 )?;
182 migrated_legacy_hex = true;
183 }
184 filters.insert(entry.name, filter);
185 }
186 }
187
188 if migrated_legacy_hex {
189 self.persist_probabilistic_encoding_marker()?;
190 }
191 self.apply_replayed_probabilistic_deltas()?;
192 Ok(())
193 }
194
195 fn latest_probabilistic_state_entries(
196 &self,
197 prefix: &str,
198 kind: &str,
199 reject_legacy_hex: bool,
200 ) -> RedDBResult<Vec<ProbabilisticStateEntry>> {
201 let Some(manager) = self.inner.db.store().get_collection("red.config") else {
202 return Ok(Vec::new());
203 };
204 let mut latest: std::collections::HashMap<String, (u64, Option<Value>)> =
205 std::collections::HashMap::new();
206 manager.try_for_each_entity(usize::MAX, |entity| {
207 let EntityData::Row(row) = &entity.data else {
208 return true;
209 };
210 let Some(named) = &row.named else {
211 return true;
212 };
213 let Some(Value::Text(key)) = named.get("key") else {
214 return true;
215 };
216 let Some(encoded_name) = key.strip_prefix(prefix) else {
217 return true;
218 };
219 let value = match named.get("value") {
220 Some(Value::Blob(value)) => Some(Value::Blob(value.clone())),
221 Some(Value::Text(value)) => Some(Value::Text(value.clone())),
222 Some(Value::Null) => None,
223 _ => return true,
224 };
225 let entity_id = entity.id.raw();
226 match latest.get(encoded_name) {
227 Some((existing_id, _)) if *existing_id > entity_id => {}
228 _ => {
229 latest.insert(encoded_name.to_string(), (entity_id, value));
230 }
231 }
232 true
233 })?;
234
235 let mut entries = Vec::new();
236 for (encoded_name, (_, value)) in latest {
237 let Some(value) = value else {
238 continue;
239 };
240 let Some(name) = hex::decode(&encoded_name)
241 .ok()
242 .and_then(|bytes| String::from_utf8(bytes).ok())
243 else {
244 continue;
245 };
246 match value {
247 Value::Blob(bytes) => entries.push(ProbabilisticStateEntry {
248 name,
249 bytes,
250 migrated_from_legacy_hex: false,
251 }),
252 Value::Text(_) if reject_legacy_hex => {
253 return Err(RedDBError::Internal(format!(
254 "legacy hex-encoded {kind} state for '{name}' is rejected after probabilistic state migrated to raw bytes"
255 )));
256 }
257 Value::Text(data_hex) => {
258 let bytes = hex::decode(data_hex.as_ref()).map_err(|err| {
259 RedDBError::Internal(format!(
260 "invalid legacy hex-encoded {kind} state for '{name}': {err}"
261 ))
262 })?;
263 entries.push(ProbabilisticStateEntry {
264 name,
265 bytes,
266 migrated_from_legacy_hex: true,
267 });
268 }
269 _ => {}
270 }
271 }
272 Ok(entries)
273 }
274
275 fn probabilistic_raw_encoding_marker(&self) -> RedDBResult<bool> {
276 let Some(manager) = self.inner.db.store().get_collection("red.config") else {
277 return Ok(false);
278 };
279 let mut latest: Option<(u64, bool)> = None;
280 manager.try_for_each_entity(usize::MAX, |entity| {
281 let EntityData::Row(row) = &entity.data else {
282 return true;
283 };
284 let Some(named) = &row.named else {
285 return true;
286 };
287 let Some(Value::Text(key)) = named.get("key") else {
288 return true;
289 };
290 if key.as_ref() != PROB_ENCODING_MARKER_KEY {
291 return true;
292 }
293 let is_raw = matches!(
294 named.get("value"),
295 Some(Value::Text(value)) if value.as_ref() == PROB_ENCODING_RAW_V1
296 );
297 let entity_id = entity.id.raw();
298 match latest {
299 Some((existing_id, _)) if existing_id > entity_id => {}
300 _ => latest = Some((entity_id, is_raw)),
301 }
302 true
303 })?;
304 Ok(latest.map(|(_, is_raw)| is_raw).unwrap_or(false))
305 }
306
307 fn persist_probabilistic_blob(
308 &self,
309 prefix: &str,
310 name: &str,
311 bytes: &[u8],
312 ) -> RedDBResult<()> {
313 let key = format!("{prefix}{}", hex::encode(name.as_bytes()));
314 self.compact_config_key(&key)?;
315 self.insert_config_value(&key, Value::Blob(bytes.to_vec()))?;
316 self.persist_probabilistic_encoding_marker()?;
317 Ok(())
318 }
319
320 fn delete_probabilistic_blob(&self, prefix: &str, name: &str) -> RedDBResult<()> {
321 let key = format!("{prefix}{}", hex::encode(name.as_bytes()));
322 self.compact_config_key(&key)?;
323 self.insert_config_value(&key, Value::Null)?;
324 Ok(())
325 }
326
327 pub(crate) fn persist_probabilistic_snapshots(&self) -> RedDBResult<()> {
328 let replayed_names = self.inner.probabilistic.replayed_names.lock().clone();
329 for (kind, name) in replayed_names {
330 let (prefix, present) = match kind {
331 PROB_WAL_KIND_HLL => (
332 PROB_HLL_STATE_PREFIX,
333 self.inner.probabilistic.hlls.read().contains_key(&name),
334 ),
335 PROB_WAL_KIND_SKETCH => (
336 PROB_SKETCH_STATE_PREFIX,
337 self.inner.probabilistic.sketches.read().contains_key(&name),
338 ),
339 PROB_WAL_KIND_FILTER => (
340 PROB_FILTER_STATE_PREFIX,
341 self.inner.probabilistic.filters.read().contains_key(&name),
342 ),
343 _ => {
344 return Err(RedDBError::Internal(
345 "unknown probabilistic WAL kind".into(),
346 ))
347 }
348 };
349 if !present {
350 self.append_probabilistic_drop_delta(kind, &name)?;
351 self.delete_probabilistic_blob(prefix, &name)?;
352 }
353 }
354 let hll_snapshots: Vec<(String, Vec<u8>)> = {
355 let hlls =
356 probabilistic_read(&self.inner.probabilistic.hlls, "probabilistic HLL store");
357 hlls.iter()
358 .map(|(name, hll)| (name.clone(), hll.as_bytes().to_vec()))
359 .collect()
360 };
361 for (name, bytes) in hll_snapshots {
362 self.append_probabilistic_snapshot_delta(PROB_WAL_KIND_HLL, &name, bytes.clone())?;
363 #[cfg(test)]
364 if std::env::var_os("REDDB_TEST_EXIT_AFTER_PROBABILISTIC_SNAPSHOT_WAL").is_some() {
365 std::process::exit(213);
366 }
367 self.persist_probabilistic_blob(PROB_HLL_STATE_PREFIX, &name, &bytes)?;
368 #[cfg(test)]
369 if std::env::var_os("REDDB_TEST_EXIT_AFTER_PROBABILISTIC_SNAPSHOT_BLOB").is_some() {
370 std::process::exit(214);
371 }
372 }
373
374 let sketch_snapshots: Vec<(String, Vec<u8>)> = {
375 let sketches = probabilistic_read(
376 &self.inner.probabilistic.sketches,
377 "probabilistic sketch store",
378 );
379 sketches
380 .iter()
381 .map(|(name, sketch)| (name.clone(), sketch.as_bytes()))
382 .collect()
383 };
384 for (name, bytes) in sketch_snapshots {
385 self.append_probabilistic_snapshot_delta(PROB_WAL_KIND_SKETCH, &name, bytes.clone())?;
386 self.persist_probabilistic_blob(PROB_SKETCH_STATE_PREFIX, &name, &bytes)?;
387 }
388
389 let filter_snapshots: Vec<(String, Vec<u8>)> = {
390 let filters = probabilistic_read(
391 &self.inner.probabilistic.filters,
392 "probabilistic filter store",
393 );
394 filters
395 .iter()
396 .map(|(name, filter)| (name.clone(), filter.as_bytes()))
397 .collect()
398 };
399 for (name, bytes) in filter_snapshots {
400 self.append_probabilistic_snapshot_delta(PROB_WAL_KIND_FILTER, &name, bytes.clone())?;
401 self.persist_probabilistic_blob(PROB_FILTER_STATE_PREFIX, &name, &bytes)?;
402 }
403
404 Ok(())
405 }
406
407 fn append_probabilistic_add_delta(
408 &self,
409 kind: u8,
410 name: &str,
411 operands: Vec<Vec<u8>>,
412 ) -> RedDBResult<()> {
413 self.append_probabilistic_delta(kind, PROB_WAL_OP_ADD, name, operands)
414 }
415
416 fn append_probabilistic_snapshot_delta(
417 &self,
418 kind: u8,
419 name: &str,
420 bytes: Vec<u8>,
421 ) -> RedDBResult<()> {
422 self.append_probabilistic_delta(kind, PROB_WAL_OP_SNAPSHOT, name, vec![bytes])
423 }
424
425 fn append_probabilistic_drop_delta(&self, kind: u8, name: &str) -> RedDBResult<()> {
426 self.append_probabilistic_delta(kind, PROB_WAL_OP_DROP, name, Vec::new())
427 }
428
429 fn append_probabilistic_delete_delta(
430 &self,
431 kind: u8,
432 name: &str,
433 operands: Vec<Vec<u8>>,
434 ) -> RedDBResult<()> {
435 self.append_probabilistic_delta(kind, PROB_WAL_OP_DELETE, name, operands)
436 }
437
438 fn append_probabilistic_delta(
439 &self,
440 kind: u8,
441 operation: u8,
442 name: &str,
443 operands: Vec<Vec<u8>>,
444 ) -> RedDBResult<()> {
445 self.inner
446 .db
447 .store()
448 .append_probabilistic_delta_record(kind, operation, name, operands)
449 .map_err(|err| RedDBError::Internal(format!("append probabilistic WAL delta: {err}")))
450 }
451
452 fn apply_replayed_probabilistic_deltas(&self) -> RedDBResult<()> {
453 let deltas = self.inner.db.store().take_replayed_probabilistic_deltas();
454 let mut reset_at = std::collections::HashMap::new();
458 for (index, (kind, operation, name, _)) in deltas.iter().enumerate() {
459 if matches!(*operation, PROB_WAL_OP_SNAPSHOT | PROB_WAL_OP_DROP) {
460 reset_at.insert((*kind, name.clone()), index);
461 }
462 }
463 for (index, (kind, operation, name, operands)) in deltas.into_iter().enumerate() {
464 if reset_at
465 .get(&(kind, name.clone()))
466 .is_some_and(|reset| index < *reset)
467 {
468 continue;
469 }
470 let present = match kind {
471 PROB_WAL_KIND_HLL => self.inner.probabilistic.hlls.read().contains_key(&name),
472 PROB_WAL_KIND_SKETCH => {
473 self.inner.probabilistic.sketches.read().contains_key(&name)
474 }
475 PROB_WAL_KIND_FILTER => self.inner.probabilistic.filters.read().contains_key(&name),
476 _ => {
477 return Err(RedDBError::Internal(
478 "unknown probabilistic WAL kind".into(),
479 ))
480 }
481 };
482 if matches!(operation, PROB_WAL_OP_ADD | PROB_WAL_OP_DELETE) && !present {
483 return Err(RedDBError::Internal(format!(
484 "probabilistic WAL mutation has no base state for '{name}'"
485 )));
486 }
487 if operation == PROB_WAL_OP_DROP && !operands.is_empty() {
488 return Err(RedDBError::Internal(
489 "invalid probabilistic WAL drop".into(),
490 ));
491 }
492 let filter_add = (kind == PROB_WAL_KIND_FILTER && operation == PROB_WAL_OP_ADD)
493 .then(|| operands.first().cloned())
494 .flatten();
495 self.apply_replayed_probabilistic_delta(kind, operation, &name, operands)?;
496 if let Some(element) = filter_add {
497 if !self
498 .inner
499 .probabilistic
500 .filters
501 .read()
502 .get(&name)
503 .is_some_and(|filter| filter.contains(&element))
504 {
505 return Err(RedDBError::Internal(
506 "probabilistic WAL filter insert failed".into(),
507 ));
508 }
509 }
510 self.inner
511 .probabilistic
512 .replayed_names
513 .lock()
514 .insert((kind, name));
515 }
516 Ok(())
517 }
518
519 pub(crate) fn persist_offline_probabilistic_recovery(
521 &self,
522 deltas: Vec<(u8, u8, String, Vec<Vec<u8>>)>,
523 ) -> RedDBResult<()> {
524 let mut touched = std::collections::BTreeSet::new();
525 for (kind, _, name, _) in deltas {
526 if !matches!(
527 kind,
528 PROB_WAL_KIND_HLL | PROB_WAL_KIND_SKETCH | PROB_WAL_KIND_FILTER
529 ) {
530 return Err(RedDBError::Internal(
531 "unknown probabilistic WAL kind".into(),
532 ));
533 }
534 touched.insert((kind, name));
535 }
536 for (kind, name) in touched {
537 let (prefix, bytes) = match kind {
538 PROB_WAL_KIND_HLL => (
539 PROB_HLL_STATE_PREFIX,
540 self.inner
541 .probabilistic
542 .hlls
543 .read()
544 .get(&name)
545 .map(|h| h.as_bytes().to_vec()),
546 ),
547 PROB_WAL_KIND_SKETCH => (
548 PROB_SKETCH_STATE_PREFIX,
549 self.inner
550 .probabilistic
551 .sketches
552 .read()
553 .get(&name)
554 .map(|s| s.as_bytes()),
555 ),
556 PROB_WAL_KIND_FILTER => (
557 PROB_FILTER_STATE_PREFIX,
558 self.inner
559 .probabilistic
560 .filters
561 .read()
562 .get(&name)
563 .map(|f| f.as_bytes()),
564 ),
565 _ => unreachable!("kind validated above"),
566 };
567 if let Some(bytes) = bytes {
568 self.append_probabilistic_snapshot_delta(kind, &name, bytes.clone())?;
569 self.persist_probabilistic_blob(prefix, &name, &bytes)?;
570 } else {
571 self.append_probabilistic_drop_delta(kind, &name)?;
572 self.delete_probabilistic_blob(prefix, &name)?;
573 }
574 }
575 Ok(())
576 }
577
578 fn apply_replayed_probabilistic_delta(
579 &self,
580 kind: u8,
581 operation: u8,
582 name: &str,
583 operands: Vec<Vec<u8>>,
584 ) -> RedDBResult<()> {
585 match (kind, operation) {
586 (PROB_WAL_KIND_HLL, PROB_WAL_OP_SNAPSHOT) => {
587 let bytes = single_probabilistic_operand(operands, "HLL snapshot")?;
588 let hll = crate::storage::primitives::hyperloglog::HyperLogLog::from_bytes(bytes)
589 .ok_or_else(|| {
590 RedDBError::Internal(format!("invalid WAL HLL snapshot for '{name}'"))
591 })?;
592 let mut hlls =
593 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
594 hlls.insert(name.to_string(), hll);
595 }
596 (PROB_WAL_KIND_HLL, PROB_WAL_OP_ADD) => {
597 let mut hlls =
598 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
599 if let Some(hll) = hlls.get_mut(name) {
600 for element in operands {
601 hll.add(&element);
602 }
603 }
604 }
605 (PROB_WAL_KIND_HLL, PROB_WAL_OP_DROP) => {
606 let mut hlls =
607 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
608 hlls.remove(name);
609 }
610 (PROB_WAL_KIND_SKETCH, PROB_WAL_OP_SNAPSHOT) => {
611 let bytes = single_probabilistic_operand(operands, "SKETCH snapshot")?;
612 let sketch =
613 crate::storage::primitives::count_min_sketch::CountMinSketch::from_bytes(
614 &bytes,
615 )
616 .ok_or_else(|| {
617 RedDBError::Internal(format!("invalid WAL SKETCH snapshot for '{name}'"))
618 })?;
619 let mut sketches = probabilistic_write(
620 &self.inner.probabilistic.sketches,
621 "probabilistic sketch store",
622 );
623 sketches.insert(name.to_string(), sketch);
624 }
625 (PROB_WAL_KIND_SKETCH, PROB_WAL_OP_ADD) => {
626 let (element, count) = sketch_add_operands(operands)?;
627 let mut sketches = probabilistic_write(
628 &self.inner.probabilistic.sketches,
629 "probabilistic sketch store",
630 );
631 if let Some(sketch) = sketches.get_mut(name) {
632 sketch.add(&element, count);
633 }
634 }
635 (PROB_WAL_KIND_SKETCH, PROB_WAL_OP_DROP) => {
636 let mut sketches = probabilistic_write(
637 &self.inner.probabilistic.sketches,
638 "probabilistic sketch store",
639 );
640 sketches.remove(name);
641 }
642 (PROB_WAL_KIND_FILTER, PROB_WAL_OP_SNAPSHOT) => {
643 let bytes = single_probabilistic_operand(operands, "FILTER snapshot")?;
644 let filter =
645 crate::storage::primitives::cuckoo_filter::CuckooFilter::from_bytes(&bytes)
646 .ok_or_else(|| {
647 RedDBError::Internal(format!(
648 "invalid WAL FILTER snapshot for '{name}'"
649 ))
650 })?;
651 let mut filters = probabilistic_write(
652 &self.inner.probabilistic.filters,
653 "probabilistic filter store",
654 );
655 filters.insert(name.to_string(), filter);
656 }
657 (PROB_WAL_KIND_FILTER, PROB_WAL_OP_ADD) => {
658 let element = single_probabilistic_operand(operands, "FILTER ADD")?;
659 let mut filters = probabilistic_write(
660 &self.inner.probabilistic.filters,
661 "probabilistic filter store",
662 );
663 if let Some(filter) = filters.get_mut(name) {
664 let _ = filter.insert(&element);
665 }
666 }
667 (PROB_WAL_KIND_FILTER, PROB_WAL_OP_DELETE) => {
668 let element = single_probabilistic_operand(operands, "FILTER DELETE")?;
669 let mut filters = probabilistic_write(
670 &self.inner.probabilistic.filters,
671 "probabilistic filter store",
672 );
673 if let Some(filter) = filters.get_mut(name) {
674 let _ = filter.delete(&element);
675 }
676 }
677 (PROB_WAL_KIND_FILTER, PROB_WAL_OP_DROP) => {
678 let mut filters = probabilistic_write(
679 &self.inner.probabilistic.filters,
680 "probabilistic filter store",
681 );
682 filters.remove(name);
683 }
684 _ => {
685 return Err(RedDBError::Internal(format!(
686 "unknown probabilistic WAL delta kind={kind} operation={operation}"
687 )));
688 }
689 }
690 Ok(())
691 }
692
693 fn persist_probabilistic_encoding_marker(&self) -> RedDBResult<()> {
694 self.compact_config_key(PROB_ENCODING_MARKER_KEY)?;
695 self.insert_config_value(
696 PROB_ENCODING_MARKER_KEY,
697 Value::text(PROB_ENCODING_RAW_V1.to_string()),
698 )
699 }
700
701 fn compact_config_key(&self, key: &str) -> RedDBResult<()> {
702 let store = self.inner.db.store();
703 let Some(manager) = store.get_collection("red.config") else {
704 return Ok(());
705 };
706 let mut ids = Vec::new();
707 manager.try_for_each_entity(usize::MAX, |entity| {
708 if let EntityData::Row(row) = &entity.data {
709 if matches!(row.named.as_ref().and_then(|named| named.get("key")), Some(Value::Text(candidate)) if candidate.as_ref() == key) {
710 ids.push(entity.id);
711 }
712 }
713 true
714 })?;
715 for id in ids {
716 store.delete("red.config", id)?;
717 }
718 Ok(())
719 }
720
721 fn insert_config_value(&self, key: &str, value: Value) -> RedDBResult<()> {
722 let store = self.inner.db.store();
723 let _ = store.get_or_create_collection("red.config");
724 let entity = UnifiedEntity::new(
725 EntityId::new(0),
726 EntityKind::TableRow {
727 table: std::sync::Arc::from("red.config"),
728 row_id: 0,
729 },
730 EntityData::Row(crate::storage::RowData {
731 columns: Vec::new(),
732 named: Some(
733 [
734 ("key".to_string(), Value::text(key.to_string())),
735 ("value".to_string(), value),
736 ]
737 .into_iter()
738 .collect(),
739 ),
740 schema: None,
741 }),
742 );
743 store
744 .insert_auto("red.config", entity)
745 .map_err(|err| RedDBError::Internal(err.to_string()))?;
746 Ok(())
747 }
748
749 fn create_probabilistic_catalog_entry(
750 &self,
751 name: &str,
752 model: crate::catalog::CollectionModel,
753 ) -> RedDBResult<()> {
754 let store = self.inner.db.store();
755 store
756 .create_collection(name)
757 .map_err(|err| RedDBError::Internal(err.to_string()))?;
758 self.inner
759 .db
760 .save_collection_contract(probabilistic_collection_contract(name, model))
761 .map_err(|err| RedDBError::Internal(err.to_string()))?;
762 if let Some(tenant_id) = crate::runtime::impl_core::current_tenant() {
763 store.set_config_tree(
764 &format!("collection_tenants.{name}"),
765 &crate::serde_json::Value::String(tenant_id),
766 );
767 }
768 self.inner
769 .db
770 .persist_metadata()
771 .map_err(|err| RedDBError::Internal(err.to_string()))?;
772 self.invalidate_result_cache();
773 Ok(())
774 }
775
776 fn drop_probabilistic_catalog_entry(&self, name: &str) -> RedDBResult<()> {
777 let store = self.inner.db.store();
778 if store.get_collection(name).is_some() {
779 store
780 .drop_collection(name)
781 .map_err(|err| RedDBError::Internal(err.to_string()))?;
782 }
783 self.inner
784 .db
785 .remove_collection_contract(name)
786 .map_err(|err| RedDBError::Internal(err.to_string()))?;
787 self.inner
788 .db
789 .persist_metadata()
790 .map_err(|err| RedDBError::Internal(err.to_string()))?;
791 self.invalidate_result_cache();
792 Ok(())
793 }
794
795 pub(crate) fn execute_probabilistic_select(
796 &self,
797 query: &TableQuery,
798 ) -> RedDBResult<Option<UnifiedResult>> {
799 let projections = reddb_rql::sql_lowering::effective_table_projections(query);
800 let mut read_projections = Vec::new();
801 for projection in &projections {
802 if let Some(read_projection) =
803 parse_probabilistic_read_projection(projection, read_projections.len())?
804 {
805 read_projections.push(read_projection);
806 }
807 }
808
809 let Some(actual_model) = self
810 .inner
811 .db
812 .collection_contract(&query.table)
813 .map(|contract| contract.declared_model)
814 else {
815 return if read_projections.is_empty() {
816 Ok(None)
817 } else {
818 Err(RedDBError::NotFound(format!(
819 "probabilistic collection '{}' not found",
820 query.table
821 )))
822 };
823 };
824
825 let is_probabilistic_model = matches!(
826 actual_model,
827 crate::catalog::CollectionModel::Hll
828 | crate::catalog::CollectionModel::Sketch
829 | crate::catalog::CollectionModel::Filter
830 );
831 if read_projections.is_empty() {
832 return if is_probabilistic_model {
833 Err(RedDBError::Query(format!(
834 "probabilistic collection '{}' supports SELECT CARDINALITY, FREQ(...), or CONTAINS(...) read forms",
835 query.table
836 )))
837 } else {
838 Ok(None)
839 };
840 }
841
842 validate_probabilistic_read_model(&query.table, actual_model, &read_projections)?;
843 let (columns, record) =
844 self.materialize_probabilistic_select_row(&query.table, &read_projections)?;
845 let mut result = UnifiedResult::with_columns(columns);
846 if probabilistic_select_row_visible(self, query, &record) {
847 result.push(record);
848 }
849 Ok(Some(result))
850 }
851
852 pub fn execute_probabilistic_command(
853 &self,
854 raw_query: &str,
855 cmd: &ProbabilisticCommand,
856 ) -> RedDBResult<RuntimeQueryResult> {
857 let is_mutation = matches!(
861 cmd,
862 ProbabilisticCommand::CreateHll { .. }
863 | ProbabilisticCommand::HllAdd { .. }
864 | ProbabilisticCommand::HllMerge { .. }
865 | ProbabilisticCommand::DropHll { .. }
866 | ProbabilisticCommand::CreateSketch { .. }
867 | ProbabilisticCommand::SketchAdd { .. }
868 | ProbabilisticCommand::SketchMerge { .. }
869 | ProbabilisticCommand::DropSketch { .. }
870 | ProbabilisticCommand::CreateFilter { .. }
871 | ProbabilisticCommand::FilterAdd { .. }
872 | ProbabilisticCommand::FilterDelete { .. }
873 | ProbabilisticCommand::DropFilter { .. }
874 );
875 if is_mutation {
876 self.check_write(crate::runtime::write_gate::WriteKind::Dml)?;
877 }
878 let store = self.inner.db.store();
881 let _publication = is_mutation.then(|| store.vault_publication_lock.lock());
882 if let ProbabilisticCommand::CreateHll { name, .. }
883 | ProbabilisticCommand::CreateSketch { name, .. }
884 | ProbabilisticCommand::CreateFilter { name, .. } = cmd
885 {
886 if super::impl_ddl::is_system_schema_name(name) {
887 return Err(RedDBError::Query("system schema is read-only".to_string()));
888 }
889 }
890 match cmd {
891 ProbabilisticCommand::CreateHll {
893 name,
894 precision,
895 if_not_exists,
896 } => {
897 let mut hlls =
898 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
899 if hlls.contains_key(name) {
900 if *if_not_exists {
901 return Ok(RuntimeQueryResult::ok_message(
902 raw_query.to_string(),
903 &format!("HLL '{}' already exists", name),
904 "create",
905 ));
906 }
907 return Err(RedDBError::Query(format!("HLL '{}' already exists", name)));
908 }
909 let hll = crate::storage::primitives::hyperloglog::HyperLogLog::with_precision(
910 *precision,
911 )
912 .ok_or_else(|| {
913 RedDBError::Query(format!(
914 "HLL precision must be between 4 and 18, got {precision}"
915 ))
916 })?;
917 self.create_probabilistic_catalog_entry(
918 name,
919 crate::catalog::CollectionModel::Hll,
920 )?;
921 self.append_probabilistic_snapshot_delta(
922 PROB_WAL_KIND_HLL,
923 name,
924 hll.as_bytes().to_vec(),
925 )?;
926 self.persist_probabilistic_blob(PROB_HLL_STATE_PREFIX, name, hll.as_bytes())?;
927 hlls.insert(name.clone(), hll);
928 Ok(RuntimeQueryResult::ok_message(
929 raw_query.to_string(),
930 &format!("HLL '{}' created", name),
931 "create",
932 ))
933 }
934 ProbabilisticCommand::HllAdd { name, elements } => {
935 let mut hlls =
936 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
937 let hll = hlls
938 .get_mut(name)
939 .ok_or_else(|| RedDBError::NotFound(format!("HLL '{}' not found", name)))?;
940 for elem in elements {
941 hll.add(elem.as_bytes());
942 }
943 self.append_probabilistic_add_delta(
944 PROB_WAL_KIND_HLL,
945 name,
946 elements
947 .iter()
948 .map(|element| element.as_bytes().to_vec())
949 .collect(),
950 )?;
951 Ok(RuntimeQueryResult::ok_message(
952 raw_query.to_string(),
953 &format!("{} element(s) added to HLL '{}'", elements.len(), name),
954 "insert",
955 ))
956 }
957 ProbabilisticCommand::HllCount { names } => {
958 let hlls =
959 probabilistic_read(&self.inner.probabilistic.hlls, "probabilistic HLL store");
960 if names.is_empty() {
961 return Err(RedDBError::Query(
962 "HLL COUNT requires at least one HLL name".to_string(),
963 ));
964 }
965 if names.len() == 1 {
966 let hll = hlls.get(&names[0]).ok_or_else(|| {
967 RedDBError::NotFound(format!("HLL '{}' not found", names[0]))
968 })?;
969 let count = hll.count();
970 let mut result = UnifiedResult::with_columns(vec!["count".into()]);
971 let mut record = UnifiedRecord::new();
972 record.set("count", Value::UnsignedInteger(count));
973 result.push(record);
974 Ok(RuntimeQueryResult {
975 query: raw_query.to_string(),
976 mode: QueryMode::Sql,
977 statement: "hll_count",
978 engine: "runtime-probabilistic",
979 result,
980 affected_rows: 0,
981 statement_type: "select",
982 bookmark: None,
983 notice: None,
984 })
985 } else {
986 let first_name = &names[0];
988 let first = hlls.get(first_name).ok_or_else(|| {
989 RedDBError::NotFound(format!("HLL '{}' not found", first_name))
990 })?;
991 let expected_precision = first.precision();
992 let mut merged =
993 crate::storage::primitives::hyperloglog::HyperLogLog::with_precision(
994 expected_precision,
995 )
996 .expect("loaded HLL precision is valid");
997 for name in names {
998 let hll = hlls.get(name).ok_or_else(|| {
999 RedDBError::NotFound(format!("HLL '{}' not found", name))
1000 })?;
1001 if hll.precision() != expected_precision {
1002 return Err(hll_precision_mismatch_error(
1003 "COUNT",
1004 first_name,
1005 expected_precision,
1006 name,
1007 hll.precision(),
1008 ));
1009 }
1010 merged.merge(hll);
1011 }
1012 let count = merged.count();
1013 let mut result = UnifiedResult::with_columns(vec!["count".into()]);
1014 let mut record = UnifiedRecord::new();
1015 record.set("count", Value::UnsignedInteger(count));
1016 result.push(record);
1017 Ok(RuntimeQueryResult {
1018 query: raw_query.to_string(),
1019 mode: QueryMode::Sql,
1020 statement: "hll_count",
1021 engine: "runtime-probabilistic",
1022 result,
1023 affected_rows: 0,
1024 statement_type: "select",
1025 bookmark: None,
1026 notice: None,
1027 })
1028 }
1029 }
1030 ProbabilisticCommand::HllMerge { dest, sources } => {
1031 let mut hlls =
1032 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
1033 let first_src = sources.first().ok_or_else(|| {
1034 RedDBError::Query("HLL MERGE requires at least one source HLL".to_string())
1035 })?;
1036 let first = hlls.get(first_src).ok_or_else(|| {
1037 RedDBError::NotFound(format!("HLL '{}' not found", first_src))
1038 })?;
1039 let expected_precision = first.precision();
1040 let mut merged =
1041 crate::storage::primitives::hyperloglog::HyperLogLog::with_precision(
1042 expected_precision,
1043 )
1044 .expect("loaded HLL precision is valid");
1045 for src in sources {
1046 let hll = hlls
1047 .get(src)
1048 .ok_or_else(|| RedDBError::NotFound(format!("HLL '{}' not found", src)))?;
1049 if hll.precision() != expected_precision {
1050 return Err(hll_precision_mismatch_error(
1051 "MERGE",
1052 first_src,
1053 expected_precision,
1054 src,
1055 hll.precision(),
1056 ));
1057 }
1058 merged.merge(hll);
1059 }
1060 self.append_probabilistic_snapshot_delta(
1061 PROB_WAL_KIND_HLL,
1062 dest,
1063 merged.as_bytes().to_vec(),
1064 )?;
1065 self.persist_probabilistic_blob(PROB_HLL_STATE_PREFIX, dest, merged.as_bytes())?;
1066 hlls.insert(dest.clone(), merged);
1067 Ok(RuntimeQueryResult::ok_message(
1068 raw_query.to_string(),
1069 &format!(
1070 "HLL '{}' created from merge of {}",
1071 dest,
1072 sources.join(", ")
1073 ),
1074 "create",
1075 ))
1076 }
1077 ProbabilisticCommand::HllInfo { name } => {
1078 let hlls =
1079 probabilistic_read(&self.inner.probabilistic.hlls, "probabilistic HLL store");
1080 let hll = hlls
1081 .get(name)
1082 .ok_or_else(|| RedDBError::NotFound(format!("HLL '{}' not found", name)))?;
1083 let mut result = UnifiedResult::with_columns(vec![
1084 "name".into(),
1085 "precision".into(),
1086 "count".into(),
1087 "memory_bytes".into(),
1088 ]);
1089 let mut record = UnifiedRecord::new();
1090 record.set("name", Value::text(name.clone()));
1091 record.set("precision", Value::UnsignedInteger(hll.precision() as u64));
1092 record.set("count", Value::UnsignedInteger(hll.count()));
1093 record.set(
1094 "memory_bytes",
1095 Value::UnsignedInteger(hll.memory_bytes() as u64),
1096 );
1097 result.push(record);
1098 Ok(RuntimeQueryResult {
1099 query: raw_query.to_string(),
1100 mode: QueryMode::Sql,
1101 statement: "hll_info",
1102 engine: "runtime-probabilistic",
1103 result,
1104 affected_rows: 0,
1105 statement_type: "select",
1106 bookmark: None,
1107 notice: None,
1108 })
1109 }
1110 ProbabilisticCommand::DropHll { name, if_exists } => {
1111 let mut hlls =
1112 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
1113 if hlls.remove(name).is_none() {
1114 if *if_exists {
1115 return Ok(RuntimeQueryResult::ok_message(
1116 raw_query.to_string(),
1117 &format!("HLL '{}' does not exist", name),
1118 "drop",
1119 ));
1120 }
1121 return Err(RedDBError::NotFound(format!("HLL '{}' not found", name)));
1122 }
1123 self.inner
1124 .probabilistic
1125 .replayed_names
1126 .lock()
1127 .insert((PROB_WAL_KIND_HLL, name.clone()));
1128 self.drop_probabilistic_catalog_entry(name)?;
1129 self.append_probabilistic_drop_delta(PROB_WAL_KIND_HLL, name)?;
1130 self.delete_probabilistic_blob(PROB_HLL_STATE_PREFIX, name)?;
1131 Ok(RuntimeQueryResult::ok_message(
1132 raw_query.to_string(),
1133 &format!("HLL '{}' dropped", name),
1134 "drop",
1135 ))
1136 }
1137
1138 ProbabilisticCommand::CreateSketch {
1140 name,
1141 width,
1142 depth,
1143 if_not_exists,
1144 } => {
1145 let mut sketches = probabilistic_write(
1146 &self.inner.probabilistic.sketches,
1147 "probabilistic sketch store",
1148 );
1149 if sketches.contains_key(name) {
1150 if *if_not_exists {
1151 return Ok(RuntimeQueryResult::ok_message(
1152 raw_query.to_string(),
1153 &format!("SKETCH '{}' already exists", name),
1154 "create",
1155 ));
1156 }
1157 return Err(RedDBError::Query(format!(
1158 "SKETCH '{}' already exists",
1159 name
1160 )));
1161 }
1162 self.create_probabilistic_catalog_entry(
1163 name,
1164 crate::catalog::CollectionModel::Sketch,
1165 )?;
1166 let sketch = crate::storage::primitives::count_min_sketch::CountMinSketch::new(
1167 *width, *depth,
1168 );
1169 self.append_probabilistic_snapshot_delta(
1170 PROB_WAL_KIND_SKETCH,
1171 name,
1172 sketch.as_bytes(),
1173 )?;
1174 self.persist_probabilistic_blob(
1175 PROB_SKETCH_STATE_PREFIX,
1176 name,
1177 &sketch.as_bytes(),
1178 )?;
1179 sketches.insert(name.clone(), sketch);
1180 Ok(RuntimeQueryResult::ok_message(
1181 raw_query.to_string(),
1182 &format!(
1183 "SKETCH '{}' created (width={}, depth={})",
1184 name, width, depth
1185 ),
1186 "create",
1187 ))
1188 }
1189 ProbabilisticCommand::SketchAdd {
1190 name,
1191 element,
1192 count,
1193 } => {
1194 let mut sketches = probabilistic_write(
1195 &self.inner.probabilistic.sketches,
1196 "probabilistic sketch store",
1197 );
1198 let sketch = sketches
1199 .get_mut(name)
1200 .ok_or_else(|| RedDBError::NotFound(format!("SKETCH '{}' not found", name)))?;
1201 sketch.add(element.as_bytes(), *count);
1202 self.append_probabilistic_add_delta(
1203 PROB_WAL_KIND_SKETCH,
1204 name,
1205 vec![element.as_bytes().to_vec(), count.to_le_bytes().to_vec()],
1206 )?;
1207 Ok(RuntimeQueryResult::ok_message(
1208 raw_query.to_string(),
1209 &format!("added {} to SKETCH '{}'", count, name),
1210 "insert",
1211 ))
1212 }
1213 ProbabilisticCommand::SketchCount { name, element } => {
1214 let sketches = probabilistic_read(
1215 &self.inner.probabilistic.sketches,
1216 "probabilistic sketch store",
1217 );
1218 let sketch = sketches
1219 .get(name)
1220 .ok_or_else(|| RedDBError::NotFound(format!("SKETCH '{}' not found", name)))?;
1221 let estimate = sketch.estimate(element.as_bytes());
1222 let mut result = UnifiedResult::with_columns(vec!["estimate".into()]);
1223 let mut record = UnifiedRecord::new();
1224 record.set("estimate", Value::UnsignedInteger(estimate));
1225 result.push(record);
1226 Ok(RuntimeQueryResult {
1227 query: raw_query.to_string(),
1228 mode: QueryMode::Sql,
1229 statement: "sketch_count",
1230 engine: "runtime-probabilistic",
1231 result,
1232 affected_rows: 0,
1233 statement_type: "select",
1234 bookmark: None,
1235 notice: None,
1236 })
1237 }
1238 ProbabilisticCommand::SketchMerge { dest, sources } => {
1239 let mut sketches = probabilistic_write(
1240 &self.inner.probabilistic.sketches,
1241 "probabilistic sketch store",
1242 );
1243 let first_src = sketches.get(&sources[0]).ok_or_else(|| {
1244 RedDBError::NotFound(format!("SKETCH '{}' not found", sources[0]))
1245 })?;
1246 let mut merged = crate::storage::primitives::count_min_sketch::CountMinSketch::new(
1247 first_src.width(),
1248 first_src.depth(),
1249 );
1250 for src in sources {
1251 let sketch = sketches.get(src).ok_or_else(|| {
1252 RedDBError::NotFound(format!("SKETCH '{}' not found", src))
1253 })?;
1254 if !merged.merge(sketch) {
1255 return Err(RedDBError::Query(format!(
1256 "SKETCH '{}' has incompatible dimensions",
1257 src
1258 )));
1259 }
1260 }
1261 self.append_probabilistic_snapshot_delta(
1262 PROB_WAL_KIND_SKETCH,
1263 dest,
1264 merged.as_bytes(),
1265 )?;
1266 self.persist_probabilistic_blob(
1267 PROB_SKETCH_STATE_PREFIX,
1268 dest,
1269 &merged.as_bytes(),
1270 )?;
1271 sketches.insert(dest.clone(), merged);
1272 Ok(RuntimeQueryResult::ok_message(
1273 raw_query.to_string(),
1274 &format!(
1275 "SKETCH '{}' created from merge of {}",
1276 dest,
1277 sources.join(", ")
1278 ),
1279 "create",
1280 ))
1281 }
1282 ProbabilisticCommand::SketchInfo { name } => {
1283 let sketches = probabilistic_read(
1284 &self.inner.probabilistic.sketches,
1285 "probabilistic sketch store",
1286 );
1287 let sketch = sketches
1288 .get(name)
1289 .ok_or_else(|| RedDBError::NotFound(format!("SKETCH '{}' not found", name)))?;
1290 let mut result = UnifiedResult::with_columns(vec![
1291 "name".into(),
1292 "width".into(),
1293 "depth".into(),
1294 "total".into(),
1295 "memory_bytes".into(),
1296 ]);
1297 let mut record = UnifiedRecord::new();
1298 record.set("name", Value::text(name.clone()));
1299 record.set("width", Value::UnsignedInteger(sketch.width() as u64));
1300 record.set("depth", Value::UnsignedInteger(sketch.depth() as u64));
1301 record.set("total", Value::UnsignedInteger(sketch.total()));
1302 record.set(
1303 "memory_bytes",
1304 Value::UnsignedInteger(sketch.memory_bytes() as u64),
1305 );
1306 result.push(record);
1307 Ok(RuntimeQueryResult {
1308 query: raw_query.to_string(),
1309 mode: QueryMode::Sql,
1310 statement: "sketch_info",
1311 engine: "runtime-probabilistic",
1312 result,
1313 affected_rows: 0,
1314 statement_type: "select",
1315 bookmark: None,
1316 notice: None,
1317 })
1318 }
1319 ProbabilisticCommand::DropSketch { name, if_exists } => {
1320 let mut sketches = probabilistic_write(
1321 &self.inner.probabilistic.sketches,
1322 "probabilistic sketch store",
1323 );
1324 if sketches.remove(name).is_none() {
1325 if *if_exists {
1326 return Ok(RuntimeQueryResult::ok_message(
1327 raw_query.to_string(),
1328 &format!("SKETCH '{}' does not exist", name),
1329 "drop",
1330 ));
1331 }
1332 return Err(RedDBError::NotFound(format!("SKETCH '{}' not found", name)));
1333 }
1334 self.inner
1335 .probabilistic
1336 .replayed_names
1337 .lock()
1338 .insert((PROB_WAL_KIND_SKETCH, name.clone()));
1339 self.drop_probabilistic_catalog_entry(name)?;
1340 self.append_probabilistic_drop_delta(PROB_WAL_KIND_SKETCH, name)?;
1341 self.delete_probabilistic_blob(PROB_SKETCH_STATE_PREFIX, name)?;
1342 Ok(RuntimeQueryResult::ok_message(
1343 raw_query.to_string(),
1344 &format!("SKETCH '{}' dropped", name),
1345 "drop",
1346 ))
1347 }
1348
1349 ProbabilisticCommand::CreateFilter {
1351 name,
1352 capacity,
1353 if_not_exists,
1354 } => {
1355 let mut filters = probabilistic_write(
1356 &self.inner.probabilistic.filters,
1357 "probabilistic filter store",
1358 );
1359 if filters.contains_key(name) {
1360 if *if_not_exists {
1361 return Ok(RuntimeQueryResult::ok_message(
1362 raw_query.to_string(),
1363 &format!("FILTER '{}' already exists", name),
1364 "create",
1365 ));
1366 }
1367 return Err(RedDBError::Query(format!(
1368 "FILTER '{}' already exists",
1369 name
1370 )));
1371 }
1372 self.create_probabilistic_catalog_entry(
1373 name,
1374 crate::catalog::CollectionModel::Filter,
1375 )?;
1376 let filter =
1377 crate::storage::primitives::cuckoo_filter::CuckooFilter::new(*capacity);
1378 self.append_probabilistic_snapshot_delta(
1379 PROB_WAL_KIND_FILTER,
1380 name,
1381 filter.as_bytes(),
1382 )?;
1383 self.persist_probabilistic_blob(
1384 PROB_FILTER_STATE_PREFIX,
1385 name,
1386 &filter.as_bytes(),
1387 )?;
1388 filters.insert(name.clone(), filter);
1389 Ok(RuntimeQueryResult::ok_message(
1390 raw_query.to_string(),
1391 &format!("FILTER '{}' created (capacity={})", name, capacity),
1392 "create",
1393 ))
1394 }
1395 ProbabilisticCommand::FilterAdd { name, element } => {
1396 let mut filters = probabilistic_write(
1397 &self.inner.probabilistic.filters,
1398 "probabilistic filter store",
1399 );
1400 let filter = filters
1401 .get_mut(name)
1402 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
1403 if !filter.insert(element.as_bytes()) {
1404 return Err(RedDBError::Query(format!("FILTER '{}' is full", name)));
1405 }
1406 self.append_probabilistic_add_delta(
1407 PROB_WAL_KIND_FILTER,
1408 name,
1409 vec![element.as_bytes().to_vec()],
1410 )?;
1411 Ok(RuntimeQueryResult::ok_message(
1412 raw_query.to_string(),
1413 &format!("element added to FILTER '{}'", name),
1414 "insert",
1415 ))
1416 }
1417 ProbabilisticCommand::FilterCheck { name, element } => {
1418 let filters = probabilistic_read(
1419 &self.inner.probabilistic.filters,
1420 "probabilistic filter store",
1421 );
1422 let filter = filters
1423 .get(name)
1424 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
1425 let exists = filter.contains(element.as_bytes());
1426 let mut result = UnifiedResult::with_columns(vec!["exists".into()]);
1427 let mut record = UnifiedRecord::new();
1428 record.set("exists", Value::Boolean(exists));
1429 result.push(record);
1430 Ok(RuntimeQueryResult {
1431 query: raw_query.to_string(),
1432 mode: QueryMode::Sql,
1433 statement: "filter_check",
1434 engine: "runtime-probabilistic",
1435 result,
1436 affected_rows: 0,
1437 statement_type: "select",
1438 bookmark: None,
1439 notice: None,
1440 })
1441 }
1442 ProbabilisticCommand::FilterDelete { name, element } => {
1443 let mut filters = probabilistic_write(
1444 &self.inner.probabilistic.filters,
1445 "probabilistic filter store",
1446 );
1447 let filter = filters
1448 .get_mut(name)
1449 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
1450 let removed = filter.delete(element.as_bytes());
1451 self.append_probabilistic_delete_delta(
1452 PROB_WAL_KIND_FILTER,
1453 name,
1454 vec![element.as_bytes().to_vec()],
1455 )?;
1456 Ok(RuntimeQueryResult::ok_message(
1457 raw_query.to_string(),
1458 &format!(
1459 "element {} from FILTER '{}'",
1460 if removed { "deleted" } else { "not found in" },
1461 name
1462 ),
1463 "delete",
1464 ))
1465 }
1466 ProbabilisticCommand::FilterCount { name } => {
1467 let filters = probabilistic_read(
1468 &self.inner.probabilistic.filters,
1469 "probabilistic filter store",
1470 );
1471 let filter = filters
1472 .get(name)
1473 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
1474 let mut result = UnifiedResult::with_columns(vec!["count".into()]);
1475 let mut record = UnifiedRecord::new();
1476 record.set("count", Value::UnsignedInteger(filter.count() as u64));
1477 result.push(record);
1478 Ok(RuntimeQueryResult {
1479 query: raw_query.to_string(),
1480 mode: QueryMode::Sql,
1481 statement: "filter_count",
1482 engine: "runtime-probabilistic",
1483 result,
1484 affected_rows: 0,
1485 statement_type: "select",
1486 bookmark: None,
1487 notice: None,
1488 })
1489 }
1490 ProbabilisticCommand::FilterInfo { name } => {
1491 let filters = probabilistic_read(
1492 &self.inner.probabilistic.filters,
1493 "probabilistic filter store",
1494 );
1495 let filter = filters
1496 .get(name)
1497 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
1498 let mut result = UnifiedResult::with_columns(vec![
1499 "name".into(),
1500 "count".into(),
1501 "load_factor".into(),
1502 "memory_bytes".into(),
1503 ]);
1504 let mut record = UnifiedRecord::new();
1505 record.set("name", Value::text(name.clone()));
1506 record.set("count", Value::UnsignedInteger(filter.count() as u64));
1507 record.set("load_factor", Value::Float(filter.load_factor()));
1508 record.set(
1509 "memory_bytes",
1510 Value::UnsignedInteger(filter.memory_bytes() as u64),
1511 );
1512 result.push(record);
1513 Ok(RuntimeQueryResult {
1514 query: raw_query.to_string(),
1515 mode: QueryMode::Sql,
1516 statement: "filter_info",
1517 engine: "runtime-probabilistic",
1518 result,
1519 affected_rows: 0,
1520 statement_type: "select",
1521 bookmark: None,
1522 notice: None,
1523 })
1524 }
1525 ProbabilisticCommand::DropFilter { name, if_exists } => {
1526 let mut filters = probabilistic_write(
1527 &self.inner.probabilistic.filters,
1528 "probabilistic filter store",
1529 );
1530 if filters.remove(name).is_none() {
1531 if *if_exists {
1532 return Ok(RuntimeQueryResult::ok_message(
1533 raw_query.to_string(),
1534 &format!("FILTER '{}' does not exist", name),
1535 "drop",
1536 ));
1537 }
1538 return Err(RedDBError::NotFound(format!("FILTER '{}' not found", name)));
1539 }
1540 self.inner
1541 .probabilistic
1542 .replayed_names
1543 .lock()
1544 .insert((PROB_WAL_KIND_FILTER, name.clone()));
1545 self.drop_probabilistic_catalog_entry(name)?;
1546 self.append_probabilistic_drop_delta(PROB_WAL_KIND_FILTER, name)?;
1547 self.delete_probabilistic_blob(PROB_FILTER_STATE_PREFIX, name)?;
1548 Ok(RuntimeQueryResult::ok_message(
1549 raw_query.to_string(),
1550 &format!("FILTER '{}' dropped", name),
1551 "drop",
1552 ))
1553 }
1554 }
1555 }
1556}
1557
1558fn parse_probabilistic_read_projection(
1559 projection: &Projection,
1560 index: usize,
1561) -> RedDBResult<Option<ProbabilisticReadProjection>> {
1562 if let Some(column) = projection_unqualified_column(projection) {
1563 if column.eq_ignore_ascii_case("CARDINALITY") {
1564 return Ok(Some(ProbabilisticReadProjection::Cardinality {
1565 label: probabilistic_projection_label(projection, "cardinality", index),
1566 }));
1567 }
1568 }
1569
1570 let Some((function, args)) = projection_function(projection) else {
1571 return Ok(None);
1572 };
1573 if function.eq_ignore_ascii_case("FREQ") {
1574 let element = projection_single_text_arg(function, args)?;
1575 return Ok(Some(ProbabilisticReadProjection::Freq {
1576 element,
1577 label: probabilistic_projection_label(projection, "freq", index),
1578 }));
1579 }
1580 if function.eq_ignore_ascii_case("CONTAINS") && args.len() == 1 {
1581 let element = projection_single_text_arg(function, args)?;
1582 return Ok(Some(ProbabilisticReadProjection::Contains {
1583 element,
1584 label: probabilistic_projection_label(projection, "contains", index),
1585 }));
1586 }
1587
1588 Ok(None)
1589}
1590
1591fn single_probabilistic_operand(mut operands: Vec<Vec<u8>>, context: &str) -> RedDBResult<Vec<u8>> {
1592 if operands.len() != 1 {
1593 return Err(RedDBError::Internal(format!(
1594 "{context} WAL delta expected one operand, got {}",
1595 operands.len()
1596 )));
1597 }
1598 Ok(operands.remove(0))
1599}
1600
1601fn sketch_add_operands(mut operands: Vec<Vec<u8>>) -> RedDBResult<(Vec<u8>, u64)> {
1602 if operands.len() != 2 {
1603 return Err(RedDBError::Internal(format!(
1604 "SKETCH ADD WAL delta expected two operands, got {}",
1605 operands.len()
1606 )));
1607 }
1608 let count_bytes = operands.remove(1);
1609 let count_array: [u8; 8] = count_bytes.as_slice().try_into().map_err(|_| {
1610 RedDBError::Internal(format!(
1611 "SKETCH ADD WAL delta count expected 8 bytes, got {}",
1612 count_bytes.len()
1613 ))
1614 })?;
1615 Ok((operands.remove(0), u64::from_le_bytes(count_array)))
1616}
1617
1618fn validate_probabilistic_read_model(
1619 collection: &str,
1620 actual_model: crate::catalog::CollectionModel,
1621 projections: &[ProbabilisticReadProjection],
1622) -> RedDBResult<()> {
1623 for projection in projections {
1624 let expected_model = match projection {
1625 ProbabilisticReadProjection::Cardinality { .. } => crate::catalog::CollectionModel::Hll,
1626 ProbabilisticReadProjection::Freq { .. } => crate::catalog::CollectionModel::Sketch,
1627 ProbabilisticReadProjection::Contains { .. } => crate::catalog::CollectionModel::Filter,
1628 };
1629 if actual_model != expected_model {
1630 return Err(RedDBError::Query(format!(
1631 "{} is only supported for {} collections; '{}' is {}",
1632 probabilistic_projection_form(projection),
1633 crate::runtime::ddl::polymorphic_resolver::model_name(expected_model),
1634 collection,
1635 crate::runtime::ddl::polymorphic_resolver::model_name(actual_model)
1636 )));
1637 }
1638 }
1639 Ok(())
1640}
1641
1642impl RedDBRuntime {
1643 fn materialize_probabilistic_select_row(
1644 &self,
1645 collection: &str,
1646 projections: &[ProbabilisticReadProjection],
1647 ) -> RedDBResult<(Vec<String>, UnifiedRecord)> {
1648 let mut columns = Vec::with_capacity(projections.len());
1649 let mut record = UnifiedRecord::new();
1650 for projection in projections {
1651 match projection {
1652 ProbabilisticReadProjection::Cardinality { label } => {
1653 let hlls = probabilistic_read(
1654 &self.inner.probabilistic.hlls,
1655 "probabilistic HLL store",
1656 );
1657 let hll = hlls.get(collection).ok_or_else(|| {
1658 RedDBError::NotFound(format!("HLL '{}' not found", collection))
1659 })?;
1660 columns.push(label.clone());
1661 record.set(label, Value::UnsignedInteger(hll.count()));
1662 }
1663 ProbabilisticReadProjection::Freq { element, label } => {
1664 let sketches = probabilistic_read(
1665 &self.inner.probabilistic.sketches,
1666 "probabilistic sketch store",
1667 );
1668 let sketch = sketches.get(collection).ok_or_else(|| {
1669 RedDBError::NotFound(format!("SKETCH '{}' not found", collection))
1670 })?;
1671 columns.push(label.clone());
1672 record.set(
1673 label,
1674 Value::UnsignedInteger(sketch.estimate(element.as_bytes())),
1675 );
1676 }
1677 ProbabilisticReadProjection::Contains { element, label } => {
1678 let filters = probabilistic_read(
1679 &self.inner.probabilistic.filters,
1680 "probabilistic filter store",
1681 );
1682 let filter = filters.get(collection).ok_or_else(|| {
1683 RedDBError::NotFound(format!("FILTER '{}' not found", collection))
1684 })?;
1685 columns.push(label.clone());
1686 record.set(label, Value::Boolean(filter.contains(element.as_bytes())));
1687 }
1688 }
1689 }
1690 Ok((columns, record))
1691 }
1692}
1693
1694fn probabilistic_select_row_visible(
1695 runtime: &RedDBRuntime,
1696 query: &TableQuery,
1697 record: &UnifiedRecord,
1698) -> bool {
1699 if query.limit == Some(0) || query.offset.is_some_and(|offset| offset > 0) {
1700 return false;
1701 }
1702 let table_name = query.table.as_str();
1703 let table_alias = query.alias.as_deref().unwrap_or(table_name);
1704 reddb_rql::sql_lowering::effective_table_filter(query).is_none_or(|filter| {
1705 super::join_filter::evaluate_runtime_filter_with_db(
1706 Some(&runtime.inner.db),
1707 record,
1708 &filter,
1709 Some(table_name),
1710 Some(table_alias),
1711 )
1712 })
1713}
1714
1715fn projection_unqualified_column(projection: &Projection) -> Option<&str> {
1716 match projection {
1717 Projection::Field(FieldRef::TableColumn { table, column }, _) if table.is_empty() => {
1718 Some(column.as_str())
1719 }
1720 Projection::Column(column) => Some(column.as_str()),
1721 Projection::Alias(column, _) => Some(column.as_str()),
1722 _ => None,
1723 }
1724}
1725
1726fn projection_function(projection: &Projection) -> Option<(&str, &[Projection])> {
1727 match projection {
1728 Projection::Function(name, args) => {
1729 let function = name.split_once(':').map(|(name, _)| name).unwrap_or(name);
1730 Some((function, args.as_slice()))
1731 }
1732 _ => None,
1733 }
1734}
1735
1736fn projection_single_text_arg(function: &str, args: &[Projection]) -> RedDBResult<String> {
1737 if args.len() != 1 {
1738 return Err(RedDBError::Query(format!(
1739 "{function}(...) expects exactly one string literal"
1740 )));
1741 }
1742 match &args[0] {
1743 Projection::Column(column) => column
1744 .strip_prefix("LIT:")
1745 .map(ToString::to_string)
1746 .ok_or_else(|| {
1747 RedDBError::Query(format!("{function}(...) expects a string literal argument"))
1748 }),
1749 _ => Err(RedDBError::Query(format!(
1750 "{function}(...) expects a string literal argument"
1751 ))),
1752 }
1753}
1754
1755fn probabilistic_projection_label(projection: &Projection, base: &str, index: usize) -> String {
1756 match projection {
1757 Projection::Field(FieldRef::TableColumn { column, .. }, Some(alias))
1758 if alias.eq_ignore_ascii_case(column) =>
1759 {
1760 numbered_probabilistic_label(base, index)
1761 }
1762 Projection::Field(_, Some(alias)) => alias.clone(),
1763 Projection::Alias(column, alias) if column.eq_ignore_ascii_case(alias) => {
1764 numbered_probabilistic_label(base, index)
1765 }
1766 Projection::Alias(_, alias) => alias.clone(),
1767 Projection::Function(name, _) => name
1768 .split_once(':')
1769 .map(|(_, alias)| {
1770 if is_generated_probabilistic_function_label(alias, base) {
1771 numbered_probabilistic_label(base, index)
1772 } else {
1773 alias.to_string()
1774 }
1775 })
1776 .unwrap_or_else(|| numbered_probabilistic_label(base, index)),
1777 _ => numbered_probabilistic_label(base, index),
1778 }
1779}
1780
1781fn is_generated_probabilistic_function_label(alias: &str, base: &str) -> bool {
1782 alias
1783 .get(..base.len())
1784 .is_some_and(|head| head.eq_ignore_ascii_case(base))
1785 && alias[base.len()..].starts_with('(')
1786}
1787
1788fn numbered_probabilistic_label(base: &str, index: usize) -> String {
1789 if index == 0 {
1790 base.to_string()
1791 } else {
1792 format!("{base}_{}", index + 1)
1793 }
1794}
1795
1796fn probabilistic_projection_form(projection: &ProbabilisticReadProjection) -> &'static str {
1797 match projection {
1798 ProbabilisticReadProjection::Cardinality { .. } => "SELECT CARDINALITY",
1799 ProbabilisticReadProjection::Freq { .. } => "FREQ(...)",
1800 ProbabilisticReadProjection::Contains { .. } => "CONTAINS(...)",
1801 }
1802}
1803
1804#[cfg(test)]
1805mod recovery_tests {
1806 use crate::{RedDBOptions, RedDBRuntime};
1807
1808 #[test]
1809 fn probabilistic_checkpoint_serializes_concurrent_mutations() {
1810 let directory = tempfile::tempdir().expect("temporary database");
1811 let path = directory.path().join("data.rdb");
1812 let options = RedDBOptions::persistent(&path)
1813 .with_memory_budget(64 * 1024 * 1024)
1814 .with_bounded_disk_residency(true);
1815 let runtime =
1816 std::sync::Arc::new(RedDBRuntime::with_options(options.clone()).expect("runtime"));
1817 runtime
1818 .execute_query("CREATE SKETCH frequencies")
1819 .expect("create sketch");
1820 runtime.checkpoint().expect("baseline checkpoint");
1821 let barrier = std::sync::Arc::new(std::sync::Barrier::new(3));
1822 let mut workers = Vec::new();
1823 for _ in 0..2 {
1824 let runtime = std::sync::Arc::clone(&runtime);
1825 let barrier = std::sync::Arc::clone(&barrier);
1826 workers.push(std::thread::spawn(move || {
1827 barrier.wait();
1828 for _ in 0..10 {
1829 runtime
1830 .execute_query("SKETCH ADD frequencies 'red' 1")
1831 .expect("add");
1832 }
1833 }));
1834 }
1835 barrier.wait();
1836 for _ in 0..3 {
1837 runtime.checkpoint().expect("concurrent checkpoint");
1838 }
1839 for worker in workers {
1840 worker.join().expect("worker");
1841 }
1842 runtime.checkpoint().expect("final checkpoint");
1843 runtime.graceful_shutdown(false).expect("shutdown");
1844 drop(runtime);
1845 let reopened = RedDBRuntime::with_options(options).expect("reopen");
1846 let count = reopened
1847 .execute_query("SKETCH COUNT frequencies 'red'")
1848 .expect("sketch count");
1849 assert_eq!(
1850 count.result.records[0]
1851 .get("estimate")
1852 .expect("estimate")
1853 .display_string(),
1854 "20"
1855 );
1856 reopened
1857 .graceful_shutdown(false)
1858 .expect("reopened shutdown");
1859 }
1860
1861 #[test]
1862 fn probabilistic_snapshot_recovery_survives_interruption() {
1863 const CHILD: &str = "REDDB_TEST_PROBABILISTIC_SNAPSHOT_PATH";
1864 const TEST: &str = "runtime::impl_probabilistic::recovery_tests::probabilistic_snapshot_recovery_survives_interruption";
1865 if let Some(path) = std::env::var_os(CHILD) {
1866 let options = RedDBOptions::persistent(std::path::PathBuf::from(path))
1867 .with_memory_budget(64 * 1024 * 1024)
1868 .with_bounded_disk_residency(true);
1869 let runtime = RedDBRuntime::with_options(options).expect("runtime");
1870 for query in [
1871 "CREATE HLL visitors",
1872 "CREATE SKETCH frequencies",
1873 "CREATE FILTER seen",
1874 ] {
1875 runtime.execute_query(query).expect("create structure");
1876 }
1877 runtime.checkpoint().expect("baseline checkpoint");
1878 for query in [
1879 "HLL ADD visitors 'alice' 'bob' 'carol'",
1880 "SKETCH ADD frequencies 'red' 3",
1881 "FILTER ADD seen 'alice'",
1882 "CREATE TABLE ordinary (id INT)",
1883 ] {
1884 runtime.execute_query(query).expect("acknowledged mutation");
1885 }
1886 let hook = std::env::var("REDDB_TEST_PROBABILISTIC_SNAPSHOT_HOOK").expect("fault hook");
1887 std::env::set_var(hook, "1");
1888 runtime.checkpoint().expect("fault injection");
1889 panic!("snapshot interruption was not reached");
1890 }
1891 for (hook, code) in [
1892 ("REDDB_TEST_EXIT_AFTER_PROBABILISTIC_SNAPSHOT_WAL", 213),
1893 ("REDDB_TEST_EXIT_AFTER_PROBABILISTIC_SNAPSHOT_BLOB", 214),
1894 ] {
1895 let directory = tempfile::tempdir().expect("temporary database");
1896 let path = directory.path().join("data.rdb");
1897 let child =
1898 std::process::Command::new(std::env::current_exe().expect("test executable"))
1899 .args(["--exact", TEST, "--nocapture"])
1900 .env(CHILD, &path)
1901 .env("REDDB_TEST_PROBABILISTIC_SNAPSHOT_HOOK", hook)
1902 .output()
1903 .expect("child process");
1904 assert_eq!(
1905 child.status.code(),
1906 Some(code),
1907 "{}",
1908 String::from_utf8_lossy(&child.stderr)
1909 );
1910 let options = RedDBOptions::persistent(&path)
1911 .with_memory_budget(64 * 1024 * 1024)
1912 .with_bounded_disk_residency(true);
1913 for _ in 0..2 {
1914 let runtime = RedDBRuntime::with_options(options.clone()).expect("recover runtime");
1915 let hll = runtime
1916 .execute_query("HLL COUNT visitors")
1917 .expect("HLL count");
1918 assert_eq!(
1919 hll.result.records[0]
1920 .get("count")
1921 .expect("count")
1922 .display_string(),
1923 "3"
1924 );
1925 let sketch = runtime
1926 .execute_query("SKETCH COUNT frequencies 'red'")
1927 .expect("sketch count");
1928 assert_eq!(
1929 sketch.result.records[0]
1930 .get("estimate")
1931 .expect("estimate")
1932 .display_string(),
1933 "3"
1934 );
1935 let filter = runtime
1936 .execute_query("FILTER CHECK seen 'alice'")
1937 .expect("filter check");
1938 assert_eq!(
1939 filter.result.records[0]
1940 .get("exists")
1941 .expect("exists")
1942 .display_string(),
1943 "true"
1944 );
1945 runtime.checkpoint().expect("recovered checkpoint");
1946 runtime.graceful_shutdown(false).expect("shutdown");
1947 }
1948 }
1949 }
1950}