hawdb-executor 0.15.0

// Copyright 2026 Nowledge
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Query-owned execution reports over storage-neutral operator events.

use super::ExecutionObserver;
use crate::binding::Binding;
use hawdb_plan_cypher::{visit_plan_with_ids, PhysicalOperatorId, PhysicalPlan};
use hawdb_storage::scan::ScanPruningReport;
use std::cell::RefCell;
use std::collections::{BTreeMap, BTreeSet};

#[derive(Default)]
pub struct QueryExecutionReports {
    pub operator_cardinality: Vec<crate::OperatorCardinalityProfile>,
    pub scan_pruning: Vec<ScanPruningReport>,
    pub vector_execution: Vec<crate::VectorExecutionReport>,
    pub graph_expansion: Vec<crate::GraphExpansionExecutionReport>,
    pub text_seed_execution: Vec<crate::TextSeedExecutionReport>,
    pub grouped_evidence: Vec<crate::GroupedEvidenceExecutionReport>,
    pub blocking_memory: Vec<crate::BlockingOperatorMemoryReport>,
    pub pipeline_memory: crate::PipelineMemoryReport,
}

/// Per-query collection over a plan whose node addresses remain stable while
/// operator events are recorded. Final host/process metrics are added outside.
pub struct QueryExecutionObserver {
    reports: RefCell<QueryExecutionReports>,
    seed_graph_input: Option<hawdb_plan_cypher::ScoringSeedGraphInput>,
    operator_ids: BTreeMap<usize, PhysicalOperatorId>,
    grouped_input_ids: BTreeMap<usize, PhysicalOperatorId>,
}

impl Default for QueryExecutionObserver {
    fn default() -> Self {
        Self {
            reports: RefCell::new(QueryExecutionReports::default()),
            seed_graph_input: None,
            operator_ids: BTreeMap::new(),
            grouped_input_ids: BTreeMap::new(),
        }
    }
}

impl QueryExecutionObserver {
    pub fn new(plan: &PhysicalPlan) -> Self {
        let mut reports = QueryExecutionReports::default();
        let mut operator_ids = BTreeMap::new();
        let mut grouped_input_ids = BTreeMap::new();
        visit_plan_with_ids(plan, &mut |operator_id, operator| {
            let previous = operator_ids.insert(plan_address(operator), operator_id);
            debug_assert!(previous.is_none());
            match operator {
                PhysicalPlan::TextSeedScan {
                    complete, top_k, ..
                } => {
                    reports
                        .text_seed_execution
                        .push(crate::TextSeedExecutionReport {
                            operator_id: Some(operator_id),
                            complete_requested: *complete,
                            candidate_limit: *top_k,
                            ..Default::default()
                        });
                }
                PhysicalPlan::AggregateExec { items, input, .. }
                    if items.iter().any(|item| {
                        matches!(
                            item.target,
                            hawdb_plan_cypher::AggregateTarget::Ranked { .. }
                        )
                    }) =>
                {
                    grouped_input_ids.insert(plan_address(input), operator_id);
                    reports
                        .grouped_evidence
                        .push(crate::GroupedEvidenceExecutionReport {
                            operator_id: Some(operator_id),
                            ..Default::default()
                        });
                }
                _ => {}
            }
            reports
                .operator_cardinality
                .push(crate::OperatorCardinalityProfile {
                    operator_id,
                    operator: operator.kind(),
                    actual_rows: None,
                });
        });
        Self {
            reports: RefCell::new(reports),
            seed_graph_input: match plan {
                PhysicalPlan::ScoringProgramExec {
                    seed_graph_input, ..
                } => seed_graph_input.clone(),
                PhysicalPlan::HostScoringExec { scoring, .. } => {
                    scoring.seed_graph_input().cloned()
                }
                _ => None,
            },
            operator_ids,
            grouped_input_ids,
        }
    }

    pub(crate) fn operator_id(&self, plan: &PhysicalPlan) -> Option<PhysicalOperatorId> {
        self.operator_ids.get(&plan_address(plan)).copied()
    }

    pub(crate) fn record_text_seed_execution(&self, report: crate::TextSeedExecutionReport) {
        let mut reports = self.reports.borrow_mut();
        if let Some(current) = reports
            .text_seed_execution
            .iter_mut()
            .find(|current| current.operator_id == report.operator_id)
        {
            *current = report;
        } else {
            reports.text_seed_execution.push(report);
        }
    }

    fn record_grouped_evidence(
        &self,
        input: &PhysicalPlan,
        mut report: crate::GroupedEvidenceExecutionReport,
    ) {
        report.operator_id = self.grouped_input_ids.get(&plan_address(input)).copied();
        let mut reports = self.reports.borrow_mut();
        if let Some(current) = reports
            .grouped_evidence
            .iter_mut()
            .find(|current| current.operator_id == report.operator_id)
        {
            *current = report;
        } else {
            reports.grouped_evidence.push(report);
        }
    }

    pub fn into_reports(self) -> QueryExecutionReports {
        self.reports.into_inner()
    }

    pub fn record_operator_start(&self, plan: &PhysicalPlan) {
        let Some(operator_id) = self.operator_ids.get(&plan_address(plan)).copied() else {
            debug_assert!(self.operator_ids.is_empty());
            return;
        };
        let mut reports = self.reports.borrow_mut();
        let report = reports
            .operator_cardinality
            .get_mut(operator_id.ordinal())
            .expect("operator cardinality profile follows plan ordinals");
        debug_assert_eq!(report.operator, plan.kind());
        report.actual_rows.get_or_insert(0);
    }

    pub fn record_operator_output(&self, plan: &PhysicalPlan, output_rows: usize) {
        let Some(operator_id) = self.operator_ids.get(&plan_address(plan)).copied() else {
            debug_assert!(self.operator_ids.is_empty());
            return;
        };
        let mut reports = self.reports.borrow_mut();
        let report = reports
            .operator_cardinality
            .get_mut(operator_id.ordinal())
            .expect("operator cardinality profile follows plan ordinals");
        debug_assert_eq!(report.operator, plan.kind());
        report.actual_rows = Some(
            report
                .actual_rows
                .unwrap_or_default()
                .saturating_add(output_rows),
        );
    }

    pub fn record_vector_execution(&self, report: crate::VectorExecutionReport) {
        self.reports.borrow_mut().vector_execution.push(report);
    }
    pub fn record_graph_expansion(&self, report: crate::GraphExpansionExecutionReport) {
        self.reports.borrow_mut().graph_expansion.push(report);
    }

    pub fn record_blocking_memory_report(&self, report: crate::BlockingOperatorMemoryReport) {
        self.reports.borrow_mut().blocking_memory.push(report);
    }

    pub fn record_pipeline_batch(&self, batch: &[Binding]) {
        let payload_bytes = batch.iter().fold(0usize, |total, binding| {
            total.saturating_add(crate::binding::binding_payload_bytes(binding))
        });
        self.record_pipeline_batch_shape(batch.len(), payload_bytes);
    }

    pub(crate) fn record_pipeline_batch_shape(&self, rows: usize, payload_bytes: usize) {
        let mut reports = self.reports.borrow_mut();
        let report = &mut reports.pipeline_memory;
        report.intermediate_rows = report.intermediate_rows.saturating_add(rows);
        report.intermediate_payload_bytes = report
            .intermediate_payload_bytes
            .saturating_add(payload_bytes);
        report.peak_batch_rows = report.peak_batch_rows.max(rows);
        report.peak_batch_payload_bytes = report.peak_batch_payload_bytes.max(payload_bytes);
    }

    pub(crate) fn record_native_projection_row(&self, payload_bytes: usize) {
        self.record_pipeline_batch_shape(1, payload_bytes);
        let mut reports = self.reports.borrow_mut();
        let report = &mut reports.pipeline_memory;
        report.native_projection_rows = report.native_projection_rows.saturating_add(1);
    }

    pub fn record_columnar_batch(&self, input_rows: usize, selected_rows: usize) {
        let mut reports = self.reports.borrow_mut();
        let report = &mut reports.pipeline_memory;
        report.columnar_batches = report.columnar_batches.saturating_add(1);
        report.columnar_input_rows = report.columnar_input_rows.saturating_add(input_rows);
        report.columnar_selected_rows = report.columnar_selected_rows.saturating_add(selected_rows);
    }

    pub fn record_morsels(&self, count: usize) {
        let mut reports = self.reports.borrow_mut();
        let report = &mut reports.pipeline_memory;
        report.morsel_count = report.morsel_count.saturating_add(count);
    }

    pub fn record_morsel_admission(&self, max_workers: usize, active_workers: usize) {
        let mut reports = self.reports.borrow_mut();
        let report = &mut reports.pipeline_memory;
        report.morsel_max_admitted_workers = report.morsel_max_admitted_workers.max(max_workers);
        report.morsel_peak_active_workers = report.morsel_peak_active_workers.max(active_workers);
    }

    pub fn record_morsel_buffering(
        &self,
        peak_outputs: usize,
        peak_output_bytes: usize,
        peak_reorder_entries: usize,
    ) {
        let mut reports = self.reports.borrow_mut();
        let report = &mut reports.pipeline_memory;
        report.morsel_peak_buffered_outputs = report.morsel_peak_buffered_outputs.max(peak_outputs);
        report.morsel_peak_buffered_output_bytes = report
            .morsel_peak_buffered_output_bytes
            .max(peak_output_bytes);
        report.morsel_peak_reorder_entries =
            report.morsel_peak_reorder_entries.max(peak_reorder_entries);
    }

    pub fn current_vector_rerank_count(&self) -> usize {
        self.reports
            .borrow()
            .vector_execution
            .last()
            .map_or(0, |report| report.reranked_candidate_count)
    }
}

fn plan_address(plan: &PhysicalPlan) -> usize {
    std::ptr::from_ref(plan).addr()
}

impl ExecutionObserver for QueryExecutionObserver {
    fn vector_graph_scoring_input(&self) -> Option<&hawdb_plan_cypher::ScoringVectorGraphInput> {
        self.seed_graph_input
            .as_ref()
            .filter(|source| source.kind() == hawdb_plan_cypher::ScoringSeedKind::Vector)
    }

    fn seed_graph_scoring_input(&self) -> Option<&hawdb_plan_cypher::ScoringSeedGraphInput> {
        self.seed_graph_input.as_ref()
    }

    fn record_grouped_evidence_report(
        &self,
        input: &PhysicalPlan,
        report: crate::GroupedEvidenceExecutionReport,
    ) {
        self.record_grouped_evidence(input, report);
    }

    fn record_borrowed_node_projection_rows(&self, rows: usize) {
        let mut reports = self.reports.borrow_mut();
        let count = &mut reports.pipeline_memory.borrowed_node_projection_rows;
        *count = count.saturating_add(rows);
    }

    fn record_scan_pruning_report(&self, report: ScanPruningReport) {
        self.reports.borrow_mut().scan_pruning.push(report);
    }

    fn record_blocking_memory_report(&self, report: crate::BlockingOperatorMemoryReport) {
        QueryExecutionObserver::record_blocking_memory_report(self, report);
    }
}

pub fn blocking_operator_kinds(plan: &PhysicalPlan) -> Vec<String> {
    let mut output = BTreeSet::new();
    collect_blocking_operator_kinds(plan, &mut output);
    output.into_iter().collect()
}

fn collect_blocking_operator_kinds(plan: &PhysicalPlan, output: &mut BTreeSet<String>) {
    visit_plan_with_ids(plan, &mut |_, operator| {
        if matches!(
            operator,
            PhysicalPlan::GraphAlgorithm { .. }
                | PhysicalPlan::GraphAlgorithmOutput(_)
                | PhysicalPlan::VectorIndexScan(_)
                | PhysicalPlan::VectorSeedScan { .. }
                | PhysicalPlan::GraphSeedScan { .. }
                | PhysicalPlan::TextSeedScan { .. }
                | PhysicalPlan::ShortestPathExec { .. }
                | PhysicalPlan::AggregateExec { .. }
                | PhysicalPlan::DistinctExec { .. }
                | PhysicalPlan::SortExec { .. }
                | PhysicalPlan::TopNExec { .. }
                | PhysicalPlan::ScoringRerankExec { .. }
                | PhysicalPlan::ScoringProgramExec { .. }
                | PhysicalPlan::HostScoringExec { .. }
                | PhysicalPlan::NodeCartesianProductExec { .. }
                | PhysicalPlan::HashJoinExec { .. }
        ) {
            output.insert(operator.kind().as_str().to_string());
        }
    });
}

#[cfg(test)]
mod tests;