graph_d/query/

aggregation.rs

//! Aggregation operations for graph queries.

use crate::error::Result;
use crate::graph::Node;
use serde_json::Value;
use std::collections::HashMap;

/// Aggregation functions for query results.
#[derive(Debug, Clone, PartialEq)]
pub enum AggregateFunction {
    /// Count the number of items
    Count,
    /// Sum numeric values for the specified property key
    Sum(
        /// Property key to sum
        String,
    ),
    /// Average numeric values for the specified property key
    Avg(
        /// Property key to average
        String,
    ),
    /// Minimum value for the specified property key
    Min(
        /// Property key to find minimum
        String,
    ),
    /// Maximum value for the specified property key
    Max(
        /// Property key to find maximum
        String,
    ),
    /// Collect unique values for the specified property key
    Distinct(
        /// Property key to collect distinct values
        String,
    ),
    /// Group by the specified property value
    GroupBy(
        /// Property key to group by
        String,
    ),
}

/// Result of an aggregation operation.
#[derive(Debug, Clone, PartialEq)]
pub enum AggregateResult {
    /// Single numeric value
    Number(f64),
    /// Single string value
    String(String),
    /// List of unique values
    List(Vec<Value>),
    /// Grouped results
    Groups(HashMap<String, Vec<Node>>),
}

/// Aggregator for performing aggregation operations on node collections.
pub struct Aggregator;

impl Aggregator {
    /// Apply an aggregation function to a collection of nodes.
    pub fn aggregate(nodes: &[Node], function: &AggregateFunction) -> Result<AggregateResult> {
        match function {
            AggregateFunction::Count => Ok(AggregateResult::Number(nodes.len() as f64)),

            AggregateFunction::Sum(property) => {
                let sum = nodes
                    .iter()
                    .filter_map(|node| node.get_property(property))
                    .filter_map(|value| value.as_f64())
                    .sum();
                Ok(AggregateResult::Number(sum))
            }

            AggregateFunction::Avg(property) => {
                let values: Vec<f64> = nodes
                    .iter()
                    .filter_map(|node| node.get_property(property))
                    .filter_map(|value| value.as_f64())
                    .collect();

                if values.is_empty() {
                    Ok(AggregateResult::Number(0.0))
                } else {
                    let avg = values.iter().sum::<f64>() / values.len() as f64;
                    Ok(AggregateResult::Number(avg))
                }
            }

            AggregateFunction::Min(property) => {
                let min = nodes
                    .iter()
                    .filter_map(|node| node.get_property(property))
                    .filter_map(|value| value.as_f64())
                    .fold(f64::INFINITY, f64::min);

                if min.is_infinite() {
                    Ok(AggregateResult::Number(0.0))
                } else {
                    Ok(AggregateResult::Number(min))
                }
            }

            AggregateFunction::Max(property) => {
                let max = nodes
                    .iter()
                    .filter_map(|node| node.get_property(property))
                    .filter_map(|value| value.as_f64())
                    .fold(f64::NEG_INFINITY, f64::max);

                if max.is_infinite() {
                    Ok(AggregateResult::Number(0.0))
                } else {
                    Ok(AggregateResult::Number(max))
                }
            }

            AggregateFunction::Distinct(property) => {
                let mut unique_values = std::collections::HashSet::new();

                for node in nodes {
                    if let Some(value) = node.get_property(property) {
                        unique_values.insert(value.clone());
                    }
                }

                let list: Vec<Value> = unique_values.into_iter().collect();
                Ok(AggregateResult::List(list))
            }

            AggregateFunction::GroupBy(property) => {
                let mut groups = HashMap::new();

                for node in nodes {
                    let group_key = if let Some(value) = node.get_property(property) {
                        match value {
                            Value::String(s) => s.clone(),
                            Value::Number(n) => n.to_string(),
                            Value::Bool(b) => b.to_string(),
                            _ => "null".to_string(),
                        }
                    } else {
                        "null".to_string()
                    };

                    groups
                        .entry(group_key)
                        .or_insert_with(Vec::new)
                        .push(node.clone());
                }

                Ok(AggregateResult::Groups(groups))
            }
        }
    }

    /// Apply multiple aggregation functions to a collection of nodes.
    pub fn multi_aggregate(
        nodes: &[Node],
        functions: &[AggregateFunction],
    ) -> Result<Vec<AggregateResult>> {
        functions
            .iter()
            .map(|func| Self::aggregate(nodes, func))
            .collect()
    }
}

/// Statistical aggregations for numeric properties.
pub struct Statistics;

impl Statistics {
    /// Calculate basic statistics for a numeric property.
    pub fn calculate(nodes: &[Node], property: &str) -> Result<StatisticsResult> {
        let values: Vec<f64> = nodes
            .iter()
            .filter_map(|node| node.get_property(property))
            .filter_map(|value| value.as_f64())
            .collect();

        if values.is_empty() {
            return Ok(StatisticsResult::default());
        }

        let count = values.len() as f64;
        let sum = values.iter().sum::<f64>();
        let mean = sum / count;

        let min = values.iter().fold(f64::INFINITY, |a, &b| a.min(b));
        let max = values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));

        // Calculate variance and standard deviation
        let variance = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / count;
        let std_dev = variance.sqrt();

        // Calculate median
        let mut sorted_values = values.clone();
        sorted_values.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let median = if sorted_values.len().is_multiple_of(2) {
            let mid = sorted_values.len() / 2;
            (sorted_values[mid - 1] + sorted_values[mid]) / 2.0
        } else {
            sorted_values[sorted_values.len() / 2]
        };

        Ok(StatisticsResult {
            count,
            sum,
            mean,
            median,
            min,
            max,
            variance,
            std_dev,
        })
    }
}

/// Result of statistical calculations.
#[derive(Debug, Clone, PartialEq)]
pub struct StatisticsResult {
    /// Number of values in the dataset
    pub count: f64,
    /// Sum of all values
    pub sum: f64,
    /// Arithmetic mean (average)
    pub mean: f64,
    /// Middle value when sorted
    pub median: f64,
    /// Smallest value in the dataset
    pub min: f64,
    /// Largest value in the dataset
    pub max: f64,
    /// Population variance
    pub variance: f64,
    /// Standard deviation (square root of variance)
    pub std_dev: f64,
}

impl Default for StatisticsResult {
    fn default() -> Self {
        StatisticsResult {
            count: 0.0,
            sum: 0.0,
            mean: 0.0,
            median: 0.0,
            min: 0.0,
            max: 0.0,
            variance: 0.0,
            std_dev: 0.0,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn create_test_nodes() -> Vec<Node> {
        vec![
            Node::new(
                1,
                [
                    ("age".to_string(), json!(25)),
                    ("name".to_string(), json!("Alice")),
                    ("category".to_string(), json!("A")),
                ]
                .into(),
            ),
            Node::new(
                2,
                [
                    ("age".to_string(), json!(30)),
                    ("name".to_string(), json!("Bob")),
                    ("category".to_string(), json!("B")),
                ]
                .into(),
            ),
            Node::new(
                3,
                [
                    ("age".to_string(), json!(35)),
                    ("name".to_string(), json!("Charlie")),
                    ("category".to_string(), json!("A")),
                ]
                .into(),
            ),
        ]
    }

    #[test]
    fn test_count_aggregation() {
        let nodes = create_test_nodes();
        let result = Aggregator::aggregate(&nodes, &AggregateFunction::Count).unwrap();
        assert_eq!(result, AggregateResult::Number(3.0));
    }

    #[test]
    fn test_sum_aggregation() {
        let nodes = create_test_nodes();
        let result =
            Aggregator::aggregate(&nodes, &AggregateFunction::Sum("age".to_string())).unwrap();
        assert_eq!(result, AggregateResult::Number(90.0));
    }

    #[test]
    fn test_avg_aggregation() {
        let nodes = create_test_nodes();
        let result =
            Aggregator::aggregate(&nodes, &AggregateFunction::Avg("age".to_string())).unwrap();
        assert_eq!(result, AggregateResult::Number(30.0));
    }

    #[test]
    fn test_group_by_aggregation() {
        let nodes = create_test_nodes();
        let result =
            Aggregator::aggregate(&nodes, &AggregateFunction::GroupBy("category".to_string()))
                .unwrap();

        if let AggregateResult::Groups(groups) = result {
            assert_eq!(groups.len(), 2);
            assert_eq!(groups.get("A").unwrap().len(), 2);
            assert_eq!(groups.get("B").unwrap().len(), 1);
        } else {
            panic!("Expected Groups result");
        }
    }

    #[test]
    fn test_statistics() {
        let nodes = create_test_nodes();
        let stats = Statistics::calculate(&nodes, "age").unwrap();

        assert_eq!(stats.count, 3.0);
        assert_eq!(stats.sum, 90.0);
        assert_eq!(stats.mean, 30.0);
        assert_eq!(stats.median, 30.0);
        assert_eq!(stats.min, 25.0);
        assert_eq!(stats.max, 35.0);
    }
}