From cad596cdf3f29eb3e9ea1100f1c02d51771c2cac Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@surrealdb.com>
Date: Sun, 26 Mar 2023 10:04:18 +0100
Subject: [PATCH] Ensure GROUP BY clauses with multi same-aggregate functions
 work correctly

Closes #1731
---
 lib/src/dbs/iterator.rs |   5 +-
 lib/src/sql/field.rs    |   2 +-
 lib/tests/group.rs      | 104 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 106 insertions(+), 5 deletions(-)

diff --git a/lib/src/dbs/iterator.rs b/lib/src/dbs/iterator.rs
index 726c5754..4aec0351 100644
--- a/lib/src/dbs/iterator.rs
+++ b/lib/src/dbs/iterator.rs
@@ -241,10 +241,7 @@ impl Iterator {
 						if let Field::Alias(v, i) = field {
 							match v {
 								Value::Function(f) if f.is_aggregate() => {
-									let x = vals
-										.all()
-										.get(ctx, opt, txn, v.to_idiom().as_ref())
-										.await?;
+									let x = vals.all().get(ctx, opt, txn, i).await?;
 									let x = f.aggregate(x).compute(ctx, opt, txn, None).await?;
 									obj.set(ctx, opt, txn, i, x).await?;
 								}
diff --git a/lib/src/sql/field.rs b/lib/src/sql/field.rs
index 970a977e..91077d2c 100644
--- a/lib/src/sql/field.rs
+++ b/lib/src/sql/field.rs
@@ -151,7 +151,7 @@ impl Fields {
 						};
 						// Check if this is a single VALUE field expression
 						match self.single().is_some() {
-							false => out.set(ctx, opt, txn, v.to_idiom().as_ref(), x).await?,
+							false => out.set(ctx, opt, txn, i, x).await?,
 							true => out = x,
 						}
 					}
diff --git a/lib/tests/group.rs b/lib/tests/group.rs
index 07c56bb6..1ea416fa 100644
--- a/lib/tests/group.rs
+++ b/lib/tests/group.rs
@@ -233,3 +233,107 @@ async fn select_limit_fetch() -> Result<(), Error> {
 	//
 	Ok(())
 }
+
+#[tokio::test]
+async fn select_multi_aggregate() -> Result<(), Error> {
+	let sql = "
+		CREATE test:1 SET group = 1, one = 1.7, two = 2.4;
+		CREATE test:2 SET group = 1, one = 4.7, two = 3.9;
+		CREATE test:3 SET group = 2, one = 3.2, two = 9.7;
+		CREATE test:4 SET group = 2, one = 4.4, two = 3.0;
+		SELECT group, math::sum(one) AS one, math::sum(two) AS two FROM test GROUP BY group;
+		SELECT group, math::sum(two) AS two, math::sum(one) AS one FROM test GROUP BY group;
+	";
+	let dbs = Datastore::new("memory").await?;
+	let ses = Session::for_kv().with_ns("test").with_db("test");
+	let res = &mut dbs.execute(&sql, &ses, None, false).await?;
+	assert_eq!(res.len(), 6);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: test:1,
+				group: 1,
+				one: 1.7,
+				two: 2.4,
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: test:2,
+				group: 1,
+				one: 4.7,
+				two: 3.9,
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: test:3,
+				group: 2,
+				one: 3.2,
+				two: 9.7,
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: test:4,
+				group: 2,
+				one: 4.4,
+				two: 3.0,
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				group: 1,
+				one: 6.4,
+				two: 6.3,
+			},
+			{
+				group: 2,
+				one: 7.6,
+				two: 12.7,
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				group: 1,
+				one: 6.4,
+				two: 6.3,
+			},
+			{
+				group: 2,
+				one: 7.6,
+				two: 12.7,
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	Ok(())
+}