Added TryFloatDiv trait for MathMean in Aggregate function (#4390)

Co-authored-by: Tobie Morgan Hitchcock <tobie@surrealdb.com>
This commit is contained in:
Anirban Halder 2024-09-05 03:12:38 +05:30 committed by GitHub
parent 0a1a99dc68
commit 405c96f78d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 197 additions and 107 deletions

View file

@ -4,7 +4,7 @@ use crate::dbs::store::MemoryCollector;
use crate::dbs::{Options, Statement};
use crate::err::Error;
use crate::sql::function::OptimisedAggregate;
use crate::sql::value::{TryAdd, TryDiv, Value};
use crate::sql::value::{TryAdd, TryFloatDiv, Value};
use crate::sql::{Array, Field, Function, Idiom};
use reblessive::tree::Stk;
use std::borrow::Cow;
@ -329,7 +329,7 @@ impl Aggregator {
OptimisedAggregate::MathSum => self.math_sum.take().unwrap_or(Value::None),
OptimisedAggregate::MathMean => {
if let Some((v, i)) = self.math_mean.take() {
v.try_div(i.into()).unwrap_or(f64::NAN.into())
v.try_float_div(i.into()).unwrap_or(f64::NAN.into())
} else {
Value::None
}

View file

@ -388,29 +388,29 @@ impl Document {
Value::Function(f) if f.is_rolling() => match f.name() {
Some("count") => {
let val = f.compute(stk, ctx, opt, Some(fdc.doc)).await?;
self.chg(&mut set_ops, &mut del_ops, &fdc.act, idiom, val);
self.chg(&mut set_ops, &mut del_ops, &fdc.act, idiom, val)?;
}
Some("math::sum") => {
let val = f.args()[0].compute(stk, ctx, opt, Some(fdc.doc)).await?;
self.chg(&mut set_ops, &mut del_ops, &fdc.act, idiom, val);
self.chg(&mut set_ops, &mut del_ops, &fdc.act, idiom, val)?;
}
Some("math::min") | Some("time::min") => {
let val = f.args()[0].compute(stk, ctx, opt, Some(fdc.doc)).await?;
self.min(&mut set_ops, &mut del_ops, fdc, field, idiom, val);
self.min(&mut set_ops, &mut del_ops, fdc, field, idiom, val)?;
}
Some("math::max") | Some("time::max") => {
let val = f.args()[0].compute(stk, ctx, opt, Some(fdc.doc)).await?;
self.max(&mut set_ops, &mut del_ops, fdc, field, idiom, val);
self.max(&mut set_ops, &mut del_ops, fdc, field, idiom, val)?;
}
Some("math::mean") => {
let val = f.args()[0].compute(stk, ctx, opt, Some(fdc.doc)).await?;
self.mean(&mut set_ops, &mut del_ops, &fdc.act, idiom, val);
self.mean(&mut set_ops, &mut del_ops, &fdc.act, idiom, val)?;
}
_ => unreachable!(),
},
_ => {
let val = expr.compute(stk, ctx, opt, Some(fdc.doc)).await?;
self.set(&mut set_ops, idiom, val);
self.set(&mut set_ops, idiom, val)?;
}
}
}
@ -419,11 +419,19 @@ impl Document {
}
/// Set the field in the foreign table
fn set(&self, ops: &mut Ops, key: Idiom, val: Value) {
fn set(&self, ops: &mut Ops, key: Idiom, val: Value) -> Result<(), Error> {
ops.push((key, Operator::Equal, val));
Ok(())
}
/// Increment or decrement the field in the foreign table
fn chg(&self, set_ops: &mut Ops, del_ops: &mut Ops, act: &FieldAction, key: Idiom, val: Value) {
fn chg(
&self,
set_ops: &mut Ops,
del_ops: &mut Ops,
act: &FieldAction,
key: Idiom,
val: Value,
) -> Result<(), Error> {
match act {
FieldAction::Add => {
set_ops.push((key.clone(), Operator::Inc, val));
@ -434,6 +442,7 @@ impl Document {
del_ops.push((key, Operator::Equal, Value::from(0)));
}
}
Ok(())
}
/// Set the new minimum value for the field in the foreign table
@ -445,7 +454,7 @@ impl Document {
field: &Field,
key: Idiom,
val: Value,
) {
) -> Result<(), Error> {
// Key for the value count
let mut key_c = Idiom::from(vec![Part::from("__")]);
key_c.0.push(Part::from(key.to_hash()));
@ -489,6 +498,7 @@ impl Document {
del_ops.push((key_c, Operator::Equal, Value::from(0)));
}
}
Ok(())
}
/// Set the new maximum value for the field in the foreign table
fn max(
@ -499,7 +509,7 @@ impl Document {
field: &Field,
key: Idiom,
val: Value,
) {
) -> Result<(), Error> {
// Key for the value count
let mut key_c = Idiom::from(vec![Part::from("__")]);
key_c.0.push(Part::from(key.to_hash()));
@ -544,6 +554,85 @@ impl Document {
del_ops.push((key_c, Operator::Equal, Value::from(0)));
}
}
Ok(())
}
/// Set the new average value for the field in the foreign table
fn mean(
&self,
set_ops: &mut Ops,
del_ops: &mut Ops,
act: &FieldAction,
key: Idiom,
val: Value,
) -> Result<(), Error> {
// Key for the value count
let mut key_c = Idiom::from(vec![Part::from("__")]);
key_c.0.push(Part::from(key.to_hash()));
key_c.0.push(Part::from("c"));
//
set_ops.push((
key.clone(),
Operator::Equal,
Value::Expression(Box::new(Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(
Box::new(Expression::Binary {
l: Value::Idiom(key),
o: Operator::Nco,
r: Value::Number(Number::Int(0)),
}),
)))),
o: Operator::Mul,
r: Value::Subquery(Box::new(Subquery::Value(Value::Expression(
Box::new(Expression::Binary {
l: Value::Idiom(key_c.clone()),
o: Operator::Nco,
r: Value::Number(Number::Int(0)),
}),
)))),
},
))))),
o: match act {
FieldAction::Sub => Operator::Sub,
FieldAction::Add => Operator::Add,
},
r: val.convert_to_decimal()?.into(),
},
))))),
o: Operator::Div,
r: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Idiom(key_c.clone()),
o: Operator::Nco,
r: Value::Number(Number::Int(0)),
},
))))),
o: match act {
FieldAction::Sub => Operator::Sub,
FieldAction::Add => Operator::Add,
},
r: Value::from(1),
},
))))),
})),
));
match act {
// Increment the number of values
FieldAction::Add => set_ops.push((key_c, Operator::Inc, Value::from(1))),
FieldAction::Sub => {
// Decrement the number of values
set_ops.push((key_c.clone(), Operator::Dec, Value::from(1)));
// Add a purge condition (delete record if the number of values is 0)
del_ops.push((key_c, Operator::Equal, Value::from(0)));
}
}
Ok(())
}
/// Recomputes the value for one group
@ -614,81 +703,4 @@ impl Document {
close: Some(Value::Idiom(key.clone())),
})))
}
/// Set the new average value for the field in the foreign table
fn mean(
&self,
set_ops: &mut Ops,
del_ops: &mut Ops,
act: &FieldAction,
key: Idiom,
val: Value,
) {
// Key for the value count
let mut key_c = Idiom::from(vec![Part::from("__")]);
key_c.0.push(Part::from(key.to_hash()));
key_c.0.push(Part::from("c"));
//
set_ops.push((
key.clone(),
Operator::Equal,
Value::Expression(Box::new(Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(
Box::new(Expression::Binary {
l: Value::Idiom(key),
o: Operator::Nco,
r: Value::Number(Number::Int(0)),
}),
)))),
o: Operator::Mul,
r: Value::Subquery(Box::new(Subquery::Value(Value::Expression(
Box::new(Expression::Binary {
l: Value::Idiom(key_c.clone()),
o: Operator::Nco,
r: Value::Number(Number::Int(0)),
}),
)))),
},
))))),
o: match act {
FieldAction::Sub => Operator::Sub,
FieldAction::Add => Operator::Add,
},
r: val,
},
))))),
o: Operator::Div,
r: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Subquery(Box::new(Subquery::Value(Value::Expression(Box::new(
Expression::Binary {
l: Value::Idiom(key_c.clone()),
o: Operator::Nco,
r: Value::Number(Number::Int(0)),
},
))))),
o: match act {
FieldAction::Sub => Operator::Sub,
FieldAction::Add => Operator::Add,
},
r: Value::from(1),
},
))))),
})),
));
match act {
// Increment the number of values
FieldAction::Add => set_ops.push((key_c, Operator::Inc, Value::from(1))),
FieldAction::Sub => {
// Decrement the number of values
set_ops.push((key_c.clone(), Operator::Dec, Value::from(1)));
// Add a purge condition (delete record if the number of values is 0)
del_ops.push((key_c, Operator::Equal, Value::from(0)));
}
}
}
}

View file

@ -1,4 +1,4 @@
use super::value::{TryAdd, TryDiv, TryMul, TryNeg, TryPow, TryRem, TrySub};
use super::value::{TryAdd, TryDiv, TryFloatDiv, TryMul, TryNeg, TryPow, TryRem, TrySub};
use crate::err::Error;
use crate::fnc::util::math::ToFloat;
use crate::sql::strand::Strand;
@ -729,6 +729,24 @@ impl TryNeg for Number {
}
}
impl TryFloatDiv for Number {
type Output = Self;
fn try_float_div(self, other: Self) -> Result<Self, Error> {
Ok(match (self, other) {
(Number::Int(v), Number::Int(w)) => {
let quotient = (v as f64).div(w as f64);
if quotient.fract() != 0.0 {
return Ok(Number::Float(quotient));
}
Number::Int(
v.checked_div(w).ok_or_else(|| Error::TryDiv(v.to_string(), w.to_string()))?,
)
}
(v, w) => v.try_div(w)?,
})
}
}
impl ops::Add for Number {
type Output = Self;
fn add(self, other: Self) -> Self {
@ -912,3 +930,20 @@ impl ToFloat for Number {
self.to_float()
}
}
#[cfg(test)]
mod tests {
use super::Number;
use super::TryFloatDiv;
#[test]
fn test_try_float_div() {
let (sum_one, count_one) = (Number::Int(5), Number::Int(2));
assert_eq!(sum_one.try_float_div(count_one).unwrap(), Number::Float(2.5));
let (sum_two, count_two) = (Number::Int(10), Number::Int(5));
assert_eq!(sum_two.try_float_div(count_two).unwrap(), Number::Int(2));
let (sum_three, count_three) = (Number::Float(6.3), Number::Int(3));
assert_eq!(sum_three.try_float_div(count_three).unwrap(), Number::Float(2.1));
}
}

View file

@ -3031,6 +3031,23 @@ impl TryDiv for Value {
// ------------------------------
pub(crate) trait TryFloatDiv<Rhs = Self> {
type Output;
fn try_float_div(self, v: Self) -> Result<Self::Output, Error>;
}
impl TryFloatDiv for Value {
type Output = Self;
fn try_float_div(self, other: Self) -> Result<Self::Output, Error> {
Ok(match (self, other) {
(Self::Number(v), Self::Number(w)) => Self::Number(v.try_float_div(w)?),
(v, w) => return Err(Error::TryDiv(v.to_raw_string(), w.to_raw_string())),
})
}
}
// ------------------------------
pub(crate) trait TryRem<Rhs = Self> {
type Output;
fn try_rem(self, v: Self) -> Result<Self::Output, Error>;

View file

@ -2256,7 +2256,9 @@ async fn function_math_mean() -> Result<(), Error> {
let sql = r#"
RETURN math::mean([]);
RETURN math::mean([101, 213, 202]);
RETURN math::mean([101.5, 213.5, 202.5]);
RETURN math::mean([101, 213, 203]);
RETURN math::mean([101, 213, 203.4]);
RETURN math::mean([101.5, 213.5, 206.5]);
"#;
let mut test = Test::new(sql).await?;
//
@ -2268,7 +2270,15 @@ async fn function_math_mean() -> Result<(), Error> {
assert_eq!(tmp, val);
//
let tmp = test.next()?.result?;
let val = Value::from(172.5);
let val = Value::from(172.33333333333334);
assert_eq!(tmp, val);
//
let tmp = test.next()?.result?;
let val = Value::from(172.46666666666667);
assert_eq!(tmp, val);
//
let tmp = test.next()?.result?;
let val = Value::from(173.83333333333334);
assert_eq!(tmp, val);
//
Ok(())

View file

@ -23,11 +23,11 @@ async fn define_foreign_table() -> Result<(), Error> {
GROUP BY age
;
INFO FOR TABLE person;
UPSERT person:one SET age = 39, score = 70;
UPSERT person:one SET age = 39, score = 72;
SELECT * FROM person_by_age;
UPSERT person:two SET age = 39, score = 80;
UPSERT person:two SET age = 39, score = 83;
SELECT * FROM person_by_age;
UPSERT person:two SET age = 39, score = 90;
UPSERT person:two SET age = 39, score = 91;
SELECT * FROM person_by_age;
";
let dbs = new_ds().await?;
@ -59,7 +59,7 @@ async fn define_foreign_table() -> Result<(), Error> {
{
age: 39,
id: person:one,
score: 70,
score: 72,
}
]",
);
@ -70,11 +70,11 @@ async fn define_foreign_table() -> Result<(), Error> {
"[
{
age: 39,
average: 70,
average: 72,
count: 1,
id: person_by_age:[39],
max: 70,
min: 70,
max: 72,
min: 72,
total: 39
}
]",
@ -82,7 +82,15 @@ async fn define_foreign_table() -> Result<(), Error> {
assert_eq!(tmp, val);
//
let tmp = res.remove(0).result?;
let val = Value::parse("[{ id: person:two, age: 39, score: 80 }]");
let val = Value::parse(
"[
{
age: 39,
id: person:two,
score: 83,
}
]",
);
assert_eq!(tmp, val);
//
let tmp = res.remove(0).result?;
@ -90,11 +98,11 @@ async fn define_foreign_table() -> Result<(), Error> {
"[
{
age: 39,
average: 75,
average: 77.5,
count: 2,
id: person_by_age:[39],
max: 80,
min: 70,
max: 83,
min: 72,
total: 78
}
]",
@ -102,7 +110,15 @@ async fn define_foreign_table() -> Result<(), Error> {
assert_eq!(tmp, val);
//
let tmp = res.remove(0).result?;
let val = Value::parse("[{ id: person:two, age: 39, score: 90 }]");
let val = Value::parse(
"[
{
age: 39,
id: person:two,
score: 91,
}
]",
);
assert_eq!(tmp, val);
//
let tmp = res.remove(0).result?;
@ -110,11 +126,11 @@ async fn define_foreign_table() -> Result<(), Error> {
"[
{
age: 39,
average: 80,
average: 81.5,
count: 2,
id: person_by_age:[39],
max: 90,
min: 70,
max: 91,
min: 72,
total: 78
}
]",