From 0880516e8fcdf49685fac5718886a66609a061f4 Mon Sep 17 00:00:00 2001 From: phil Date: Mon, 9 Jun 2025 16:38:00 -0400 Subject: [PATCH] include all counts in reponses --- ufos/src/lib.rs | 32 ++++++++++++++++++-- ufos/src/storage_fjall.rs | 62 +++++++++++++++++---------------------- 2 files changed, 56 insertions(+), 38 deletions(-) diff --git a/ufos/src/lib.rs b/ufos/src/lib.rs index 0f8c4ce..fbc5f3a 100644 --- a/ufos/src/lib.rs +++ b/ufos/src/lib.rs @@ -10,7 +10,7 @@ pub mod store_types; use crate::db_types::{EncodingError, EncodingResult}; use crate::error::BatchInsertError; -use crate::store_types::SketchSecretPrefix; +use crate::store_types::{CountsValue, SketchSecretPrefix}; use cardinality_estimator_safe::{Element, Sketch}; use error::FirehoseEventError; use jetstream::events::{CommitEvent, CommitOp, Cursor}; @@ -281,17 +281,43 @@ pub enum ConsumerInfo { pub struct NsidCount { nsid: String, creates: u64, - // TODO: add updates and deletes + updates: u64, + deletes: u64, dids_estimate: u64, } +impl NsidCount { + pub fn new(nsid: &Nsid, counts: &CountsValue) -> Self { + let crud = counts.counts(); + Self { + nsid: nsid.to_string(), + creates: crud.creates, + updates: crud.updates, + deletes: crud.deletes, + dids_estimate: counts.dids().estimate() as u64, + } + } +} #[derive(Debug, PartialEq, Serialize, JsonSchema)] pub struct PrefixCount { prefix: String, creates: u64, - // TODO: add updates and deletes + updates: u64, + deletes: u64, dids_estimate: u64, } +impl PrefixCount { + pub fn new(prefix: &str, counts: &CountsValue) -> Self { + let crud = counts.counts(); + Self { + prefix: prefix.to_string(), + creates: crud.creates, + updates: crud.updates, + deletes: crud.deletes, + dids_estimate: counts.dids().estimate() as u64, + } + } +} #[derive(Debug, PartialEq, Serialize, JsonSchema)] #[serde(tag = "type", rename_all = "camelCase")] diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index 57e6e39..743dfa7 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -507,11 +507,7 @@ impl FjallReader { merged.merge(&counts); } } - out.push(NsidCount { - nsid: nsid.to_string(), - creates: merged.counts().creates, - dids_estimate: merged.dids().estimate() as u64, - }); + out.push(NsidCount::new(&nsid, &merged)); } let next_cursor = current_nsid.map(|s| s.to_db_bytes()).transpose()?; @@ -617,11 +613,7 @@ impl FjallReader { .into_iter() .rev() .take(limit) - .map(|(nsid, cv)| NsidCount { - nsid: nsid.to_string(), - creates: cv.counts().creates, - dids_estimate: cv.dids().estimate() as u64, - }) + .map(|(nsid, cv)| NsidCount::new(&nsid, &cv)) .collect(); Ok(counts) } @@ -727,13 +719,13 @@ impl FjallReader { let mut prefix_count = CountsValue::default(); #[derive(Debug, Clone, PartialEq)] enum Child { - FullNsid(String), + FullNsid(Nsid), ChildPrefix(String), } impl Child { fn from_prefix(nsid: &Nsid, prefix: &NsidPrefix) -> Option { if prefix.is_group_of(nsid) { - return Some(Child::FullNsid(nsid.to_string())); + return Some(Child::FullNsid(nsid.clone())); } let suffix = nsid.as_str().strip_prefix(&format!("{}.", prefix.0))?; let (segment, _) = suffix.split_once('.').unwrap(); @@ -742,17 +734,17 @@ impl FjallReader { } fn is_before(&self, other: &Child) -> bool { match (self, other) { - (Child::FullNsid(s), Child::ChildPrefix(o)) if s == o => true, - (Child::ChildPrefix(s), Child::FullNsid(o)) if s == o => false, - (Child::FullNsid(s), Child::FullNsid(o)) => s < o, + (Child::FullNsid(s), Child::ChildPrefix(o)) if s.as_str() == o => true, + (Child::ChildPrefix(s), Child::FullNsid(o)) if s == o.as_str() => false, + (Child::FullNsid(s), Child::FullNsid(o)) => s.as_str() < o.as_str(), (Child::ChildPrefix(s), Child::ChildPrefix(o)) => s < o, - (Child::FullNsid(s), Child::ChildPrefix(o)) => s < o, - (Child::ChildPrefix(s), Child::FullNsid(o)) => s < o, + (Child::FullNsid(s), Child::ChildPrefix(o)) => s.to_string() < *o, + (Child::ChildPrefix(s), Child::FullNsid(o)) => *s < o.to_string(), } } fn into_inner(self) -> String { match self { - Child::FullNsid(s) => s, + Child::FullNsid(s) => s.to_string(), Child::ChildPrefix(s) => s, } } @@ -791,16 +783,10 @@ impl FjallReader { } } items.push(match child { - Child::FullNsid(nsid) => PrefixChild::Collection(NsidCount { - nsid, - creates: merged.counts().creates, - dids_estimate: merged.dids().estimate() as u64, - }), - Child::ChildPrefix(prefix) => PrefixChild::Prefix(PrefixCount { - prefix, - creates: merged.counts().creates, - dids_estimate: merged.dids().estimate() as u64, - }), + Child::FullNsid(nsid) => PrefixChild::Collection(NsidCount::new(&nsid, &merged)), + Child::ChildPrefix(prefix) => { + PrefixChild::Prefix(PrefixCount::new(&prefix, &merged)) + } }); } @@ -991,15 +977,11 @@ impl FjallReader { for kv in self.rollups.range((start, end)) { let (key_bytes, val_bytes) = kv?; let key = db_complete::(&key_bytes)?; - let nsid = key.collection().as_str().to_string(); + let nsid = key.collection(); for term in &terms { if nsid.contains(term) { let counts = db_complete::(&val_bytes)?; - matches.push(NsidCount { - nsid: nsid.clone(), - creates: counts.counts().creates, - dids_estimate: counts.dids().estimate() as u64, - }); + matches.push(NsidCount::new(nsid, &counts)); break; } } @@ -2649,6 +2631,8 @@ mod tests { vec![PrefixChild::Collection(NsidCount { nsid: "a.a.a".to_string(), creates: 1, + updates: 0, + deletes: 0, dids_estimate: 1 }),] ); @@ -2695,7 +2679,9 @@ mod tests { vec![PrefixChild::Prefix(PrefixCount { prefix: "a.a.a".to_string(), creates: 1, - dids_estimate: 1 + updates: 0, + deletes: 0, + dids_estimate: 1, }),] ); assert_eq!(cursor, None); @@ -2750,6 +2736,8 @@ mod tests { vec![PrefixChild::Prefix(PrefixCount { prefix: "a.a.a".to_string(), creates: 2, + updates: 0, + deletes: 0, dids_estimate: 1 }),] ); @@ -2818,11 +2806,15 @@ mod tests { PrefixChild::Collection(NsidCount { nsid: "a.a.a.a".to_string(), creates: 1, + updates: 0, + deletes: 0, dids_estimate: 1 }), PrefixChild::Prefix(PrefixCount { prefix: "a.a.a.a".to_string(), creates: 1, + updates: 0, + deletes: 0, dids_estimate: 1 }), ] -- 2.51.2