From 0e3fb7b3656a956939567b0c558a99484c09e9e0 Mon Sep 17 00:00:00 2001 From: Emmanuel Keller Date: Tue, 12 Mar 2024 08:58:22 +0000 Subject: [PATCH] Bug fix: Implements the union strategy on unique indexes (#3674) --- core/src/dbs/processor.rs | 4 +++ core/src/idx/planner/executor.rs | 4 +++ core/src/idx/planner/iterators.rs | 41 +++++++++++++++++++++ lib/tests/planner.rs | 59 ++++++++++++++++++++++++++++++- 4 files changed, 107 insertions(+), 1 deletion(-) diff --git a/core/src/dbs/processor.rs b/core/src/dbs/processor.rs index 22191d23..7292e205 100644 --- a/core/src/dbs/processor.rs +++ b/core/src/dbs/processor.rs @@ -605,6 +605,10 @@ impl<'a> Processor<'a> { } // Everything ok return Ok(()); + } else { + return Err(Error::QueryNotExecutedDetail { + message: "No Iterator has been found.".to_string(), + }); } } Err(Error::QueryNotExecutedDetail { diff --git a/core/src/idx/planner/executor.rs b/core/src/idx/planner/executor.rs index 4eb86b86..a9d521e7 100644 --- a/core/src/idx/planner/executor.rs +++ b/core/src/idx/planner/executor.rs @@ -10,6 +10,7 @@ use crate::idx::ft::{FtIndex, MatchRef}; use crate::idx::planner::iterators::{ DocIdsIterator, IndexEqualThingIterator, IndexRangeThingIterator, IndexUnionThingIterator, MatchesThingIterator, ThingIterator, UniqueEqualThingIterator, UniqueRangeThingIterator, + UniqueUnionThingIterator, }; use crate::idx::planner::knn::KnnPriorityList; use crate::idx::planner::plan::IndexOperator::Matches; @@ -338,6 +339,9 @@ impl QueryExecutor { IndexOperator::Equality(value) => { Some(ThingIterator::UniqueEqual(UniqueEqualThingIterator::new(opt, ix, value))) } + IndexOperator::Union(value) => { + Some(ThingIterator::UniqueUnion(UniqueUnionThingIterator::new(opt, ix, value))) + } _ => None, } } diff --git a/core/src/idx/planner/iterators.rs b/core/src/idx/planner/iterators.rs index cd717db8..63cd4218 100644 --- a/core/src/idx/planner/iterators.rs +++ b/core/src/idx/planner/iterators.rs @@ -18,6 +18,7 @@ pub(crate) enum ThingIterator { IndexUnion(IndexUnionThingIterator), UniqueEqual(UniqueEqualThingIterator), UniqueRange(UniqueRangeThingIterator), + UniqueUnion(UniqueUnionThingIterator), Matches(MatchesThingIterator), Knn(DocIdsIterator), } @@ -34,6 +35,7 @@ impl ThingIterator { ThingIterator::IndexRange(i) => i.next_batch(tx, size).await, ThingIterator::UniqueRange(i) => i.next_batch(tx, size).await, ThingIterator::IndexUnion(i) => i.next_batch(tx, size).await, + ThingIterator::UniqueUnion(i) => i.next_batch(tx, size).await, ThingIterator::Matches(i) => i.next_batch(tx, size).await, ThingIterator::Knn(i) => i.next_batch(tx, size).await, } @@ -368,6 +370,45 @@ impl UniqueRangeThingIterator { } } +pub(crate) struct UniqueUnionThingIterator { + keys: VecDeque, +} + +impl UniqueUnionThingIterator { + pub(super) fn new(opt: &Options, ix: &DefineIndexStatement, a: &Array) -> Self { + // We create a VecDeque to hold the key for each value in the array. + let keys: VecDeque = + a.0.iter() + .map(|v| { + let a = Array::from(v.clone()); + let key = Index::new(opt.ns(), opt.db(), &ix.what, &ix.name, &a, None).into(); + key + }) + .collect(); + Self { + keys, + } + } + + async fn next_batch( + &mut self, + txn: &Transaction, + limit: u32, + ) -> Result)>, Error> { + let mut run = txn.lock().await; + let mut res = vec![]; + while let Some(key) = self.keys.pop_front() { + if let Some(val) = run.get(key).await? { + res.push((val.into(), None)); + } + if res.len() >= limit as usize { + return Ok(res); + } + } + Ok(res) + } +} + pub(crate) struct MatchesThingIterator { hits: Option, } diff --git a/lib/tests/planner.rs b/lib/tests/planner.rs index d553444d..20c44742 100644 --- a/lib/tests/planner.rs +++ b/lib/tests/planner.rs @@ -1312,6 +1312,63 @@ async fn select_with_in_operator() -> Result<(), Error> { ); assert_eq!(format!("{:#}", tmp), format!("{:#}", val)); } - + Ok(()) +} + +#[tokio::test] +async fn select_with_in_operator_uniq_index() -> Result<(), Error> { + let dbs = new_ds().await?; + let ses = Session::owner().with_ns("test").with_db("test"); + + let sql = r#" + DEFINE INDEX apprenantUid ON apprenants FIELDS apprenantUid UNIQUE; + CREATE apprenants:1 CONTENT { apprenantUid: "00013483-fedd-43e3-a94e-80728d896f6e" }; + SELECT apprenantUid FROM apprenants WHERE apprenantUid in []; + SELECT apprenantUid FROM apprenants WHERE apprenantUid IN ["00013483-fedd-43e3-a94e-80728d896f6e"]; + SELECT apprenantUid FROM apprenants WHERE apprenantUid IN ["99999999-aaaa-1111-8888-abcdef012345", "00013483-fedd-43e3-a94e-80728d896f6e"]; + SELECT apprenantUid FROM apprenants WHERE apprenantUid IN ["00013483-fedd-43e3-a94e-80728d896f6e", "99999999-aaaa-1111-8888-abcdef012345"]; + SELECT apprenantUid FROM apprenants WHERE apprenantUid IN ["99999999-aaaa-1111-8888-abcdef012345", "00013483-fedd-43e3-a94e-80728d896f6e", "99999999-aaaa-1111-8888-abcdef012345"]; + SELECT apprenantUid FROM apprenants WHERE apprenantUid IN ["00013483-fedd-43e3-a94e-80728d896f6e"] EXPLAIN; + "#; + let mut res = dbs.execute(&sql, &ses, None).await?; + + assert_eq!(res.len(), 8); + skip_ok(&mut res, 2)?; + + let tmp = res.remove(0).result?; + let val = Value::parse(r#"[]"#); + assert_eq!(format!("{:#}", tmp), format!("{:#}", val)); + + for _ in 0..4 { + let tmp = res.remove(0).result?; + let val = Value::parse( + r#"[ + { + apprenantUid: '00013483-fedd-43e3-a94e-80728d896f6e' + } + ]"#, + ); + assert_eq!(format!("{:#}", tmp), format!("{:#}", val)); + } + + let tmp = res.remove(0).result?; + let val = Value::parse( + r#"[ + { + detail: { + plan: { + index: 'apprenantUid', + operator: 'union', + value: [ + '00013483-fedd-43e3-a94e-80728d896f6e' + ] + }, + table: 'apprenants' + }, + operation: 'Iterate Index' + } + ]"#, + ); + assert_eq!(format!("{:#}", tmp), format!("{:#}", val)); Ok(()) }