fosk 0.2.0

In-memory SQL-like query engine and lightweight data store for testing and prototyping.
Documentation
use serde_json::Value;

use crate::{
    JsonPrimitive,
    parser::{
        aggregators_helper::{Accumulator, AggregateImpl},
        analyzer::{AnalysisContext, AnalyzerError},
        ast::Function,
    },
};

pub struct AvgImpl;
impl AggregateImpl for AvgImpl {
    fn name(&self) -> &'static str {
        "avg"
    }
    fn infer_type(
        &self,
        fun: &Function,
        ctx: &AnalysisContext,
    ) -> Result<(JsonPrimitive, bool), AnalyzerError> {
        if let Some(arg) = fun.args.first() {
            let (t, _n) = crate::parser::analyzer::TypeInference::infer_scalar(arg, ctx)?;
            match t {
                JsonPrimitive::Int | JsonPrimitive::Float => Ok((JsonPrimitive::Float, true)),
                other => Err(AnalyzerError::FunctionArgMismatch {
                    name: fun.name.clone(),
                    expected: "numeric".into(),
                    got: vec![other],
                }),
            }
        } else {
            Err(AnalyzerError::FunctionArgMismatch {
                name: fun.name.clone(),
                expected: "AVG(arg)".into(),
                got: vec![],
            })
        }
    }
    fn create_accumulator(&self) -> Box<dyn Accumulator> {
        Box::new(AvgAcc { sum: 0.0, cnt: 0 })
    }
}

struct AvgAcc {
    sum: f64,
    cnt: i64,
}
impl Accumulator for AvgAcc {
    fn update(&mut self, args: &[Value]) -> Result<(), AnalyzerError> {
        let [v] = args else {
            return Err(AnalyzerError::FunctionArgMismatch {
                name: "AVG".into(),
                expected: "AVG(expr)".into(),
                got: vec![],
            });
        };
        match v {
            Value::Null => {}
            Value::Number(n) => {
                if let Some(i) = n.as_i64() {
                    self.sum += i as f64;
                    self.cnt += 1;
                } else if let Some(f) = n.as_f64() {
                    self.sum += f;
                    self.cnt += 1;
                } else {
                    return Err(AnalyzerError::Other("AVG got non numeric number".into()));
                }
            }
            _ => return Err(AnalyzerError::Other("AVG got non numeric arg".into())),
        }
        Ok(())
    }
    fn finalize(&self) -> Value {
        if self.cnt == 0 {
            Value::Null
        } else {
            let avg = self.sum / (self.cnt as f64);
            serde_json::Number::from_f64(avg)
                .map(Value::Number)
                .unwrap_or(Value::Null)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::AvgImpl;
    use crate::parser::aggregators_helper::AggregateImpl;
    use serde_json::json;

    #[test]
    fn accumulator_averages_numbers_and_ignores_nulls() {
        let mut acc = AvgImpl.create_accumulator();

        acc.update(&[json!(2)]).unwrap();
        acc.update(&[json!(null)]).unwrap();
        acc.update(&[json!(4.0)]).unwrap();

        assert_eq!(acc.finalize(), json!(3.0));
    }

    #[test]
    fn accumulator_returns_null_when_no_values_seen() {
        let acc = AvgImpl.create_accumulator();

        assert_eq!(acc.finalize(), json!(null));
    }

    #[test]
    fn accumulator_rejects_bad_arg_count_and_non_numeric_values() {
        let mut acc = AvgImpl.create_accumulator();

        assert!(acc.update(&[]).is_err());
        assert!(acc.update(&[json!("bad")]).is_err());
    }
}