Add support for multi-yield graph traversal expressions

lib/src/doc/pluck.rs

@@ -7,6 +7,7 @@ use crate::err::Error;
 use crate::sql::field::Field;
 use crate::sql::idiom::Idiom;
 use crate::sql::output::Output;
+use crate::sql::part::Part;
 use crate::sql::permission::Permission;
 use crate::sql::value::Value;
 
@@ -36,14 +37,79 @@ impl<'a> Document<'a> {
 					for v in v.other() {
 						match v {
 							Field::All => (),
-							Field::Alone(v) => {
+							Field::Alone(v) => match v {
-								let x = v.compute(ctx, opt, txn, Some(&self.current)).await?;
+								// This expression is a multi-output graph traversal
-								out.set(ctx, opt, txn, v.to_idiom().as_ref(), x).await?;
+								Value::Idiom(v) if v.is_multi_yield() => {
-							}
+									// Store the different output yields here
-							Field::Alias(v, i) => {
+									let mut res: Vec<(&[Part], Value)> = Vec::new();
-								let x = v.compute(ctx, opt, txn, Some(&self.current)).await?;
+									// Split the expression by each output alias
-								out.set(ctx, opt, txn, i, x).await?;
+									for v in v.split_inclusive(Idiom::split_multi_yield) {
-							}
+										// Use the last fetched value for each fetch
+										let x = match res.last() {
+											Some((_, r)) => r,
+											None => self.current.as_ref(),
+										};
+										// Continue fetching the next idiom part
+										let x = x
+											.get(ctx, opt, txn, v)
+											.await?
+											.compute(ctx, opt, txn, Some(&self.current))
+											.await?;
+										// Add the result to the temporary store
+										res.push((v, x));
+									}
+									// Assign each fetched yield to the output
+									for (p, x) in res {
+										match p.last().unwrap().alias() {
+											// This is an alias expression part
+											Some(i) => out.set(ctx, opt, txn, i, x).await?,
+											// This is the end of the expression
+											None => out.set(ctx, opt, txn, v, x).await?,
+										}
+									}
+								}
+								// This expression is a normal field expression
+								_ => {
+									let x = v.compute(ctx, opt, txn, Some(&self.current)).await?;
+									out.set(ctx, opt, txn, v.to_idiom().as_ref(), x).await?
+								}
+							},
+							Field::Alias(v, i) => match v {
+								// This expression is a multi-output graph traversal
+								Value::Idiom(v) if v.is_multi_yield() => {
+									// Store the different output yields here
+									let mut res: Vec<(&[Part], Value)> = Vec::new();
+									// Split the expression by each output alias
+									for v in v.split_inclusive(Idiom::split_multi_yield) {
+										// Use the last fetched value for each fetch
+										let x = match res.last() {
+											Some((_, r)) => r,
+											None => self.current.as_ref(),
+										};
+										// Continue fetching the next idiom part
+										let x = x
+											.get(ctx, opt, txn, v)
+											.await?
+											.compute(ctx, opt, txn, Some(&self.current))
+											.await?;
+										// Add the result to the temporary store
+										res.push((v, x));
+									}
+									// Assign each fetched yield to the output
+									for (p, x) in res {
+										match p.last().unwrap().alias() {
+											// This is an alias expression part
+											Some(i) => out.set(ctx, opt, txn, i, x).await?,
+											// This is the end of the expression
+											None => out.set(ctx, opt, txn, i, x).await?,
+										}
+									}
+								}
+								_ => {
+									let x = v.compute(ctx, opt, txn, Some(&self.current)).await?;
+									out.set(ctx, opt, txn, i, x).await?
+								}
+							},
 						}
 					}
 					Ok(out)
@@ -59,9 +125,12 @@ impl<'a> Document<'a> {
 						match v {
 							Field::All => (),
 							Field::Alone(v) => match v {
+								// This expression is a grouped aggregate function
 								Value::Function(f) if s.group.is_some() && f.is_aggregate() => {
 									let x = match f.args().len() {
+										// If no function arguments, then compute the result
 										0 => f.compute(ctx, opt, txn, Some(&self.current)).await?,
+										// If arguments, then pass the first value through
 										_ => {
 											f.args()[0]
 												.compute(ctx, opt, txn, Some(&self.current))
@@ -70,15 +139,49 @@ impl<'a> Document<'a> {
 									};
 									out.set(ctx, opt, txn, v.to_idiom().as_ref(), x).await?;
 								}
+								// This expression is a multi-output graph traversal
+								Value::Idiom(v) if v.is_multi_yield() => {
+									// Store the different output yields here
+									let mut res: Vec<(&[Part], Value)> = Vec::new();
+									// Split the expression by each output alias
+									for v in v.split_inclusive(Idiom::split_multi_yield) {
+										// Use the last fetched value for each fetch
+										let x = match res.last() {
+											Some((_, r)) => r,
+											None => self.current.as_ref(),
+										};
+										// Continue fetching the next idiom part
+										let x = x
+											.get(ctx, opt, txn, v)
+											.await?
+											.compute(ctx, opt, txn, Some(&self.current))
+											.await?;
+										// Add the result to the temporary store
+										res.push((v, x));
+									}
+									// Assign each fetched yield to the output
+									for (p, x) in res {
+										match p.last().unwrap().alias() {
+											// This is an alias expression part
+											Some(i) => out.set(ctx, opt, txn, i, x).await?,
+											// This is the end of the expression
+											None => out.set(ctx, opt, txn, v, x).await?,
+										}
+									}
+								}
+								// This expression is a normal field expression
 								_ => {
 									let x = v.compute(ctx, opt, txn, Some(&self.current)).await?;
 									out.set(ctx, opt, txn, v.to_idiom().as_ref(), x).await?;
 								}
 							},
 							Field::Alias(v, i) => match v {
+								// This expression is a grouped aggregate function
 								Value::Function(f) if s.group.is_some() && f.is_aggregate() => {
 									let x = match f.args().len() {
+										// If no function arguments, then compute the result
 										0 => f.compute(ctx, opt, txn, Some(&self.current)).await?,
+										// If arguments, then pass the first value through
 										_ => {
 											f.args()[0]
 												.compute(ctx, opt, txn, Some(&self.current))
@@ -87,6 +190,37 @@ impl<'a> Document<'a> {
 									};
 									out.set(ctx, opt, txn, i, x).await?;
 								}
+								// This expression is a multi-output graph traversal
+								Value::Idiom(v) if v.is_multi_yield() => {
+									// Store the different output yields here
+									let mut res: Vec<(&[Part], Value)> = Vec::new();
+									// Split the expression by each output alias
+									for v in v.split_inclusive(Idiom::split_multi_yield) {
+										// Use the last fetched value for each fetch
+										let x = match res.last() {
+											Some((_, r)) => r,
+											None => self.current.as_ref(),
+										};
+										// Continue fetching the next idiom part
+										let x = x
+											.get(ctx, opt, txn, v)
+											.await?
+											.compute(ctx, opt, txn, Some(&self.current))
+											.await?;
+										// Add the result to the temporary store
+										res.push((v, x));
+									}
+									// Assign each fetched yield to the output
+									for (p, x) in res {
+										match p.last().unwrap().alias() {
+											// This is an alias expression part
+											Some(i) => out.set(ctx, opt, txn, i, x).await?,
+											// This is the end of the expression
+											None => out.set(ctx, opt, txn, i, x).await?,
+										}
+									}
+								}
+								// This expression is a normal field expression
 								_ => {
 									let x = v.compute(ctx, opt, txn, Some(&self.current)).await?;
 									out.set(ctx, opt, txn, i, x).await?;

lib/src/sql/idiom.rs

@@ -58,17 +58,17 @@ impl From<Vec<Part>> for Idiom {
 }
 
 impl Idiom {
-	///
+	// Appends a part to the end of this Idiom
-	pub fn push(mut self, n: Part) -> Idiom {
+	pub(crate) fn push(mut self, n: Part) -> Idiom {
 		self.0.push(n);
 		self
 	}
-	///
+	// Convert this Idiom to a JSON Path string
-	pub fn to_path(&self) -> String {
+	pub(crate) fn to_path(&self) -> String {
 		format!("/{}", self).replace(']', "").replace(&['.', '['][..], "/")
 	}
-	///
+	// Simplifies this Idiom for use in object keys
-	pub fn simplify(&self) -> Idiom {
+	pub(crate) fn simplify(&self) -> Idiom {
 		self.0
 			.iter()
 			.cloned()
@@ -76,6 +76,14 @@ impl Idiom {
 			.collect::<Vec<_>>()
 			.into()
 	}
+	// Check if this is an expression with multiple yields
+	pub(crate) fn is_multi_yield(&self) -> bool {
+		self.iter().any(Self::split_multi_yield)
+	}
+	// Check if the path part is a yield in a multi-yield expression
+	pub(crate) fn split_multi_yield(v: &Part) -> bool {
+		matches!(v, Part::Graph(g) if g.alias.is_some())
+	}
 }
 
 impl Idiom {

lib/src/sql/part.rs

@@ -2,6 +2,7 @@ use crate::sql::comment::shouldbespace;
 use crate::sql::error::IResult;
 use crate::sql::graph::{graph as graph_raw, Graph};
 use crate::sql::ident::{ident, Ident};
+use crate::sql::idiom::Idiom;
 use crate::sql::number::{number, Number};
 use crate::sql::value::{value, Value};
 use nom::branch::alt;
@@ -80,6 +81,16 @@ impl From<&str> for Part {
 	}
 }
 
+impl Part {
+	// Returns a yield if an alias is specified
+	pub(crate) fn alias(&self) -> Option<&Idiom> {
+		match self {
+			Part::Graph(v) => v.alias.as_ref(),
+			_ => None,
+		}
+	}
+}
+
 impl fmt::Display for Part {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 		match self {