[query] Implement select key count #556

src/agdb/collections/multi_map.rs

@@ -234,6 +234,30 @@ where
         Ok(values)
     }
 
+    pub fn values_count(&self, key: &K) -> Result<u64, DbError> {
+        if self.capacity() == 0 {
+            return Ok(0);
+        }
+
+        let hash = key.stable_hash();
+        let mut pos = hash % self.capacity();
+        let mut result = 0;
+
+        loop {
+            match self.data.state(pos)? {
+                MapValueState::Empty => break,
+                MapValueState::Valid if self.data.key(pos)? == *key => {
+                    result += 1;
+                }
+                MapValueState::Valid | MapValueState::Deleted => {}
+            }
+
+            pos = self.next_pos(pos)
+        }
+
+        Ok(result)
+    }
+
     fn drop_value(&mut self, pos: u64) -> Result<(), DbError> {
         self.data.set_state(pos, MapValueState::Deleted)?;
         self.data.set_key(pos, &K::default())?;
@@ -527,6 +551,30 @@ mod tests {
             .is_ok());
     }
 
+    #[test]
+    fn values_count() {
+        let test_file = TestFile::new();
+        let storage = Rc::new(RefCell::new(
+            FileStorage::new(test_file.file_name()).unwrap(),
+        ));
+
+        let mut map = MultiMapStorage::<u64, String>::new(storage).unwrap();
+
+        assert_eq!(map.values_count(&4), Ok(0));
+
+        map.insert(&1, &"Hello".to_string()).unwrap();
+        map.insert(&1, &"World".to_string()).unwrap();
+        map.insert(&1, &"!".to_string()).unwrap();
+        map.insert(&2, &"a".to_string()).unwrap();
+        map.insert(&3, &"b".to_string()).unwrap();
+        map.remove_value(&1, &"World".to_string()).unwrap();
+
+        assert_eq!(map.values_count(&1), Ok(2));
+        assert_eq!(map.values_count(&2), Ok(1));
+        assert_eq!(map.values_count(&3), Ok(1));
+        assert_eq!(map.values_count(&4), Ok(0));
+    }
+
     #[test]
     fn from_storage() {
         let test_file = TestFile::new();

src/agdb/db.rs

@@ -257,6 +257,10 @@ impl Db {
             .collect())
     }
 
+    pub(crate) fn key_count(&self, db_id: DbId) -> Result<u64, DbError> {
+        self.values.values_count(&db_id)
+    }
+
     pub(crate) fn remove(&mut self, query_id: &QueryId) -> Result<bool, QueryError> {
         match query_id {
             QueryId::Id(db_id) => {

src/agdb/query/select_key_count_query.rs

@@ -1,3 +1,30 @@
 use super::query_ids::QueryIds;
+use crate::Db;
+use crate::DbElement;
+use crate::Query;
+use crate::QueryError;
+use crate::QueryResult;
 
 pub struct SelectKeyCountQuery(pub QueryIds);
+
+impl Query for SelectKeyCountQuery {
+    fn process(&self, db: &Db, result: &mut QueryResult) -> Result<(), QueryError> {
+        match &self.0 {
+            QueryIds::Ids(ids) => {
+                result.elements.reserve(ids.len());
+                result.result += ids.len() as i64;
+
+                for id in ids {
+                    let db_id = db.db_id(id)?;
+                    result.elements.push(DbElement {
+                        index: db_id,
+                        values: vec![("key_count", db.key_count(db_id)?).into()],
+                    });
+                }
+
+                Ok(())
+            }
+            QueryIds::Search(_) => Err(QueryError::from("Invalid select key count query")),
+        }
+    }
+}

tests/select_key_count_test.rs

@@ -1,17 +1,67 @@
+mod test_db;
+
+use agdb::DbElement;
+use agdb::DbId;
 use agdb::QueryBuilder;
+use test_db::TestDb;
 
 #[test]
 fn select_key_count_ids() {
-    let _query = QueryBuilder::select()
-        .key_count()
-        .ids(&["alias".into()])
-        .query();
+    let mut db = TestDb::new();
+    db.exec_mut(
+        QueryBuilder::insert()
+            .nodes()
+            .aliases(&["alias".into()])
+            .values(&[&[
+                ("key", 100).into(),
+                (1, "value").into(),
+                (vec![1.1_f64], 1).into(),
+            ]])
+            .query(),
+        1,
+    );
+    db.exec_elements(
+        QueryBuilder::select()
+            .key_count()
+            .ids(&["alias".into()])
+            .query(),
+        &[DbElement {
+            index: DbId(1),
+            values: vec![("key_count", 3_u64).into()],
+        }],
+    );
+}
+
+#[test]
+fn select_keys_count_no_keys() {
+    let mut db = TestDb::new();
+    db.exec_mut(
+        QueryBuilder::insert()
+            .nodes()
+            .aliases(&["alias".into()])
+            .query(),
+        1,
+    );
+    db.exec_elements(
+        QueryBuilder::select()
+            .key_count()
+            .ids(&["alias".into()])
+            .query(),
+        &[DbElement {
+            index: DbId(1),
+            values: vec![("key_count", 0_u64).into()],
+        }],
+    );
 }
 
 #[test]
-fn select_key_count_search() {
-    let _query = QueryBuilder::select()
-        .key_count()
-        .search(QueryBuilder::search().from("alias".into()).query())
-        .query();
+fn select_keys_search() {
+    let db = TestDb::new();
+    db.exec_error(
+        QueryBuilder::select()
+            .key_count()
+            .search(QueryBuilder::search().from("alias".into()).query())
+            .query(),
+        "Invalid select key count query",
+    );
 }