tellaro-query-language 3.0.1

//! Stats evaluator for TQL aggregation queries.
//!
//! Provides statistical aggregation functions with grouping support.

use crate::error::{Result, TqlError};
use crate::field_accessor;
use crate::parser::Aggregation;
use serde_json::{json, Value as JsonValue};
use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};

/// Build the `params` map for an [`AggregationSpec`] from a parsed [`Aggregation`].
///
/// Every construction site used to pass `HashMap::new()`, so
/// `params.get("percentile_values")` was unconditionally `None` and the
/// evaluator fell back to its `vec![50.0]` default. `percentile(n, 90)`,
/// `percentile(n, 50)` and `p(n, 10)` all returned the SAME number -- a
/// plausible one, silently answering a question nobody asked. The values were
/// parsed correctly the whole time; only the hand-off dropped them.
pub fn agg_params(agg: &Aggregation) -> HashMap<String, JsonValue> {
    let mut params = HashMap::new();
    if let Some(values) = &agg.percentile_values {
        params.insert("percentile_values".to_string(), json!(values));
    }
    if let Some(values) = &agg.rank_values {
        params.insert("rank_values".to_string(), json!(values));
    }
    params
}

/// Render one group-key part as text. Used only to ORDER non-scalar values in
/// [`sort_key`]; bucket IDENTITY is [`group_key_class`], which is a different
/// question and used to be conflated with this one.
fn render_group_key_part(value: &JsonValue) -> String {
    match value {
        JsonValue::String(s) => s.clone(),
        JsonValue::Number(n) => n.to_string(),
        JsonValue::Bool(b) => b.to_string(),
        JsonValue::Null => "null".to_string(),
        _ => serde_json::to_string(value).unwrap_or_default(),
    }
}

/// The identity two group-key values share: same class string, same bucket.
///
/// Bucket identity used to be `render_group_key_part`, i.e. equality of the
/// RENDERED text, and that is not the question. Rendering collapses values of
/// different types that happen to print alike and separates values of the same
/// type that do not:
///
/// ```text
///                     rendered text     Python `==` / `hash`     OpenSearch
///   1  vs  "1"        ONE bucket        two buckets              two buckets
///   1  vs  1.0        two buckets       ONE bucket               ONE bucket
///   1  vs  true       two buckets       ONE bucket               n/a
///   true vs "true"    ONE bucket        two buckets              two buckets
/// ```
///
/// Python groups with `_make_hashable_key`, which uses the value itself
/// whenever it is hashable — so its buckets are Python's own equality classes,
/// where `True == 1 == 1.0` and no number is ever equal to a string. That is
/// the rule reproduced here, and it is also the rule a `terms` aggregation
/// applies for the cases OpenSearch can express: a field's mapping decides how
/// a literal is read, so `1` and `1.0` cannot be distinct terms on a numeric
/// field while `"1"` on a keyword field is a different term entirely.
///
/// The regression was new on this branch: the commit that stopped EMITTING the
/// rendered key (so `| stats count() by n` reports `{"n": 1}` rather than
/// `{"n": "1"}`) left the rendering as the grouping key, where the type
/// information it had just been taught to preserve is exactly what gets thrown
/// away. Emitting and identity are separate concerns and now use separate
/// functions.
///
/// The one-character prefix is what keeps the classes disjoint: a string `"#1"`
/// classes as `s#1`, never as the number `1`.
fn group_key_class(value: &JsonValue) -> String {
    match value {
        // `True == 1` and `False == 0` in Python, and `hash` agrees, so a
        // boolean shares a class with the number it equals.
        JsonValue::Bool(b) => format!("#{}", if *b { 1 } else { 0 }),
        JsonValue::Number(n) => format!("#{}", canonical_number_class(n)),
        JsonValue::String(s) => format!("s{}", s),
        JsonValue::Null => "n".to_string(),
        other => format!("j{}", serde_json::to_string(other).unwrap_or_default()),
    }
}

/// One spelling per numeric VALUE, so `1`, `1.0` and `true` all class alike.
///
/// Integral values are written as integers whatever JSON type carried them.
/// `as_i64` / `as_u64` are tried before `as_f64` so a magnitude beyond f64's
/// exact-integer range is not rounded into a neighbour's class.
fn canonical_number_class(n: &serde_json::Number) -> String {
    if let Some(i) = n.as_i64() {
        return i.to_string();
    }
    if let Some(u) = n.as_u64() {
        return u.to_string();
    }
    match n.as_f64() {
        Some(f)
            if f.is_finite()
                && f.fract() == 0.0
                && (i64::MIN as f64..=i64::MAX as f64).contains(&f) =>
        {
            (f as i64).to_string()
        }
        Some(f) => format!("{:?}", f),
        None => n.to_string(),
    }
}

/// Ordering key for the listing family (`values` / `unique` / `distinct`).
///
/// `(type_rank, numeric, text)`. Numbers and booleans share rank 0 and order by
/// their numeric value, so `[1, 2, 10]` stays `[1, 2, 10]` rather than becoming
/// `[1, 10, 2]`; strings take rank 1 and order lexicographically. That agrees
/// with Python's `sorted` for every homogeneous list, which is the only case
/// Python can express -- `sorted([1, "a"])` raises `TypeError`.
fn sort_key(value: &JsonValue) -> (u8, f64, String) {
    match value {
        JsonValue::Number(n) => (0, n.as_f64().unwrap_or(0.0), String::new()),
        JsonValue::Bool(b) => (0, if *b { 1.0 } else { 0.0 }, String::new()),
        JsonValue::String(s) => (1, 0.0, s.clone()),
        other => (2, 0.0, render_group_key_part(other)),
    }
}

/// The buckets one record contributes to, given its group-by field values.
///
/// This is OpenSearch's `terms` aggregation contract, which is the third engine
/// that settles what `| stats ... by <field>` must mean. Verified against a live
/// OpenSearch 2.19.4 cluster over the documents
/// `[{a:"x"},{a:"x"},{a:"y"},{a:null},{},{a:"z"}]` and `[{s:["p","q"]},{s:"p"}]`:
///
/// * **A missing or null value produces NO bucket.** `terms` on `a` returned
///   exactly `x`(2), `y`(1), `z`(1) -- the null-valued and absent documents are
///   not counted anywhere. Rust invented a bucket keyed with the literal STRING
///   `"null"`, which is indistinguishable from a genuine value `"null"`, and
///   Python invented one keyed with JSON `null`. Both were wrong, in different
///   ways, which is why neither engine's disagreement with the other pointed at
///   the answer. A phantom group is worse than a missing one: an analyst reads
///   it as data, and a threshold rule counts it.
/// * **An array produces one bucket per distinct element.** `terms` on `s`
///   returned `p`(2) -- both documents -- and `q`(1), because OpenSearch indexes
///   each element as its own term. Rust collapsed the array into a single bucket
///   keyed with the JSON text `"[\"p\",\"q\"]"`, a group no document could ever
///   be said to be in. Deduplicated within one record so a repeated element
///   cannot count that record twice.
///
/// Several group-by fields give the cross product, which is what nested `terms`
/// aggregations produce.
fn group_key_combinations(field_values: &[JsonValue]) -> Vec<Vec<JsonValue>> {
    let mut combinations: Vec<Vec<JsonValue>> = vec![Vec::new()];

    for value in field_values {
        let alternatives: Vec<JsonValue> = match value {
            JsonValue::Null => Vec::new(),
            JsonValue::Array(items) => {
                let mut seen = HashSet::new();
                items
                    .iter()
                    .filter(|item| !item.is_null())
                    .filter(|item| seen.insert(group_key_class(item)))
                    .cloned()
                    .collect()
            }
            other => vec![other.clone()],
        };

        if alternatives.is_empty() {
            // This record has no value for this field, so it belongs to no
            // bucket at all -- not to a bucket named "null".
            return Vec::new();
        }

        combinations = combinations
            .iter()
            .flat_map(|prefix| {
                alternatives.iter().map(move |alt| {
                    let mut extended = prefix.clone();
                    extended.push(alt.clone());
                    extended
                })
            })
            .collect();
    }

    combinations
}

/// Stats evaluator for aggregation queries
pub struct StatsEvaluator {
    _max_depth: usize,
}

/// Aggregation specification
///
/// `modifier` / `limit` carry the AGGREGATION-level top-N, which both engines
/// accept in two spellings -- `sum(x) top 3 by y` and `sum(x, top 3) by y` --
/// and which the parser folds onto the same pair of fields. It is a DIFFERENT
/// modifier from [`GroupBySpec::bucket_size`]: this one ranks buckets by the
/// AGGREGATE VALUE, that one by `doc_count`. See [`StatsEvaluator::apply_modifiers`].
#[derive(Debug, Clone, Default)]
pub struct AggregationSpec {
    pub function: String,
    pub field: String,
    pub alias: Option<String>,
    /// `"top"` or `"bottom"`, when the query carried an aggregation-level top-N.
    pub modifier: Option<String>,
    /// How many buckets `modifier` keeps. Absent with a modifier present means
    /// 10, which is Python's `agg.get("limit", 10)`.
    pub limit: Option<usize>,
    pub params: HashMap<String, JsonValue>,
}

/// One group-by field, with the per-field bucket limit `by <field> top N` sets.
///
/// This type exists because `StatsQuery.group_by` was a `Vec<String>`: the
/// parser read `bucket_size` onto the `GroupBy` AST node and BOTH lowering
/// sites -- `lib.rs` and `file_ops.rs` -- dropped it on the way in, exactly as
/// they dropped `params` before `agg_params` was introduced (see the note on
/// [`agg_params`], and `stats_agg_name_reachability_tests.rs`). A limit that is
/// parsed, carried and then discarded returns MORE buckets than the user asked
/// for, which reads as a complete answer rather than as an error.
///
/// Both lowering sites now go through [`From`] impls in this module so a new
/// field on the AST node cannot be silently dropped at one site and not the other.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct GroupBySpec {
    pub field: String,
    /// `N` from `by <field> top N`. Ranks by `doc_count`, not by any aggregate.
    pub bucket_size: Option<usize>,
}

impl From<&str> for GroupBySpec {
    fn from(field: &str) -> Self {
        Self {
            field: field.to_string(),
            bucket_size: None,
        }
    }
}

impl From<String> for GroupBySpec {
    fn from(field: String) -> Self {
        Self {
            field,
            bucket_size: None,
        }
    }
}

impl From<&crate::parser::GroupBy> for GroupBySpec {
    fn from(group_by: &crate::parser::GroupBy) -> Self {
        Self {
            field: group_by.field.clone(),
            bucket_size: group_by.bucket_size,
        }
    }
}

impl From<&Aggregation> for AggregationSpec {
    fn from(agg: &Aggregation) -> Self {
        Self {
            function: agg.function.clone(),
            // `count()` and `count(*)` both reach the evaluator as field `*`.
            field: agg.field.clone().unwrap_or_else(|| "*".to_string()),
            alias: agg.alias.clone(),
            modifier: agg.modifier.clone(),
            limit: agg.limit,
            params: agg_params(agg),
        }
    }
}

/// Distinct group-key values already reserved at one grouping level, keyed by
/// the rendered values of the levels ABOVE it. `None` is a group-by field the
/// bucket's key does not carry, which is a distinct member from any string.
type ReservedValues = HashMap<Vec<Option<String>>, HashSet<Option<String>>>;

/// Stats query specification
#[derive(Debug, Clone)]
pub struct StatsQuery {
    pub aggregations: Vec<AggregationSpec>,
    pub group_by: Vec<GroupBySpec>,
}

impl Default for StatsEvaluator {
    fn default() -> Self {
        Self::new()
    }
}

impl StatsEvaluator {
    /// Create a new stats evaluator
    pub fn new() -> Self {
        Self { _max_depth: 100 }
    }

    /// Evaluate stats query against records
    ///
    /// Returns aggregated results in a structured format
    pub fn evaluate_stats(&self, records: &[JsonValue], query: &StatsQuery) -> Result<JsonValue> {
        if query.group_by.is_empty() {
            self.simple_aggregation(records, &query.aggregations)
        } else {
            self.grouped_aggregation(records, &query.aggregations, &query.group_by)
        }
    }

    /// Perform aggregation without grouping
    fn simple_aggregation(
        &self,
        records: &[JsonValue],
        aggregations: &[AggregationSpec],
    ) -> Result<JsonValue> {
        if aggregations.len() == 1 {
            // Single aggregation
            let agg = &aggregations[0];
            let value = self.calculate_aggregation(records, agg)?;

            Ok(json!({
                "type": "simple_aggregation",
                "function": agg.function,
                "field": agg.field,
                "alias": agg.alias,
                "value": value
            }))
        } else {
            // Multiple aggregations
            let mut results = HashMap::new();
            for agg in aggregations {
                let value = self.calculate_aggregation(records, agg)?;
                let key = agg
                    .alias
                    .clone()
                    .unwrap_or_else(|| format!("{}_{}", agg.function, agg.field));
                results.insert(key, value);
            }

            Ok(json!({
                "type": "multiple_aggregations",
                "results": results
            }))
        }
    }

    /// Perform aggregation with grouping
    fn grouped_aggregation(
        &self,
        records: &[JsonValue],
        aggregations: &[AggregationSpec],
        group_by: &[GroupBySpec],
    ) -> Result<JsonValue> {
        let group_by_fields: Vec<String> = group_by.iter().map(|g| g.field.clone()).collect();

        // Group records by field values, emitted in FIRST-APPEARANCE order --
        // the order of the record stream, which is what Python's `defaultdict`
        // produces and therefore the reference ordering.
        //
        // This was a `BTreeMap`, which ordered buckets lexicographically by the
        // rendered key. Determinism was the stated reason and is preserved here
        // (`records` is an ordered slice), but lexicographic order is not
        // Python's order, and that difference is not cosmetic once a top-N
        // modifier exists: both top-N passes below sort STABLY, so the
        // pre-sort order is exactly what breaks ties. Over the shared
        // `user_records.json` corpus, `| stats sum(salary) by department top 5`
        // cuts inside a four-way `doc_count = 2` tie -- Python's fifth bucket is
        // `Finance` (first appearance), a lexicographic pre-sort's would be
        // `Data Science`. Same count, different answer.
        //
        // The map is keyed by a STRING rendering of the group key purely so it can be
        // ordered and hashed; the JSON value each part came from is carried alongside
        // and is what the emitted `key` reports. Rendering the key and then emitting
        // the rendering is what made `| stats count() by n` over numeric `n` report
        // `{"n": "1"}` where Python reports `{"n": 1}` -- a TYPE divergence in the one
        // field a consumer keys its buckets on.
        let mut order: Vec<Vec<String>> = Vec::new();
        let mut groups: HashMap<Vec<String>, Vec<&JsonValue>> = HashMap::new();
        let mut key_values: HashMap<Vec<String>, Vec<JsonValue>> = HashMap::new();

        for record in records {
            let field_values: Vec<JsonValue> = group_by_fields
                .iter()
                .map(|field| match field_accessor::get_field(record, field) {
                    Ok(Some(v)) => v.clone(),
                    Ok(None) | Err(_) => JsonValue::Null,
                })
                .collect();

            for key_parts in group_key_combinations(&field_values) {
                let rendered: Vec<String> = key_parts.iter().map(group_key_class).collect();
                if !groups.contains_key(&rendered) {
                    order.push(rendered.clone());
                    key_values.insert(rendered.clone(), key_parts);
                }
                groups.entry(rendered).or_default().push(record);
            }
        }

        // Calculate aggregations for each group
        let mut results = Vec::new();
        for rendered_key in &order {
            let group_records = &groups[rendered_key];
            let mut group_result: HashMap<String, JsonValue> = HashMap::new();

            // Add group key, as the JSON values it was built from rather than as
            // the string rendering used to order the map.
            let key_parts = key_values
                .get(rendered_key)
                .expect("every group has a recorded key");
            let mut key_map = HashMap::new();
            for (i, field) in group_by_fields.iter().enumerate() {
                key_map.insert(field.clone(), key_parts[i].clone());
            }
            group_result.insert("key".to_string(), json!(key_map));
            group_result.insert("doc_count".to_string(), json!(group_records.len()));

            if aggregations.len() == 1 {
                // Single aggregation
                let agg = &aggregations[0];
                let owned_records: Vec<JsonValue> =
                    group_records.iter().map(|&r| r.clone()).collect();
                let value = self.calculate_aggregation(&owned_records, agg)?;
                let agg_key = agg.alias.clone().unwrap_or_else(|| agg.function.clone());
                group_result.insert(agg_key, value);
            } else {
                // Multiple aggregations
                let mut agg_results = HashMap::new();
                for agg in aggregations {
                    let owned_records: Vec<JsonValue> =
                        group_records.iter().map(|&r| r.clone()).collect();
                    let value = self.calculate_aggregation(&owned_records, agg)?;
                    let agg_key = agg
                        .alias
                        .clone()
                        .unwrap_or_else(|| format!("{}_{}", agg.function, agg.field));
                    agg_results.insert(agg_key, value);
                }
                group_result.insert("aggregations".to_string(), json!(agg_results));
            }

            results.push(json!(group_result));
        }

        // Order matters here: Python applies the aggregation-level modifier
        // FIRST and the group-by bucket limit SECOND, so a query carrying both
        // (`sum(x) top 3 by y top 5`) narrows by aggregate value and then
        // RE-ORDERS the survivors by `doc_count`. Swapping the two passes gives
        // a different set, not just a different order.
        let results = Self::apply_modifiers(results, aggregations)?;
        let results = Self::apply_bucket_limits(results, group_by);

        Ok(json!({
            "type": "grouped_aggregation",
            "group_by": group_by_fields,
            "results": results
        }))
    }

    /// The map a modifier ranks in: the nested `aggregations` when the bucket
    /// carries several, the group result itself when it carries one.
    ///
    /// Python does NOT fall back from the nested map to the top level, so this
    /// is a `match` on the presence of `"aggregations"` rather than a chained
    /// `or_else` -- a fallback would hide exactly the naming mismatch below.
    fn modifier_container(result: &JsonValue) -> &JsonValue {
        match result.get("aggregations") {
            Some(aggs) => aggs,
            None => result,
        }
    }

    /// The value a modifier ranks a bucket by, or an error naming the miss.
    ///
    /// An ABSENT key is a wiring gap between the emitter and this sort, not
    /// missing data, and the two must not look alike. Both engines defaulted it
    /// to `0`, which made every bucket compare equal and turned a stable sort's
    /// "top N" into "the first N in record order" -- no error, no empty result,
    /// just a different answer. That is the absence-read-as-a-benign-value
    /// shape, and the fix is to refuse rather than to guess.
    ///
    /// A present JSON `null` is NOT that case: it is what `avg`/`min`/`max`
    /// return for a group with no numeric values, and it ranks lowest. That is a
    /// deliberate divergence from Python, which raises `TypeError: '<' not
    /// supported between instances of 'NoneType' and 'NoneType'` (measured on
    /// `| stats avg(bogus_field) top 3 by department`) -- there is no Python
    /// semantics to match there, only a crash, and a deterministic answer beats
    /// an exception.
    fn modifier_value<'a>(result: &'a JsonValue, agg_key: &str) -> Result<&'a JsonValue> {
        let container = Self::modifier_container(result);
        container.get(agg_key).ok_or_else(|| {
            let found: Vec<&str> = match container.as_object() {
                Some(map) => {
                    let mut keys: Vec<&str> = map.keys().map(String::as_str).collect();
                    keys.sort_unstable();
                    keys
                }
                None => Vec::new(),
            };
            TqlError::ExecutionError(format!(
                "Cannot rank buckets by '{}': no aggregation under that name in the \
                 group result (found {:?}). This is a naming mismatch between the \
                 aggregation emitter and this sort, not missing data. Give the \
                 aggregation an explicit alias with `as <name>`.",
                agg_key, found
            ))
        })
    }

    /// Total order over modifier sort values. See `modifier_value` for why a
    /// `null` ranks lowest and an absent key never reaches here.
    fn modifier_sort_key(value: &JsonValue) -> (u8, f64, String) {
        match value {
            JsonValue::Null => (0, 0.0, String::new()),
            other => sort_key(other),
        }
    }

    /// Apply the AGGREGATION-level top-N (`sum(x) top 3 by y`, `sum(x, top 3) by y`).
    ///
    /// # The key rule
    ///
    /// Ranks buckets by the aggregate's OWN emitted value, and the name that
    /// value was emitted under is not one expression. `grouped_aggregation`
    /// above names it:
    ///
    /// ```text
    ///   alias                   if the query supplied one (`... as total`)
    ///   function                if the query has exactly ONE aggregation --
    ///                           written at the TOP LEVEL of the group result
    ///   "{function}_{field}"    if it has TWO OR MORE -- written inside the
    ///                           group result's `aggregations` map
    /// ```
    ///
    /// So the branch is on `aggregations.len()`, not on the individual
    /// aggregation and not on whether a given result happens to carry an
    /// `aggregations` key. This computed `alias or function` unconditionally,
    /// so with two or more aggregations it asked for `sum` where the emitter had
    /// written `sum_salary`, the lookup missed for every bucket, and the top-N
    /// became a no-op:
    ///
    /// ```text
    ///   | stats count(), sum(salary) top 2 by dept   in-memory: record order
    ///   the same query pushed to OpenSearch          order: {sum_salary_1: desc}
    /// ```
    ///
    /// One saved query, two answers, decided by where it ran -- and silently,
    /// because a missed lookup defaulted to `0` and a stable sort left the
    /// buckets where they were. One aggregation alone was correct all along,
    /// which is why this survived. `modifier_value` now refuses an absent key
    /// rather than defaulting it.
    ///
    /// The OpenSearch sub-aggregation alias (`sum_salary_1`, index-suffixed) is
    /// a THIRD naming and deliberately untouched: it is internal to the DSL, and
    /// the two engines only have to agree on the ORDER they return.
    ///
    /// Only the FIRST aggregation carrying a modifier is applied -- Python
    /// `break`s out of the loop -- so a second modifier on a later aggregation is
    /// ignored rather than composed. `limit` defaults to 10.
    ///
    /// The sort is STABLE, and Python's `sorted(..., reverse=True)` is stable
    /// too (it does not reverse equal elements), so ties in both engines are
    /// broken by first-appearance order. That is why `grouped_aggregation`
    /// emits groups in record order.
    ///
    /// Never applied without a `by` clause: Python calls this only from its
    /// grouped paths, so `| stats sum(salary) top 3` is the plain sum in both
    /// engines.
    fn apply_modifiers(
        results: Vec<JsonValue>,
        aggregations: &[AggregationSpec],
    ) -> Result<Vec<JsonValue>> {
        for agg in aggregations {
            let Some(modifier) = agg.modifier.as_deref() else {
                continue;
            };
            let agg_key = agg.alias.clone().unwrap_or_else(|| {
                if aggregations.len() == 1 {
                    agg.function.clone()
                } else {
                    format!("{}_{}", agg.function, agg.field)
                }
            });
            let descending = modifier == "top";

            // Resolve every bucket's rank value BEFORE sorting, so a missing key
            // is reported rather than swallowed by a comparator that cannot fail.
            let mut ranked: Vec<((u8, f64, String), JsonValue)> = Vec::with_capacity(results.len());
            for result in results {
                let key = Self::modifier_sort_key(Self::modifier_value(&result, &agg_key)?);
                ranked.push((key, result));
            }

            ranked.sort_by(|(ka, _), (kb, _)| {
                let ordering = ka.partial_cmp(kb).unwrap_or(Ordering::Equal);
                if descending {
                    ordering.reverse()
                } else {
                    ordering
                }
            });

            ranked.truncate(agg.limit.unwrap_or(10));
            return Ok(ranked.into_iter().map(|(_, result)| result).collect());
        }
        Ok(results)
    }

    /// Apply the GROUP-BY bucket limit (`by <field> top N`).
    ///
    /// Ranks by `doc_count` descending -- NOT by any aggregate, which is what
    /// separates this modifier from `apply_modifiers` above. With one group-by
    /// field it is a sort and a truncate; with several it is a greedy pass that
    /// reserves at most `bucket_size` distinct values per (level, parent-key)
    /// pair, which is what nested `terms` aggregations with per-level `size`
    /// produce.
    ///
    /// `top 0` on a group-by field is a no-op for single-level grouping and
    /// empties the result for multi-level grouping. That asymmetry is Python's
    /// (`if bucket_size:` is falsy at 0, `if bucket_size is not None:` is not)
    /// and is reproduced rather than corrected, because the two engines
    /// disagreeing is worse than one odd answer. It is pinned by a test.
    fn apply_bucket_limits(
        mut results: Vec<JsonValue>,
        group_by: &[GroupBySpec],
    ) -> Vec<JsonValue> {
        if !group_by.iter().any(|g| g.bucket_size.is_some()) {
            return results;
        }

        let doc_count = |r: &JsonValue| r.get("doc_count").and_then(|v| v.as_u64()).unwrap_or(0);

        if group_by.len() == 1 {
            // The `> 0` guard covers the SORT as well as the truncate. Python
            // writes `if bucket_size:`, which is falsy at 0, so `top 0` on a
            // single group-by field leaves the buckets in their original order
            // -- not merely unlimited. Guarding only the truncate re-sorted
            // them by `doc_count` and diverged on ordering while agreeing on count.
            if let Some(bucket_size) = group_by[0].bucket_size {
                if bucket_size > 0 {
                    results.sort_by_key(|r| std::cmp::Reverse(doc_count(r)));
                    results.truncate(bucket_size);
                }
            }
            return results;
        }

        // Multi-level: Python sorts unconditionally once ANY field carries a
        // limit, including when the field carrying it is `top 0`.
        results.sort_by_key(|r| std::cmp::Reverse(doc_count(r)));

        // Hierarchical: reserve at most `bucket_size` distinct values per level,
        // scoped to the values already chosen at the levels above.
        let mut level_values: Vec<ReservedValues> = vec![HashMap::new(); group_by.len()];
        let mut filtered = Vec::new();

        for result in results {
            let mut should_include = true;
            let mut key_path: Vec<Option<String>> = Vec::new();

            for (level, spec) in group_by.iter().enumerate() {
                let field_value = result
                    .get("key")
                    .and_then(|k| k.get(&spec.field))
                    .map(group_key_class);
                key_path.push(field_value.clone());

                let Some(bucket_size) = spec.bucket_size else {
                    continue;
                };
                let parent_key = key_path[..level].to_vec();
                let reserved = level_values[level].entry(parent_key).or_default();
                if !reserved.contains(&field_value) {
                    if reserved.len() >= bucket_size {
                        // Python `break`s here, KEEPING the reservations this
                        // result already made at shallower levels.
                        should_include = false;
                        break;
                    }
                    reserved.insert(field_value);
                }
            }

            if should_include {
                filtered.push(result);
            }
        }

        filtered
    }

    /// Calculate a single aggregation value
    fn calculate_aggregation(
        &self,
        records: &[JsonValue],
        agg_spec: &AggregationSpec,
    ) -> Result<JsonValue> {
        let func = &agg_spec.function;
        let field = &agg_spec.field;

        // Handle count(*)
        if func == "count" && field == "*" {
            return Ok(json!(records.len()));
        }

        // Extract field values
        let values: Vec<JsonValue> = records
            .iter()
            .filter_map(|record| match field_accessor::get_field(record, field) {
                Ok(Some(value)) => {
                    if value != &JsonValue::Null {
                        Some(value.clone())
                    } else {
                        None
                    }
                }
                _ => None,
            })
            .collect();

        // Calculate aggregation based on function
        match func.to_lowercase().as_str() {
            "count" => Ok(json!(values.len())),

            "unique_count" | "cardinality" => {
                let unique: HashSet<String> = values
                    .iter()
                    .map(|v| serde_json::to_string(v).unwrap_or_default())
                    .collect();
                Ok(json!(unique.len()))
            }

            "sum" => {
                let sum: f64 = self.numeric_values(&values, field)?.iter().sum();
                // `0.0`, never `-0.0`: an empty population sums to positive
                // zero in OpenSearch and in Python, and `-0.0` serialises as
                // `-0.0` in the JSON a consumer reads.
                Ok(json!(if sum == 0.0 { 0.0 } else { sum }))
            }

            "min" => {
                let min = self
                    .numeric_values(&values, field)?
                    .into_iter()
                    .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                Ok(json!(min))
            }

            "max" => {
                let max = self
                    .numeric_values(&values, field)?
                    .into_iter()
                    .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                Ok(json!(max))
            }

            "average" | "avg" | "mean" => {
                let numeric_values: Vec<f64> = self.numeric_values(&values, field)?;

                if numeric_values.is_empty() {
                    Ok(JsonValue::Null)
                } else {
                    let avg = numeric_values.iter().sum::<f64>() / numeric_values.len() as f64;
                    Ok(json!(avg))
                }
            }

            "median" | "med" => {
                let mut numeric_values: Vec<f64> = self.numeric_values(&values, field)?;

                if numeric_values.is_empty() {
                    return Ok(JsonValue::Null);
                }

                numeric_values
                    .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                let len = numeric_values.len();

                let median = if len.is_multiple_of(2) {
                    (numeric_values[len / 2 - 1] + numeric_values[len / 2]) / 2.0
                } else {
                    numeric_values[len / 2]
                };

                Ok(json!(median))
            }

            "std" | "standard_deviation" => {
                let numeric_values: Vec<f64> = self.numeric_values(&values, field)?;

                if numeric_values.len() < 2 {
                    return Ok(JsonValue::Null);
                }

                let mean = numeric_values.iter().sum::<f64>() / numeric_values.len() as f64;
                let variance = numeric_values
                    .iter()
                    .map(|v| (v - mean).powi(2))
                    .sum::<f64>()
                    / (numeric_values.len() - 1) as f64;
                let std_dev = variance.sqrt();

                Ok(json!(std_dev))
            }

            "percentile" | "percentiles" | "p" | "pct" => {
                let mut numeric_values: Vec<f64> = self.numeric_values(&values, field)?;

                if numeric_values.is_empty() {
                    return Ok(JsonValue::Null);
                }

                numeric_values
                    .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

                // Get percentile values from params (default to 50 for median)
                let percentile_values = agg_spec
                    .params
                    .get("percentile_values")
                    .and_then(|v| v.as_array())
                    .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect::<Vec<f64>>())
                    .unwrap_or_else(|| vec![50.0]);

                if percentile_values.len() == 1 {
                    // Single percentile
                    let p = self.calculate_percentile(&numeric_values, percentile_values[0])?;
                    Ok(json!(p))
                } else {
                    // Multiple percentiles
                    let mut result = HashMap::new();
                    for p_val in percentile_values {
                        let p = self.calculate_percentile(&numeric_values, p_val)?;
                        result.insert(format!("p{}", p_val as i32), json!(p));
                    }
                    Ok(json!(result))
                }
            }

            "percentile_rank" | "percentile_ranks" | "pct_rank" | "pct_ranks" => {
                let mut numeric_values: Vec<f64> = self.numeric_values(&values, field)?;

                if numeric_values.is_empty() {
                    return Ok(JsonValue::Null);
                }

                numeric_values
                    .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

                let rank_values = agg_spec
                    .params
                    .get("rank_values")
                    .and_then(|v| v.as_array())
                    .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect::<Vec<f64>>())
                    .unwrap_or_default();

                if rank_values.is_empty() {
                    return Err(TqlError::ExecutionError(
                        "percentile_rank requires at least one value".to_string(),
                    ));
                }

                if rank_values.len() == 1 {
                    Ok(json!(self.calculate_percentile_rank(
                        &numeric_values,
                        rank_values[0]
                    )))
                } else {
                    let mut result = HashMap::new();
                    for v in rank_values {
                        result.insert(
                            format!("rank_{:?}", v),
                            json!(self.calculate_percentile_rank(&numeric_values, v)),
                        );
                    }
                    Ok(json!(result))
                }
            }

            "values" | "unique" | "distinct" => {
                // The distinct values, KEEPING THEIR JSON TYPE.
                //
                // This rendered every value to a string first, so `distinct(n)`
                // over `[1, 2, 3]` returned `["1", "2", "3"]` where Python
                // returns `[1, 2, 3]`. OpenSearch settles it: a `terms`
                // aggregation -- which is what these three names are mapped to
                // in `AGGREGATION_MAPPING`, so it is the same operation
                // executed elsewhere -- returns `{"key": 1}` on a `long` field,
                // not `{"key": "1"}` (measured on a live 2.19.4 cluster). A
                // consumer comparing a listed value against a numeric threshold
                // got a string from this engine and a number from the other two.
                //
                // Rendering also decided the ORDER, lexicographically, so
                // `[1, 2, 10]` came back as `["1", "10", "2"]`. Numbers now sort
                // as numbers, which is both what Python's `sorted` does and what
                // a reader expects of a numeric field.
                let mut unique: Vec<JsonValue> = Vec::new();
                let mut seen = HashSet::new();
                for value in &values {
                    if seen.insert(group_key_class(value)) {
                        unique.push(value.clone());
                    }
                }

                // Sort numbers numerically and strings lexicographically, which
                // agrees with Python for any homogeneous list. A MIXED list is
                // ordered by type first -- Python's `sorted` raises TypeError
                // there, so there is no behaviour to match, and a deterministic
                // order beats an arbitrary one.
                unique.sort_by(|a, b| {
                    sort_key(a)
                        .partial_cmp(&sort_key(b))
                        .unwrap_or(Ordering::Equal)
                });
                Ok(json!(unique))
            }

            _ => Err(TqlError::ExecutionError(format!(
                "Unsupported aggregation function: {}",
                func
            ))),
        }
    }

    /// Every value of a numeric aggregate, or an error naming the first value
    /// that is not one.
    ///
    /// # Why this refuses rather than skipping
    ///
    /// These arms were written as `filter_map(to_numeric)`, which silently
    /// dropped anything that would not convert. Over
    /// `n = [1, 2, 3, null, "notnum"]`, `sum(n)` returned `6.0` and `avg(n)`
    /// returned `2.0` -- correct arithmetic over a population the analyst never
    /// chose, with nothing to say the population had been reduced. Python
    /// raised on the same data, so the two engines answered one query with a
    /// number and an error.
    ///
    /// This is a contract decision, not an obvious bug, so the reasoning is
    /// recorded rather than assumed:
    ///
    /// * **OpenSearch cannot arbitrate directly** -- it rejects `"notnum"` at
    ///   index time, so a cluster never sees this case. But that IS the
    ///   arbitration, one level up: the platform's contract is that a field you
    ///   aggregate numerically is MAPPED numeric. A string arriving here is a
    ///   schema violation, not dirty data to be tolerated, and it means the
    ///   in-memory and cluster answers to one query are computed over different
    ///   populations.
    /// * **The failure modes are not symmetric.** Refusing costs one visible
    ///   error, which an analyst fixes with a filter, a mutator, or a mapping
    ///   change. Skipping costs a threshold rule that compares against a
    ///   plausible wrong number indefinitely, with no signal that anything
    ///   happened -- which is this repository's dominant defect shape.
    /// * **Python already refused**, so converging on refusal removes the
    ///   divergence without weakening a guard that already existed.
    ///
    /// `null` and absent values are NOT an error and never reach here: they are
    /// dropped in `calculate_aggregation` before this is called, because
    /// OpenSearch's metric aggregations skip documents that have no value. Only
    /// a value that is PRESENT and not numeric is refused.
    fn numeric_values(&self, values: &[JsonValue], field: &str) -> Result<Vec<f64>> {
        values
            .iter()
            .map(|value| {
                self.to_numeric(value).ok_or_else(|| {
                    TqlError::ExecutionError(format!(
                        "Cannot convert {} to numeric value for aggregation on field '{}'. Ensure the field contains numeric data.",
                        serde_json::to_string(value).unwrap_or_else(|_| "value".to_string()),
                        field
                    ))
                })
            })
            .collect()
    }

    /// Convert value to numeric
    fn to_numeric(&self, value: &JsonValue) -> Option<f64> {
        match value {
            JsonValue::Number(n) => n.as_f64(),
            JsonValue::String(s) => s.parse::<f64>().ok(),
            JsonValue::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
            _ => None,
        }
    }

    /// Calculate a percentile by LINEAR INTERPOLATION between the two
    /// neighbouring ranks.
    ///
    /// This was nearest-rank, which is a different statistic, not a rounding
    /// difference: over `1..=10`, nearest-rank puts p50 at `6.0` while Python
    /// (`stats_evaluator.py::_calculate_percentile`) and OpenSearch's
    /// `percentiles` aggregation both report `5.5`. The engines disagreed on
    /// the VALUE for every dataset with an even element count, and Rust's own
    /// `median` arm -- which averages the two middle elements -- disagreed with
    /// its own `percentile(n, 50)`.
    fn calculate_percentile(&self, sorted_values: &[f64], percentile: f64) -> Result<Option<f64>> {
        if sorted_values.is_empty() {
            return Ok(None);
        }

        if !(0.0..=100.0).contains(&percentile) {
            return Err(TqlError::ExecutionError(format!(
                "Percentile must be between 0 and 100, got {}",
                percentile
            )));
        }

        let n = sorted_values.len();
        if n == 1 {
            return Ok(Some(sorted_values[0]));
        }

        let pos = (n as f64 - 1.0) * (percentile / 100.0);
        let lower_idx = pos.floor() as usize;
        let upper_idx = (lower_idx + 1).min(n - 1);

        if lower_idx == upper_idx {
            return Ok(Some(sorted_values[lower_idx]));
        }

        let lower = sorted_values[lower_idx];
        let upper = sorted_values[upper_idx];
        let fraction = pos - lower_idx as f64;

        Ok(Some(lower + fraction * (upper - lower)))
    }

    /// Calculate the percentile rank of `value` within `sorted_values`.
    ///
    /// Mirrors `stats_evaluator.py::_calculate_percentile_rank`, including the
    /// midpoint treatment of ties and the 2-dp rounding, so the two engines
    /// return the same number rather than merely the same shape.
    fn calculate_percentile_rank(&self, sorted_values: &[f64], value: f64) -> Option<f64> {
        if sorted_values.is_empty() {
            return None;
        }

        let n = sorted_values.len() as f64;
        let count_less = sorted_values.iter().filter(|v| **v < value).count() as f64;
        let count_equal = sorted_values.iter().filter(|v| **v == value).count() as f64;

        let rank = if count_equal > 0.0 {
            (count_less + count_equal / 2.0) / n * 100.0
        } else {
            count_less / n * 100.0
        };

        Some((rank * 100.0).round() / 100.0)
    }
}

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

    /// A rank key the emitter never wrote is REFUSED, not defaulted.
    ///
    /// Both engines used to answer `0` for a missed lookup. That made a wiring
    /// gap between the emitter and this sort indistinguishable from a group that
    /// genuinely aggregated to zero -- and since every bucket then compared
    /// equal, a stable sort turned `top N` into "the first N in record order".
    /// The whole B3 defect was invisible for exactly that reason, so the
    /// default is the thing under test here.
    ///
    /// This cannot be reached through a query any more: emitter and reader now
    /// derive the same name. It is a guard against them drifting apart again,
    /// which is why it calls `apply_modifiers` directly.
    #[test]
    fn a_rank_key_the_emitter_never_wrote_is_refused() {
        let results = vec![
            json!({"key": {"dept": "a"}, "doc_count": 1, "aggregations": {"sum_salary": 100}}),
            json!({"key": {"dept": "b"}, "doc_count": 1, "aggregations": {"sum_salary": 200}}),
        ];
        let aggregations = vec![
            AggregationSpec {
                function: "count".to_string(),
                field: "*".to_string(),
                ..Default::default()
            },
            AggregationSpec {
                function: "sum".to_string(),
                field: "salary".to_string(),
                // The name the emitter would have written is `sum_salary`.
                alias: Some("never_emitted".to_string()),
                modifier: Some("top".to_string()),
                limit: Some(1),
                ..Default::default()
            },
        ];

        let err = StatsEvaluator::apply_modifiers(results, &aggregations)
            .expect_err("a missing rank key must be reported, not defaulted to 0");
        let message = err.to_string();
        assert!(
            message.contains("never_emitted"),
            "the error must name the key that missed: {}",
            message
        );
        assert!(
            message.contains("sum_salary"),
            "the error must name what WAS present, so the mismatch is readable: {}",
            message
        );
    }

    /// A present JSON `null` is not a missing key: it is what `avg`/`min`/`max`
    /// return for a group with no numeric values, and it ranks lowest rather
    /// than aborting the query.
    #[test]
    fn a_null_aggregate_ranks_lowest_rather_than_refusing() {
        let results = vec![
            json!({"key": {"dept": "a"}, "doc_count": 1, "aggregations": {"avg_x": null, "count_*": 1}}),
            json!({"key": {"dept": "b"}, "doc_count": 1, "aggregations": {"avg_x": 5.0, "count_*": 1}}),
        ];
        let aggregations = vec![
            AggregationSpec {
                function: "count".to_string(),
                field: "*".to_string(),
                ..Default::default()
            },
            AggregationSpec {
                function: "avg".to_string(),
                field: "x".to_string(),
                modifier: Some("top".to_string()),
                limit: Some(1),
                ..Default::default()
            },
        ];

        let ranked = StatsEvaluator::apply_modifiers(results, &aggregations)
            .expect("a null aggregate must not abort the query");
        assert_eq!(ranked.len(), 1);
        assert_eq!(ranked[0]["key"]["dept"], json!("b"));
    }

    #[test]
    fn test_simple_count() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"name": "Alice", "age": 30}),
            json!({"name": "Bob", "age": 25}),
            json!({"name": "Charlie", "age": 35}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "count".to_string(),
                field: "*".to_string(),
                alias: None,
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(3));
    }

    #[test]
    fn test_sum_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"name": "Alice", "score": 90}),
            json!({"name": "Bob", "score": 85}),
            json!({"name": "Charlie", "score": 95}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "sum".to_string(),
                field: "score".to_string(),
                alias: None,
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(270.0));
    }

    #[test]
    fn test_average_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"name": "Alice", "age": 30}),
            json!({"name": "Bob", "age": 20}),
            json!({"name": "Charlie", "age": 40}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "avg".to_string(),
                field: "age".to_string(),
                alias: None,
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(30.0));
    }

    #[test]
    fn test_min_max_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"value": 10}),
            json!({"value": 50}),
            json!({"value": 30}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "min".to_string(),
                field: "value".to_string(),
                alias: Some("min_value".to_string()),
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(10.0));
    }

    #[test]
    fn test_grouped_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"city": "NYC", "sales": 100}),
            json!({"city": "LA", "sales": 150}),
            json!({"city": "NYC", "sales": 200}),
            json!({"city": "LA", "sales": 250}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "sum".to_string(),
                field: "sales".to_string(),
                alias: None,
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![GroupBySpec::from("city")],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["type"], "grouped_aggregation");

        let results = result["results"].as_array().unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_multiple_aggregations() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"score": 90}),
            json!({"score": 85}),
            json!({"score": 95}),
        ];

        let query = StatsQuery {
            aggregations: vec![
                AggregationSpec {
                    function: "sum".to_string(),
                    field: "score".to_string(),
                    alias: Some("total".to_string()),
                    params: HashMap::new(),
                    ..Default::default()
                },
                AggregationSpec {
                    function: "avg".to_string(),
                    field: "score".to_string(),
                    alias: Some("average".to_string()),
                    params: HashMap::new(),
                    ..Default::default()
                },
            ],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["type"], "multiple_aggregations");
        assert_eq!(result["results"]["total"], json!(270.0));
        assert_eq!(result["results"]["average"], json!(90.0));
    }

    #[test]
    fn test_median() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"value": 10}),
            json!({"value": 20}),
            json!({"value": 30}),
            json!({"value": 40}),
            json!({"value": 50}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "median".to_string(),
                field: "value".to_string(),
                alias: None,
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(30.0));
    }

    #[test]
    fn test_unique_count() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"city": "NYC"}),
            json!({"city": "LA"}),
            json!({"city": "NYC"}),
            json!({"city": "SF"}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "unique_count".to_string(),
                field: "city".to_string(),
                alias: None,
                params: HashMap::new(),
                ..Default::default()
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(3));
    }
}