1 files changed, 57 insertions, 18 deletions

diff --git a/database/src/kv/mod.rs b/database/src/kv/mod.rs
index 2d2f3cc..9cb05fc 100644
--- a/database/src/kv/mod.rs
+++ b/database/src/kv/mod.rs
@@ -8,17 +8,19 @@ pub mod binning;
 pub mod helpers;
 pub mod index;
 pub mod index_key;
+pub mod merge_iterator;
 pub mod prefix_iterator;
 pub mod sort;
 
 use std::borrow::Cow;
 
 use crate::{
-    Database, Query, RowIter, RowNum, Transaction,
+    Database, Query, RowNum, Transaction,
     kv::{
         helpers::{read_counter, write_counter},
-        index::{read_count_index, update_index},
+        index::{iter_index, read_count_index, update_index},
         index_key::{IndexKey, SortKey},
+        merge_iterator::MergeIterator,
         prefix_iterator::PrefixIterator,
     },
 };
@@ -91,27 +93,36 @@ impl Transaction for &mut dyn jellykv::Transaction {
         Ok(jellykv::Transaction::get(*self, &row_key(row))?.map(ObjectBuffer::from))
     }
 
-    fn query(&mut self, query: Query) -> Result<RowIter> {
-        // query
-        //     .filter
-        //     .get_bins()
-        //     .into_iter()
-        //     .flat_map(|b| {
-        //         let ikey = (b, query.sort.key());
-        //         self.indices.get(&ikey)
-        //     })
-        //     .map(|i| i.query(txn, &query.sort))
-        todo!()
+    fn query<'a>(
+        &'a mut self,
+        query: Query,
+    ) -> Result<Box<dyn Iterator<Item = Result<(RowNum, Vec<u8>)>> + 'a>> {
+        let mut prefixes = Vec::new();
+        for (binning, mut prefix) in query.filter.get_bins() {
+            let ik = IndexKey(binning, query.sort.key());
+            let is = get_or_create_index(*self, &ik)?;
+            prefix.splice(0..0, is.to_be_bytes());
+            prefixes.push(prefix);
+        }
+        let mut iters = Vec::new();
+        for prefix in prefixes {
+            iters.push(iter_index(*self, prefix, &query.sort)?)
+        }
+        if iters.len() == 1 {
+            return Ok(iters.pop().unwrap());
+        }
+        Ok(Box::new(MergeIterator::new(iters)))
     }
     fn query_single(&mut self, query: Query) -> Result<Option<RowNum>> {
-        self.query(query)?.next().transpose()
+        Ok(self.query(query)?.next().transpose()?.map(|(e, _)| e))
     }
     fn count(&mut self, query: Query) -> Result<u64> {
         let mut total = 0;
-        for (binning, prefix) in query.filter.get_bins() {
+        for (binning, mut prefix) in query.filter.get_bins() {
             let ik = IndexKey(binning, SortKey::Count);
             let is = get_or_create_index(*self, &ik)?;
-            total += read_count_index(*self, is, prefix)?;
+            prefix.splice(0..0, is.to_be_bytes());
+            total += read_count_index(*self, prefix)?;
         }
         Ok(total)
     }
@@ -183,10 +194,11 @@ fn inv_row_key(k: &[u8]) -> RowNum {
 #[cfg(test)]
 mod test {
     use crate::{
-        Database,
-        test_shared::{NAME, new_bob},
+        Database, Filter, Query, Sort,
+        test_shared::{AGE, NAME, new_alice, new_bob, new_charlie},
     };
     use anyhow::Result;
+    use jellyobject::Path;
 
     #[test]
     pub fn insert_get() -> Result<()> {
@@ -232,4 +244,31 @@ mod test {
         assert_eq!(bob.unwrap().as_object().get(NAME).unwrap(), "Better Bob");
         Ok(())
     }
+
+    #[test]
+    pub fn filter() -> Result<()> {
+        let db = jellykv::memory::new();
+
+        let mut rows = [0, 0, 0];
+        let mut result = None;
+
+        db.transaction(&mut |txn| {
+            rows = [
+                txn.insert(new_bob())?,
+                txn.insert(new_alice())?,
+                txn.insert(new_charlie())?,
+            ];
+            Ok(())
+        })?;
+        db.transaction(&mut |txn| {
+            result = txn.query_single(Query {
+                filter: Filter::Match(Path(vec![AGE.0]), 35_u32.to_be_bytes().to_vec()),
+                sort: Sort::None,
+            })?;
+            Ok(())
+        })?;
+
+        assert_eq!(result.unwrap(), 0);
+        Ok(())
+    }
 }