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