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.";
9
10fn probabilistic_read<'a, T>(lock: &'a RwLock<T>, _name: &str) -> RwLockReadGuard<'a, T> {
11 lock.read()
12}
13
14fn probabilistic_write<'a, T>(lock: &'a RwLock<T>, _name: &str) -> RwLockWriteGuard<'a, T> {
15 lock.write()
16}
17
18fn probabilistic_collection_contract(
19 name: &str,
20 model: crate::catalog::CollectionModel,
21) -> crate::physical::CollectionContract {
22 let now = crate::utils::now_unix_millis() as u128;
23 crate::physical::CollectionContract {
24 name: name.to_string(),
25 declared_model: model,
26 schema_mode: crate::catalog::SchemaMode::Dynamic,
27 origin: crate::physical::ContractOrigin::Explicit,
28 version: 1,
29 created_at_unix_ms: now,
30 updated_at_unix_ms: now,
31 default_ttl_ms: None,
32 vector_dimension: None,
33 vector_metric: None,
34 context_index_fields: Vec::new(),
35 declared_columns: Vec::new(),
36 table_def: None,
37 timestamps_enabled: false,
38 context_index_enabled: false,
39 metrics_raw_retention_ms: None,
40 metrics_rollup_policies: Vec::new(),
41 metrics_tenant_identity: None,
42 metrics_namespace: None,
43 append_only: false,
44 subscriptions: Vec::new(),
45 analytics_config: Vec::new(),
46 session_key: None,
47 session_gap_ms: None,
48 retention_duration_ms: None,
49 analytical_storage: None,
50
51 ai_policy: None,
52 }
53}
54
55fn hll_precision_mismatch_error(
56 operation: &str,
57 expected_name: &str,
58 expected_precision: u8,
59 actual_name: &str,
60 actual_precision: u8,
61) -> RedDBError {
62 RedDBError::Query(format!(
63 "HLL {operation} requires matching precision; '{expected_name}' uses precision {expected_precision}, but '{actual_name}' uses precision {actual_precision}"
64 ))
65}
66
67enum ProbabilisticReadProjection {
68 Cardinality { label: String },
69 Freq { element: String, label: String },
70 Contains { element: String, label: String },
71}
72
73impl RedDBRuntime {
74 pub(crate) fn load_probabilistic_state(&self) -> RedDBResult<()> {
75 {
76 let entries = self.latest_probabilistic_state_entries(PROB_HLL_STATE_PREFIX);
77 let mut hlls =
78 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
79 for (name, data_hex) in entries {
80 let bytes = hex::decode(&data_hex).map_err(|err| {
81 RedDBError::Internal(format!("invalid persisted HLL state for '{name}': {err}"))
82 })?;
83 let Some(hll) =
84 crate::storage::primitives::hyperloglog::HyperLogLog::from_bytes(bytes)
85 else {
86 return Err(RedDBError::Internal(format!(
87 "invalid persisted HLL state for '{name}'"
88 )));
89 };
90 hlls.insert(name, hll);
91 }
92 }
93
94 {
95 let entries = self.latest_probabilistic_state_entries(PROB_SKETCH_STATE_PREFIX);
96 let mut sketches = probabilistic_write(
97 &self.inner.probabilistic.sketches,
98 "probabilistic sketch store",
99 );
100 for (name, data_hex) in entries {
101 let bytes = hex::decode(&data_hex).map_err(|err| {
102 RedDBError::Internal(format!(
103 "invalid persisted SKETCH state for '{name}': {err}"
104 ))
105 })?;
106 let sketch =
107 crate::storage::primitives::count_min_sketch::CountMinSketch::from_bytes(
108 &bytes,
109 )
110 .ok_or_else(|| {
111 RedDBError::Internal(format!("invalid persisted SKETCH state for '{name}'"))
112 })?;
113 sketches.insert(name, sketch);
114 }
115 }
116
117 {
118 let entries = self.latest_probabilistic_state_entries(PROB_FILTER_STATE_PREFIX);
119 let mut filters = probabilistic_write(
120 &self.inner.probabilistic.filters,
121 "probabilistic filter store",
122 );
123 for (name, data_hex) in entries {
124 let bytes = hex::decode(&data_hex).map_err(|err| {
125 RedDBError::Internal(format!(
126 "invalid persisted FILTER state for '{name}': {err}"
127 ))
128 })?;
129 let filter =
130 crate::storage::primitives::cuckoo_filter::CuckooFilter::from_bytes(&bytes)
131 .ok_or_else(|| {
132 RedDBError::Internal(format!(
133 "invalid persisted FILTER state for '{name}'"
134 ))
135 })?;
136 filters.insert(name, filter);
137 }
138 }
139
140 Ok(())
141 }
142
143 fn latest_probabilistic_state_entries(&self, prefix: &str) -> Vec<(String, String)> {
144 let Some(manager) = self.inner.db.store().get_collection("red_config") else {
145 return Vec::new();
146 };
147 let mut latest: std::collections::HashMap<String, (u64, Option<String>)> =
148 std::collections::HashMap::new();
149 for entity in manager.query_all(|_| true) {
150 let EntityData::Row(row) = &entity.data else {
151 continue;
152 };
153 let Some(named) = &row.named else {
154 continue;
155 };
156 let Some(Value::Text(key)) = named.get("key") else {
157 continue;
158 };
159 let Some(encoded_name) = key.strip_prefix(prefix) else {
160 continue;
161 };
162 let value = match named.get("value") {
163 Some(Value::Text(value)) => Some(value.to_string()),
164 Some(Value::Null) => None,
165 _ => continue,
166 };
167 let entity_id = entity.id.raw();
168 match latest.get(encoded_name) {
169 Some((existing_id, _)) if *existing_id > entity_id => {}
170 _ => {
171 latest.insert(encoded_name.to_string(), (entity_id, value));
172 }
173 }
174 }
175
176 latest
177 .into_iter()
178 .filter_map(|(encoded_name, (_, value))| {
179 let value = value?;
180 let bytes = hex::decode(encoded_name).ok()?;
181 let name = String::from_utf8(bytes).ok()?;
182 Some((name, value))
183 })
184 .collect()
185 }
186
187 fn persist_probabilistic_blob(
188 &self,
189 prefix: &str,
190 name: &str,
191 bytes: &[u8],
192 ) -> RedDBResult<()> {
193 let key = format!("{prefix}{}", hex::encode(name.as_bytes()));
194 self.inner
195 .db
196 .store()
197 .set_config_tree(&key, &crate::serde_json::Value::String(hex::encode(bytes)));
198 Ok(())
199 }
200
201 fn delete_probabilistic_blob(&self, prefix: &str, name: &str) -> RedDBResult<()> {
202 let key = format!("{prefix}{}", hex::encode(name.as_bytes()));
203 self.inner
204 .db
205 .store()
206 .set_config_tree(&key, &crate::serde_json::Value::Null);
207 Ok(())
208 }
209
210 fn create_probabilistic_catalog_entry(
211 &self,
212 name: &str,
213 model: crate::catalog::CollectionModel,
214 ) -> RedDBResult<()> {
215 let store = self.inner.db.store();
216 store
217 .create_collection(name)
218 .map_err(|err| RedDBError::Internal(err.to_string()))?;
219 self.inner
220 .db
221 .save_collection_contract(probabilistic_collection_contract(name, model))
222 .map_err(|err| RedDBError::Internal(err.to_string()))?;
223 if let Some(tenant_id) = crate::runtime::impl_core::current_tenant() {
224 store.set_config_tree(
225 &format!("red.collection_tenants.{name}"),
226 &crate::serde_json::Value::String(tenant_id),
227 );
228 }
229 self.inner
230 .db
231 .persist_metadata()
232 .map_err(|err| RedDBError::Internal(err.to_string()))?;
233 self.invalidate_result_cache();
234 Ok(())
235 }
236
237 fn drop_probabilistic_catalog_entry(&self, name: &str) -> RedDBResult<()> {
238 let store = self.inner.db.store();
239 if store.get_collection(name).is_some() {
240 store
241 .drop_collection(name)
242 .map_err(|err| RedDBError::Internal(err.to_string()))?;
243 }
244 self.inner
245 .db
246 .remove_collection_contract(name)
247 .map_err(|err| RedDBError::Internal(err.to_string()))?;
248 self.inner
249 .db
250 .persist_metadata()
251 .map_err(|err| RedDBError::Internal(err.to_string()))?;
252 self.invalidate_result_cache();
253 Ok(())
254 }
255
256 pub(crate) fn execute_probabilistic_select(
257 &self,
258 query: &TableQuery,
259 ) -> RedDBResult<Option<UnifiedResult>> {
260 let projections = crate::storage::query::sql_lowering::effective_table_projections(query);
261 let mut read_projections = Vec::new();
262 for projection in &projections {
263 if let Some(read_projection) =
264 parse_probabilistic_read_projection(projection, read_projections.len())?
265 {
266 read_projections.push(read_projection);
267 }
268 }
269
270 let Some(actual_model) = self
271 .inner
272 .db
273 .collection_contract(&query.table)
274 .map(|contract| contract.declared_model)
275 else {
276 return if read_projections.is_empty() {
277 Ok(None)
278 } else {
279 Err(RedDBError::NotFound(format!(
280 "probabilistic collection '{}' not found",
281 query.table
282 )))
283 };
284 };
285
286 let is_probabilistic_model = matches!(
287 actual_model,
288 crate::catalog::CollectionModel::Hll
289 | crate::catalog::CollectionModel::Sketch
290 | crate::catalog::CollectionModel::Filter
291 );
292 if read_projections.is_empty() {
293 return if is_probabilistic_model {
294 Err(RedDBError::Query(format!(
295 "probabilistic collection '{}' supports SELECT CARDINALITY, FREQ(...), or CONTAINS(...) read forms",
296 query.table
297 )))
298 } else {
299 Ok(None)
300 };
301 }
302
303 validate_probabilistic_read_model(&query.table, actual_model, &read_projections)?;
304 let (columns, record) =
305 self.materialize_probabilistic_select_row(&query.table, &read_projections)?;
306 let mut result = UnifiedResult::with_columns(columns);
307 if probabilistic_select_row_visible(self, query, &record) {
308 result.push(record);
309 }
310 Ok(Some(result))
311 }
312
313 pub fn execute_probabilistic_command(
314 &self,
315 raw_query: &str,
316 cmd: &ProbabilisticCommand,
317 ) -> RedDBResult<RuntimeQueryResult> {
318 let is_mutation = matches!(
322 cmd,
323 ProbabilisticCommand::CreateHll { .. }
324 | ProbabilisticCommand::HllAdd { .. }
325 | ProbabilisticCommand::HllMerge { .. }
326 | ProbabilisticCommand::DropHll { .. }
327 | ProbabilisticCommand::CreateSketch { .. }
328 | ProbabilisticCommand::SketchAdd { .. }
329 | ProbabilisticCommand::SketchMerge { .. }
330 | ProbabilisticCommand::DropSketch { .. }
331 | ProbabilisticCommand::CreateFilter { .. }
332 | ProbabilisticCommand::FilterAdd { .. }
333 | ProbabilisticCommand::FilterDelete { .. }
334 | ProbabilisticCommand::DropFilter { .. }
335 );
336 if is_mutation {
337 self.check_write(crate::runtime::write_gate::WriteKind::Dml)?;
338 }
339 match cmd {
340 ProbabilisticCommand::CreateHll {
342 name,
343 precision,
344 if_not_exists,
345 } => {
346 let mut hlls =
347 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
348 if hlls.contains_key(name) {
349 if *if_not_exists {
350 return Ok(RuntimeQueryResult::ok_message(
351 raw_query.to_string(),
352 &format!("HLL '{}' already exists", name),
353 "create",
354 ));
355 }
356 return Err(RedDBError::Query(format!("HLL '{}' already exists", name)));
357 }
358 let hll = crate::storage::primitives::hyperloglog::HyperLogLog::with_precision(
359 *precision,
360 )
361 .ok_or_else(|| {
362 RedDBError::Query(format!(
363 "HLL precision must be between 4 and 18, got {precision}"
364 ))
365 })?;
366 self.create_probabilistic_catalog_entry(
367 name,
368 crate::catalog::CollectionModel::Hll,
369 )?;
370 self.persist_probabilistic_blob(PROB_HLL_STATE_PREFIX, name, hll.as_bytes())?;
371 hlls.insert(name.clone(), hll);
372 Ok(RuntimeQueryResult::ok_message(
373 raw_query.to_string(),
374 &format!("HLL '{}' created", name),
375 "create",
376 ))
377 }
378 ProbabilisticCommand::HllAdd { name, elements } => {
379 let mut hlls =
380 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
381 let hll = hlls
382 .get_mut(name)
383 .ok_or_else(|| RedDBError::NotFound(format!("HLL '{}' not found", name)))?;
384 for elem in elements {
385 hll.add(elem.as_bytes());
386 }
387 self.persist_probabilistic_blob(PROB_HLL_STATE_PREFIX, name, hll.as_bytes())?;
388 Ok(RuntimeQueryResult::ok_message(
389 raw_query.to_string(),
390 &format!("{} element(s) added to HLL '{}'", elements.len(), name),
391 "insert",
392 ))
393 }
394 ProbabilisticCommand::HllCount { names } => {
395 let hlls =
396 probabilistic_read(&self.inner.probabilistic.hlls, "probabilistic HLL store");
397 if names.is_empty() {
398 return Err(RedDBError::Query(
399 "HLL COUNT requires at least one HLL name".to_string(),
400 ));
401 }
402 if names.len() == 1 {
403 let hll = hlls.get(&names[0]).ok_or_else(|| {
404 RedDBError::NotFound(format!("HLL '{}' not found", names[0]))
405 })?;
406 let count = hll.count();
407 let mut result = UnifiedResult::with_columns(vec!["count".into()]);
408 let mut record = UnifiedRecord::new();
409 record.set("count", Value::UnsignedInteger(count));
410 result.push(record);
411 Ok(RuntimeQueryResult {
412 query: raw_query.to_string(),
413 mode: QueryMode::Sql,
414 statement: "hll_count",
415 engine: "runtime-probabilistic",
416 result,
417 affected_rows: 0,
418 statement_type: "select",
419 bookmark: None,
420 })
421 } else {
422 let first_name = &names[0];
424 let first = hlls.get(first_name).ok_or_else(|| {
425 RedDBError::NotFound(format!("HLL '{}' not found", first_name))
426 })?;
427 let expected_precision = first.precision();
428 let mut merged =
429 crate::storage::primitives::hyperloglog::HyperLogLog::with_precision(
430 expected_precision,
431 )
432 .expect("loaded HLL precision is valid");
433 for name in names {
434 let hll = hlls.get(name).ok_or_else(|| {
435 RedDBError::NotFound(format!("HLL '{}' not found", name))
436 })?;
437 if hll.precision() != expected_precision {
438 return Err(hll_precision_mismatch_error(
439 "COUNT",
440 first_name,
441 expected_precision,
442 name,
443 hll.precision(),
444 ));
445 }
446 merged.merge(hll);
447 }
448 let count = merged.count();
449 let mut result = UnifiedResult::with_columns(vec!["count".into()]);
450 let mut record = UnifiedRecord::new();
451 record.set("count", Value::UnsignedInteger(count));
452 result.push(record);
453 Ok(RuntimeQueryResult {
454 query: raw_query.to_string(),
455 mode: QueryMode::Sql,
456 statement: "hll_count",
457 engine: "runtime-probabilistic",
458 result,
459 affected_rows: 0,
460 statement_type: "select",
461 bookmark: None,
462 })
463 }
464 }
465 ProbabilisticCommand::HllMerge { dest, sources } => {
466 let mut hlls =
467 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
468 let first_src = sources.first().ok_or_else(|| {
469 RedDBError::Query("HLL MERGE requires at least one source HLL".to_string())
470 })?;
471 let first = hlls.get(first_src).ok_or_else(|| {
472 RedDBError::NotFound(format!("HLL '{}' not found", first_src))
473 })?;
474 let expected_precision = first.precision();
475 let mut merged =
476 crate::storage::primitives::hyperloglog::HyperLogLog::with_precision(
477 expected_precision,
478 )
479 .expect("loaded HLL precision is valid");
480 for src in sources {
481 let hll = hlls
482 .get(src)
483 .ok_or_else(|| RedDBError::NotFound(format!("HLL '{}' not found", src)))?;
484 if hll.precision() != expected_precision {
485 return Err(hll_precision_mismatch_error(
486 "MERGE",
487 first_src,
488 expected_precision,
489 src,
490 hll.precision(),
491 ));
492 }
493 merged.merge(hll);
494 }
495 self.persist_probabilistic_blob(PROB_HLL_STATE_PREFIX, dest, merged.as_bytes())?;
496 hlls.insert(dest.clone(), merged);
497 Ok(RuntimeQueryResult::ok_message(
498 raw_query.to_string(),
499 &format!(
500 "HLL '{}' created from merge of {}",
501 dest,
502 sources.join(", ")
503 ),
504 "create",
505 ))
506 }
507 ProbabilisticCommand::HllInfo { name } => {
508 let hlls =
509 probabilistic_read(&self.inner.probabilistic.hlls, "probabilistic HLL store");
510 let hll = hlls
511 .get(name)
512 .ok_or_else(|| RedDBError::NotFound(format!("HLL '{}' not found", name)))?;
513 let mut result = UnifiedResult::with_columns(vec![
514 "name".into(),
515 "precision".into(),
516 "count".into(),
517 "memory_bytes".into(),
518 ]);
519 let mut record = UnifiedRecord::new();
520 record.set("name", Value::text(name.clone()));
521 record.set("precision", Value::UnsignedInteger(hll.precision() as u64));
522 record.set("count", Value::UnsignedInteger(hll.count()));
523 record.set(
524 "memory_bytes",
525 Value::UnsignedInteger(hll.memory_bytes() as u64),
526 );
527 result.push(record);
528 Ok(RuntimeQueryResult {
529 query: raw_query.to_string(),
530 mode: QueryMode::Sql,
531 statement: "hll_info",
532 engine: "runtime-probabilistic",
533 result,
534 affected_rows: 0,
535 statement_type: "select",
536 bookmark: None,
537 })
538 }
539 ProbabilisticCommand::DropHll { name, if_exists } => {
540 let mut hlls =
541 probabilistic_write(&self.inner.probabilistic.hlls, "probabilistic HLL store");
542 if hlls.remove(name).is_none() {
543 if *if_exists {
544 return Ok(RuntimeQueryResult::ok_message(
545 raw_query.to_string(),
546 &format!("HLL '{}' does not exist", name),
547 "drop",
548 ));
549 }
550 return Err(RedDBError::NotFound(format!("HLL '{}' not found", name)));
551 }
552 self.drop_probabilistic_catalog_entry(name)?;
553 self.delete_probabilistic_blob(PROB_HLL_STATE_PREFIX, name)?;
554 Ok(RuntimeQueryResult::ok_message(
555 raw_query.to_string(),
556 &format!("HLL '{}' dropped", name),
557 "drop",
558 ))
559 }
560
561 ProbabilisticCommand::CreateSketch {
563 name,
564 width,
565 depth,
566 if_not_exists,
567 } => {
568 let mut sketches = probabilistic_write(
569 &self.inner.probabilistic.sketches,
570 "probabilistic sketch store",
571 );
572 if sketches.contains_key(name) {
573 if *if_not_exists {
574 return Ok(RuntimeQueryResult::ok_message(
575 raw_query.to_string(),
576 &format!("SKETCH '{}' already exists", name),
577 "create",
578 ));
579 }
580 return Err(RedDBError::Query(format!(
581 "SKETCH '{}' already exists",
582 name
583 )));
584 }
585 self.create_probabilistic_catalog_entry(
586 name,
587 crate::catalog::CollectionModel::Sketch,
588 )?;
589 let sketch = crate::storage::primitives::count_min_sketch::CountMinSketch::new(
590 *width, *depth,
591 );
592 self.persist_probabilistic_blob(
593 PROB_SKETCH_STATE_PREFIX,
594 name,
595 &sketch.as_bytes(),
596 )?;
597 sketches.insert(name.clone(), sketch);
598 Ok(RuntimeQueryResult::ok_message(
599 raw_query.to_string(),
600 &format!(
601 "SKETCH '{}' created (width={}, depth={})",
602 name, width, depth
603 ),
604 "create",
605 ))
606 }
607 ProbabilisticCommand::SketchAdd {
608 name,
609 element,
610 count,
611 } => {
612 let mut sketches = probabilistic_write(
613 &self.inner.probabilistic.sketches,
614 "probabilistic sketch store",
615 );
616 let sketch = sketches
617 .get_mut(name)
618 .ok_or_else(|| RedDBError::NotFound(format!("SKETCH '{}' not found", name)))?;
619 sketch.add(element.as_bytes(), *count);
620 self.persist_probabilistic_blob(
621 PROB_SKETCH_STATE_PREFIX,
622 name,
623 &sketch.as_bytes(),
624 )?;
625 Ok(RuntimeQueryResult::ok_message(
626 raw_query.to_string(),
627 &format!("added {} to SKETCH '{}'", count, name),
628 "insert",
629 ))
630 }
631 ProbabilisticCommand::SketchCount { name, element } => {
632 let sketches = probabilistic_read(
633 &self.inner.probabilistic.sketches,
634 "probabilistic sketch store",
635 );
636 let sketch = sketches
637 .get(name)
638 .ok_or_else(|| RedDBError::NotFound(format!("SKETCH '{}' not found", name)))?;
639 let estimate = sketch.estimate(element.as_bytes());
640 let mut result = UnifiedResult::with_columns(vec!["estimate".into()]);
641 let mut record = UnifiedRecord::new();
642 record.set("estimate", Value::UnsignedInteger(estimate));
643 result.push(record);
644 Ok(RuntimeQueryResult {
645 query: raw_query.to_string(),
646 mode: QueryMode::Sql,
647 statement: "sketch_count",
648 engine: "runtime-probabilistic",
649 result,
650 affected_rows: 0,
651 statement_type: "select",
652 bookmark: None,
653 })
654 }
655 ProbabilisticCommand::SketchMerge { dest, sources } => {
656 let mut sketches = probabilistic_write(
657 &self.inner.probabilistic.sketches,
658 "probabilistic sketch store",
659 );
660 let first_src = sketches.get(&sources[0]).ok_or_else(|| {
661 RedDBError::NotFound(format!("SKETCH '{}' not found", sources[0]))
662 })?;
663 let mut merged = crate::storage::primitives::count_min_sketch::CountMinSketch::new(
664 first_src.width(),
665 first_src.depth(),
666 );
667 for src in sources {
668 let sketch = sketches.get(src).ok_or_else(|| {
669 RedDBError::NotFound(format!("SKETCH '{}' not found", src))
670 })?;
671 if !merged.merge(sketch) {
672 return Err(RedDBError::Query(format!(
673 "SKETCH '{}' has incompatible dimensions",
674 src
675 )));
676 }
677 }
678 self.persist_probabilistic_blob(
679 PROB_SKETCH_STATE_PREFIX,
680 dest,
681 &merged.as_bytes(),
682 )?;
683 sketches.insert(dest.clone(), merged);
684 Ok(RuntimeQueryResult::ok_message(
685 raw_query.to_string(),
686 &format!(
687 "SKETCH '{}' created from merge of {}",
688 dest,
689 sources.join(", ")
690 ),
691 "create",
692 ))
693 }
694 ProbabilisticCommand::SketchInfo { name } => {
695 let sketches = probabilistic_read(
696 &self.inner.probabilistic.sketches,
697 "probabilistic sketch store",
698 );
699 let sketch = sketches
700 .get(name)
701 .ok_or_else(|| RedDBError::NotFound(format!("SKETCH '{}' not found", name)))?;
702 let mut result = UnifiedResult::with_columns(vec![
703 "name".into(),
704 "width".into(),
705 "depth".into(),
706 "total".into(),
707 "memory_bytes".into(),
708 ]);
709 let mut record = UnifiedRecord::new();
710 record.set("name", Value::text(name.clone()));
711 record.set("width", Value::UnsignedInteger(sketch.width() as u64));
712 record.set("depth", Value::UnsignedInteger(sketch.depth() as u64));
713 record.set("total", Value::UnsignedInteger(sketch.total()));
714 record.set(
715 "memory_bytes",
716 Value::UnsignedInteger(sketch.memory_bytes() as u64),
717 );
718 result.push(record);
719 Ok(RuntimeQueryResult {
720 query: raw_query.to_string(),
721 mode: QueryMode::Sql,
722 statement: "sketch_info",
723 engine: "runtime-probabilistic",
724 result,
725 affected_rows: 0,
726 statement_type: "select",
727 bookmark: None,
728 })
729 }
730 ProbabilisticCommand::DropSketch { name, if_exists } => {
731 let mut sketches = probabilistic_write(
732 &self.inner.probabilistic.sketches,
733 "probabilistic sketch store",
734 );
735 if sketches.remove(name).is_none() {
736 if *if_exists {
737 return Ok(RuntimeQueryResult::ok_message(
738 raw_query.to_string(),
739 &format!("SKETCH '{}' does not exist", name),
740 "drop",
741 ));
742 }
743 return Err(RedDBError::NotFound(format!("SKETCH '{}' not found", name)));
744 }
745 self.drop_probabilistic_catalog_entry(name)?;
746 self.delete_probabilistic_blob(PROB_SKETCH_STATE_PREFIX, name)?;
747 Ok(RuntimeQueryResult::ok_message(
748 raw_query.to_string(),
749 &format!("SKETCH '{}' dropped", name),
750 "drop",
751 ))
752 }
753
754 ProbabilisticCommand::CreateFilter {
756 name,
757 capacity,
758 if_not_exists,
759 } => {
760 let mut filters = probabilistic_write(
761 &self.inner.probabilistic.filters,
762 "probabilistic filter store",
763 );
764 if filters.contains_key(name) {
765 if *if_not_exists {
766 return Ok(RuntimeQueryResult::ok_message(
767 raw_query.to_string(),
768 &format!("FILTER '{}' already exists", name),
769 "create",
770 ));
771 }
772 return Err(RedDBError::Query(format!(
773 "FILTER '{}' already exists",
774 name
775 )));
776 }
777 self.create_probabilistic_catalog_entry(
778 name,
779 crate::catalog::CollectionModel::Filter,
780 )?;
781 let filter =
782 crate::storage::primitives::cuckoo_filter::CuckooFilter::new(*capacity);
783 self.persist_probabilistic_blob(
784 PROB_FILTER_STATE_PREFIX,
785 name,
786 &filter.as_bytes(),
787 )?;
788 filters.insert(name.clone(), filter);
789 Ok(RuntimeQueryResult::ok_message(
790 raw_query.to_string(),
791 &format!("FILTER '{}' created (capacity={})", name, capacity),
792 "create",
793 ))
794 }
795 ProbabilisticCommand::FilterAdd { name, element } => {
796 let mut filters = probabilistic_write(
797 &self.inner.probabilistic.filters,
798 "probabilistic filter store",
799 );
800 let filter = filters
801 .get_mut(name)
802 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
803 if !filter.insert(element.as_bytes()) {
804 return Err(RedDBError::Query(format!("FILTER '{}' is full", name)));
805 }
806 self.persist_probabilistic_blob(
807 PROB_FILTER_STATE_PREFIX,
808 name,
809 &filter.as_bytes(),
810 )?;
811 Ok(RuntimeQueryResult::ok_message(
812 raw_query.to_string(),
813 &format!("element added to FILTER '{}'", name),
814 "insert",
815 ))
816 }
817 ProbabilisticCommand::FilterCheck { name, element } => {
818 let filters = probabilistic_read(
819 &self.inner.probabilistic.filters,
820 "probabilistic filter store",
821 );
822 let filter = filters
823 .get(name)
824 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
825 let exists = filter.contains(element.as_bytes());
826 let mut result = UnifiedResult::with_columns(vec!["exists".into()]);
827 let mut record = UnifiedRecord::new();
828 record.set("exists", Value::Boolean(exists));
829 result.push(record);
830 Ok(RuntimeQueryResult {
831 query: raw_query.to_string(),
832 mode: QueryMode::Sql,
833 statement: "filter_check",
834 engine: "runtime-probabilistic",
835 result,
836 affected_rows: 0,
837 statement_type: "select",
838 bookmark: None,
839 })
840 }
841 ProbabilisticCommand::FilterDelete { name, element } => {
842 let mut filters = probabilistic_write(
843 &self.inner.probabilistic.filters,
844 "probabilistic filter store",
845 );
846 let filter = filters
847 .get_mut(name)
848 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
849 let removed = filter.delete(element.as_bytes());
850 self.persist_probabilistic_blob(
851 PROB_FILTER_STATE_PREFIX,
852 name,
853 &filter.as_bytes(),
854 )?;
855 Ok(RuntimeQueryResult::ok_message(
856 raw_query.to_string(),
857 &format!(
858 "element {} from FILTER '{}'",
859 if removed { "deleted" } else { "not found in" },
860 name
861 ),
862 "delete",
863 ))
864 }
865 ProbabilisticCommand::FilterCount { name } => {
866 let filters = probabilistic_read(
867 &self.inner.probabilistic.filters,
868 "probabilistic filter store",
869 );
870 let filter = filters
871 .get(name)
872 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
873 let mut result = UnifiedResult::with_columns(vec!["count".into()]);
874 let mut record = UnifiedRecord::new();
875 record.set("count", Value::UnsignedInteger(filter.count() as u64));
876 result.push(record);
877 Ok(RuntimeQueryResult {
878 query: raw_query.to_string(),
879 mode: QueryMode::Sql,
880 statement: "filter_count",
881 engine: "runtime-probabilistic",
882 result,
883 affected_rows: 0,
884 statement_type: "select",
885 bookmark: None,
886 })
887 }
888 ProbabilisticCommand::FilterInfo { name } => {
889 let filters = probabilistic_read(
890 &self.inner.probabilistic.filters,
891 "probabilistic filter store",
892 );
893 let filter = filters
894 .get(name)
895 .ok_or_else(|| RedDBError::NotFound(format!("FILTER '{}' not found", name)))?;
896 let mut result = UnifiedResult::with_columns(vec![
897 "name".into(),
898 "count".into(),
899 "load_factor".into(),
900 "memory_bytes".into(),
901 ]);
902 let mut record = UnifiedRecord::new();
903 record.set("name", Value::text(name.clone()));
904 record.set("count", Value::UnsignedInteger(filter.count() as u64));
905 record.set("load_factor", Value::Float(filter.load_factor()));
906 record.set(
907 "memory_bytes",
908 Value::UnsignedInteger(filter.memory_bytes() as u64),
909 );
910 result.push(record);
911 Ok(RuntimeQueryResult {
912 query: raw_query.to_string(),
913 mode: QueryMode::Sql,
914 statement: "filter_info",
915 engine: "runtime-probabilistic",
916 result,
917 affected_rows: 0,
918 statement_type: "select",
919 bookmark: None,
920 })
921 }
922 ProbabilisticCommand::DropFilter { name, if_exists } => {
923 let mut filters = probabilistic_write(
924 &self.inner.probabilistic.filters,
925 "probabilistic filter store",
926 );
927 if filters.remove(name).is_none() {
928 if *if_exists {
929 return Ok(RuntimeQueryResult::ok_message(
930 raw_query.to_string(),
931 &format!("FILTER '{}' does not exist", name),
932 "drop",
933 ));
934 }
935 return Err(RedDBError::NotFound(format!("FILTER '{}' not found", name)));
936 }
937 self.drop_probabilistic_catalog_entry(name)?;
938 self.delete_probabilistic_blob(PROB_FILTER_STATE_PREFIX, name)?;
939 Ok(RuntimeQueryResult::ok_message(
940 raw_query.to_string(),
941 &format!("FILTER '{}' dropped", name),
942 "drop",
943 ))
944 }
945 }
946 }
947}
948
949fn parse_probabilistic_read_projection(
950 projection: &Projection,
951 index: usize,
952) -> RedDBResult<Option<ProbabilisticReadProjection>> {
953 if let Some(column) = projection_unqualified_column(projection) {
954 if column.eq_ignore_ascii_case("CARDINALITY") {
955 return Ok(Some(ProbabilisticReadProjection::Cardinality {
956 label: probabilistic_projection_label(projection, "cardinality", index),
957 }));
958 }
959 }
960
961 let Some((function, args)) = projection_function(projection) else {
962 return Ok(None);
963 };
964 if function.eq_ignore_ascii_case("FREQ") {
965 let element = projection_single_text_arg(function, args)?;
966 return Ok(Some(ProbabilisticReadProjection::Freq {
967 element,
968 label: probabilistic_projection_label(projection, "freq", index),
969 }));
970 }
971 if function.eq_ignore_ascii_case("CONTAINS") {
972 let element = projection_single_text_arg(function, args)?;
973 return Ok(Some(ProbabilisticReadProjection::Contains {
974 element,
975 label: probabilistic_projection_label(projection, "contains", index),
976 }));
977 }
978
979 Ok(None)
980}
981
982fn validate_probabilistic_read_model(
983 collection: &str,
984 actual_model: crate::catalog::CollectionModel,
985 projections: &[ProbabilisticReadProjection],
986) -> RedDBResult<()> {
987 for projection in projections {
988 let expected_model = match projection {
989 ProbabilisticReadProjection::Cardinality { .. } => crate::catalog::CollectionModel::Hll,
990 ProbabilisticReadProjection::Freq { .. } => crate::catalog::CollectionModel::Sketch,
991 ProbabilisticReadProjection::Contains { .. } => crate::catalog::CollectionModel::Filter,
992 };
993 if actual_model != expected_model {
994 return Err(RedDBError::Query(format!(
995 "{} is only supported for {} collections; '{}' is {}",
996 probabilistic_projection_form(projection),
997 crate::runtime::ddl::polymorphic_resolver::model_name(expected_model),
998 collection,
999 crate::runtime::ddl::polymorphic_resolver::model_name(actual_model)
1000 )));
1001 }
1002 }
1003 Ok(())
1004}
1005
1006impl RedDBRuntime {
1007 fn materialize_probabilistic_select_row(
1008 &self,
1009 collection: &str,
1010 projections: &[ProbabilisticReadProjection],
1011 ) -> RedDBResult<(Vec<String>, UnifiedRecord)> {
1012 let mut columns = Vec::with_capacity(projections.len());
1013 let mut record = UnifiedRecord::new();
1014 for projection in projections {
1015 match projection {
1016 ProbabilisticReadProjection::Cardinality { label } => {
1017 let hlls = probabilistic_read(
1018 &self.inner.probabilistic.hlls,
1019 "probabilistic HLL store",
1020 );
1021 let hll = hlls.get(collection).ok_or_else(|| {
1022 RedDBError::NotFound(format!("HLL '{}' not found", collection))
1023 })?;
1024 columns.push(label.clone());
1025 record.set(label, Value::UnsignedInteger(hll.count()));
1026 }
1027 ProbabilisticReadProjection::Freq { element, label } => {
1028 let sketches = probabilistic_read(
1029 &self.inner.probabilistic.sketches,
1030 "probabilistic sketch store",
1031 );
1032 let sketch = sketches.get(collection).ok_or_else(|| {
1033 RedDBError::NotFound(format!("SKETCH '{}' not found", collection))
1034 })?;
1035 columns.push(label.clone());
1036 record.set(
1037 label,
1038 Value::UnsignedInteger(sketch.estimate(element.as_bytes())),
1039 );
1040 }
1041 ProbabilisticReadProjection::Contains { element, label } => {
1042 let filters = probabilistic_read(
1043 &self.inner.probabilistic.filters,
1044 "probabilistic filter store",
1045 );
1046 let filter = filters.get(collection).ok_or_else(|| {
1047 RedDBError::NotFound(format!("FILTER '{}' not found", collection))
1048 })?;
1049 columns.push(label.clone());
1050 record.set(label, Value::Boolean(filter.contains(element.as_bytes())));
1051 }
1052 }
1053 }
1054 Ok((columns, record))
1055 }
1056}
1057
1058fn probabilistic_select_row_visible(
1059 runtime: &RedDBRuntime,
1060 query: &TableQuery,
1061 record: &UnifiedRecord,
1062) -> bool {
1063 if query.limit == Some(0) || query.offset.is_some_and(|offset| offset > 0) {
1064 return false;
1065 }
1066 let table_name = query.table.as_str();
1067 let table_alias = query.alias.as_deref().unwrap_or(table_name);
1068 crate::storage::query::sql_lowering::effective_table_filter(query).is_none_or(|filter| {
1069 super::join_filter::evaluate_runtime_filter_with_db(
1070 Some(&runtime.inner.db),
1071 record,
1072 &filter,
1073 Some(table_name),
1074 Some(table_alias),
1075 )
1076 })
1077}
1078
1079fn projection_unqualified_column(projection: &Projection) -> Option<&str> {
1080 match projection {
1081 Projection::Field(FieldRef::TableColumn { table, column }, _) if table.is_empty() => {
1082 Some(column.as_str())
1083 }
1084 Projection::Column(column) => Some(column.as_str()),
1085 Projection::Alias(column, _) => Some(column.as_str()),
1086 _ => None,
1087 }
1088}
1089
1090fn projection_function(projection: &Projection) -> Option<(&str, &[Projection])> {
1091 match projection {
1092 Projection::Function(name, args) => {
1093 let function = name.split_once(':').map(|(name, _)| name).unwrap_or(name);
1094 Some((function, args.as_slice()))
1095 }
1096 _ => None,
1097 }
1098}
1099
1100fn projection_single_text_arg(function: &str, args: &[Projection]) -> RedDBResult<String> {
1101 if args.len() != 1 {
1102 return Err(RedDBError::Query(format!(
1103 "{function}(...) expects exactly one string literal"
1104 )));
1105 }
1106 match &args[0] {
1107 Projection::Column(column) => column
1108 .strip_prefix("LIT:")
1109 .map(ToString::to_string)
1110 .ok_or_else(|| {
1111 RedDBError::Query(format!("{function}(...) expects a string literal argument"))
1112 }),
1113 _ => Err(RedDBError::Query(format!(
1114 "{function}(...) expects a string literal argument"
1115 ))),
1116 }
1117}
1118
1119fn probabilistic_projection_label(projection: &Projection, base: &str, index: usize) -> String {
1120 match projection {
1121 Projection::Field(FieldRef::TableColumn { column, .. }, Some(alias))
1122 if alias.eq_ignore_ascii_case(column) =>
1123 {
1124 numbered_probabilistic_label(base, index)
1125 }
1126 Projection::Field(_, Some(alias)) => alias.clone(),
1127 Projection::Alias(column, alias) if column.eq_ignore_ascii_case(alias) => {
1128 numbered_probabilistic_label(base, index)
1129 }
1130 Projection::Alias(_, alias) => alias.clone(),
1131 Projection::Function(name, _) => name
1132 .split_once(':')
1133 .map(|(_, alias)| {
1134 if is_generated_probabilistic_function_label(alias, base) {
1135 numbered_probabilistic_label(base, index)
1136 } else {
1137 alias.to_string()
1138 }
1139 })
1140 .unwrap_or_else(|| numbered_probabilistic_label(base, index)),
1141 _ => numbered_probabilistic_label(base, index),
1142 }
1143}
1144
1145fn is_generated_probabilistic_function_label(alias: &str, base: &str) -> bool {
1146 alias
1147 .get(..base.len())
1148 .is_some_and(|head| head.eq_ignore_ascii_case(base))
1149 && alias[base.len()..].starts_with('(')
1150}
1151
1152fn numbered_probabilistic_label(base: &str, index: usize) -> String {
1153 if index == 0 {
1154 base.to_string()
1155 } else {
1156 format!("{base}_{}", index + 1)
1157 }
1158}
1159
1160fn probabilistic_projection_form(projection: &ProbabilisticReadProjection) -> &'static str {
1161 match projection {
1162 ProbabilisticReadProjection::Cardinality { .. } => "SELECT CARDINALITY",
1163 ProbabilisticReadProjection::Freq { .. } => "FREQ(...)",
1164 ProbabilisticReadProjection::Contains { .. } => "CONTAINS(...)",
1165 }
1166}