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