heliosdb-nano 3.37.3

//! Top-K operator: combined `Sort → Limit` for large inputs.
//!
//! Instead of fully sorting N rows (O(N log N)) and discarding all but the
//! first `k = limit + offset`, `TopKOperator` streams the input through a
//! bounded max-heap of size `k`, yielding O(N log k) complexity and O(k)
//! memory. The outer `LimitOperator` still applies the `offset` skip.
//!
//! This lands as part of the pagination work (see
//! FEATURE_REQUEST_pagination.md) so that `ORDER BY ... LIMIT N OFFSET M`
//! stays fast on large tables when the application can't use keyset
//! pagination.

use super::{compare_values, PhysicalOperator, TimeoutContext};
use crate::{Result, Schema, Tuple, Value};
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::sync::Arc;

/// Heap entry: the sort key vector + the tuple it came from. Ordered so
/// that `BinaryHeap` (a max-heap) keeps the *largest* sort key at the
/// top. When we hit capacity and the next tuple's sort key is smaller
/// than the heap's max, we pop the max and push the new entry — leaving
/// the `k` smallest at the end. Descending order (`asc = false`) is
/// handled by reversing per-field comparison when building the key.
struct HeapEntry {
    key: Vec<Value>,
    asc: Arc<Vec<bool>>,
    tuple: Tuple,
}

impl PartialEq for HeapEntry {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}
impl Eq for HeapEntry {}
impl PartialOrd for HeapEntry {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}
impl Ord for HeapEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        for (i, (a, b)) in self.key.iter().zip(other.key.iter()).enumerate() {
            let mut c = compare_values(a, b);
            // Apply asc/desc per column
            if !self.asc.get(i).copied().unwrap_or(true) {
                c = c.reverse();
            }
            if c != Ordering::Equal {
                return c;
            }
        }
        Ordering::Equal
    }
}

/// Streaming Top-K operator. Materialises only `k` tuples instead of the
/// full input.
pub struct TopKOperator {
    sorted: Vec<Tuple>,
    cursor: usize,
    schema: Arc<Schema>,
}

enum SortKeyAccessor {
    Column(usize),
    Expr(crate::sql::LogicalExpr),
}

impl TopKOperator {
    /// `k` is the number of rows to keep (typically `limit + offset`).
    /// The caller is responsible for wrapping in a `LimitOperator` to
    /// apply the actual `offset`/`limit` window.
    pub fn new(
        mut input: Box<dyn PhysicalOperator>,
        exprs: Vec<crate::sql::LogicalExpr>,
        asc: Vec<bool>,
        k: usize,
        parameters: Vec<Value>,
        timeout_ctx: Option<TimeoutContext>,
    ) -> Result<Self> {
        let schema = input.schema();
        // Must carry the query parameters: an `ORDER BY <expr>` key may
        // reference a bound parameter (e.g. the kNN idiom
        // `ORDER BY embedding <=> $1`). Without them the parameter resolves
        // to an error, the sort key collapses to NULL for every row, and the
        // heap returns rows in arbitrary order — silently wrong nearest
        // neighbors. See NANO-DEFICIENCIES A17.
        let evaluator = crate::sql::Evaluator::with_parameters(schema.clone(), parameters);
        let key_accessors: Vec<SortKeyAccessor> = exprs
            .into_iter()
            .map(|expr| match &expr {
                crate::sql::LogicalExpr::Column { table, name } => schema
                    .get_qualified_column_index(table.as_deref(), name)
                    .map(SortKeyAccessor::Column)
                    .unwrap_or_else(|| SortKeyAccessor::Expr(expr)),
                _ => SortKeyAccessor::Expr(expr),
            })
            .collect();
        let asc = Arc::new(asc);

        // Degenerate case: k == 0 ⇒ no rows to return.
        if k == 0 {
            return Ok(Self {
                sorted: Vec::new(),
                cursor: 0,
                schema,
            });
        }

        let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::with_capacity(k + 1);
        while let Some(tuple) = input.next()? {
            if let Some(ref ctx) = timeout_ctx {
                ctx.check_timeout()?;
            }
            // Build the sort key for this tuple
            let mut key = Vec::with_capacity(key_accessors.len());
            for accessor in &key_accessors {
                match accessor {
                    SortKeyAccessor::Column(index) => {
                        key.push(tuple.get(*index).cloned().unwrap_or(Value::Null));
                    }
                    SortKeyAccessor::Expr(expr) => match evaluator.evaluate(expr, &tuple) {
                        Ok(v) => key.push(v),
                        Err(_) => key.push(Value::Null),
                    },
                }
            }
            let entry = HeapEntry {
                key,
                asc: Arc::clone(&asc),
                tuple,
            };
            if heap.len() < k {
                heap.push(entry);
            } else {
                // Peek at heap max (largest-so-far)
                let replace = heap.peek().map(|top| entry < *top).unwrap_or(false);
                if replace {
                    heap.pop();
                    heap.push(entry);
                }
            }
        }

        // Drain heap in ascending order of the comparator — this is the
        // final top-k result in correct order.
        let mut entries: Vec<HeapEntry> = heap.into_sorted_vec();
        // `into_sorted_vec` returns ascending order for a max-heap, which
        // (because our Ord already applies asc/desc) gives the user-visible
        // sort direction directly.
        let sorted: Vec<Tuple> = entries.drain(..).map(|e| e.tuple).collect();

        Ok(Self {
            sorted,
            cursor: 0,
            schema,
        })
    }

    /// No-op — timeout is handled during construction, kept for API symmetry.
    pub fn with_timeout(self, _: Option<TimeoutContext>) -> Self {
        self
    }
}

impl PhysicalOperator for TopKOperator {
    fn next(&mut self) -> Result<Option<Tuple>> {
        let out = self.sorted.get(self.cursor).cloned();
        if out.is_some() {
            self.cursor += 1;
        }
        Ok(out)
    }

    fn schema(&self) -> Arc<Schema> {
        self.schema.clone()
    }
}