spiral-sdk 0.1.11

//! Transport-independent descriptions of planned and executed operations.

use std::collections::BTreeMap;
use std::fmt;

use arrow_schema::Schema;
use serde::{Deserialize, Serialize};

use super::ScandalError;
use super::profile::{EXECUTION_OPTIONS_KEY, ExecutionProfile, PROFILE_METADATA_KEY};

/// Reserved Arrow schema metadata carrying an operation's versioned plan description.
/// It describes execution, not the logical type of the rows.
pub const PLAN_METADATA_KEY: &str = "spiral.execution.plan";

/// A node in a logical or physical plan. IDs are local to that graph.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PlanNode {
    /// Stable position within this plan, independent of allocation addresses.
    pub id: String,
    /// Operator or implementation name, without literal payloads.
    pub operator: String,
    /// Dependencies in this graph, in evaluation order.
    pub inputs: Vec<String>,
    /// Operator-specific facts. Unknown facts are omitted, not reported as zero.
    pub properties: BTreeMap<String, String>,
}

/// One candidate considered during physical selection.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PlanDecision {
    /// Node in the logical graph, absent for a compiler-synthesized expression.
    pub node: Option<String>,
    /// Compiler or implementation considered.
    pub candidate: String,
    /// Machine-readable outcome, such as `accepted` or `declined`.
    pub outcome: String,
    /// Stable reason code; an unspecified decline is not a capability diagnosis.
    pub reason: String,
}

/// The shared description of a scan, projection or other prepared operation.
///
/// Bind/scan planning returns a preview (`execution_id` is absent). Execution returns the
/// description it actually used, with a fresh execution ID. Fingerprints are scoped to the
/// implementation/build and description version; they are not durable cross-version IDs.
/// Literal payloads and credentials must never appear in nodes or decision reasons.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PlanDescription {
    /// Description format version. This implementation reads version 1.
    pub version: u32,
    /// Fingerprint of the semantic request, including literal options but no credentials.
    pub semantic_id: String,
    /// Opaque identity of the bound implementation and physical description. A merged
    /// delivery can name multiple physical variants when its hosts selected different plans.
    pub physical_id: String,
    /// An actual execution, or the observed execution IDs combined by a delivery merge.
    /// Absent for a planning preview.
    pub execution_id: Option<String>,
    /// Original operators and output provenance, before implementation lowering.
    pub logical: Vec<PlanNode>,
    /// The operations the current executor performs, including forced transfers.
    pub physical: Vec<PlanNode>,
    /// Optional aggregate observations, absent when profiling was not requested.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub profile: Option<ExecutionProfile>,
    /// Selection decisions in the order candidates were offered.
    pub decisions: Vec<PlanDecision>,
}

impl PlanDescription {
    /// Read the common plan attachment from a bind, scan or execution schema.
    /// Older hosts may omit it. Unknown versions fail instead of being misinterpreted.
    pub fn from_schema(schema: &Schema) -> Result<Option<Self>, ScandalError> {
        let Some(encoded) = schema.metadata().get(PLAN_METADATA_KEY) else {
            return Ok(None);
        };
        if encoded.len() > 4 * 1024 * 1024 {
            return Err(invalid("plan description exceeds 4 MiB"));
        }
        let plan: Self = serde_json::from_str(encoded).map_err(invalid)?;
        if plan.version != 1 {
            return Err(invalid(format!(
                "unsupported plan description version {}",
                plan.version
            )));
        }
        Ok(Some(plan))
    }

    /// Attach this description without changing fields or application metadata.
    pub fn with_schema(&self, schema: &Schema) -> Result<Schema, ScandalError> {
        let mut metadata = schema.metadata().clone();
        if self.profile.is_some() {
            metadata.insert(PROFILE_METADATA_KEY.into(), "1".into());
        } else {
            metadata.remove(PROFILE_METADATA_KEY);
        }
        let encoded = serde_json::to_string(self).map_err(invalid)?;
        if encoded.len() > 4 * 1024 * 1024 {
            return Err(invalid("plan description exceeds 4 MiB"));
        }
        metadata.insert(PLAN_METADATA_KEY.into(), encoded);
        Ok(schema.clone().with_metadata(metadata))
    }

    /// Remove execution diagnostics when comparing or binding logical row types.
    pub fn row_schema(schema: &Schema) -> Schema {
        let metadata = schema
            .metadata()
            .iter()
            .filter(|(key, _)| {
                key.as_str() != PLAN_METADATA_KEY
                    && key.as_str() != EXECUTION_OPTIONS_KEY
                    && key.as_str() != PROFILE_METADATA_KEY
            })
            .map(|(key, value)| (key.clone(), value.clone()))
            .collect();
        Schema::new_with_metadata(schema.fields().clone(), metadata)
    }
}

impl fmt::Display for PlanDescription {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "plan={} semantic={}", self.physical_id, self.semantic_id)?;
        match &self.execution_id {
            Some(id) => write!(f, " execution={id}")?,
            None => f.write_str(" preview")?,
        }
        for node in &self.physical {
            write!(f, "\n  {} {}", node.id, node.operator)?;
            if !node.inputs.is_empty() {
                write!(f, " <- [{}]", node.inputs.join(", "))?;
            }
            for (name, value) in &node.properties {
                write!(f, " {name}=\"{}\"", value.escape_debug())?;
            }
        }
        if let Some(profile) = &self.profile {
            write!(f, "\n{profile}")?;
        }
        if f.alternate() {
            for node in &self.logical {
                write!(
                    f,
                    "\n  logical {} {} <- [{}]",
                    node.id,
                    node.operator,
                    node.inputs.join(", ")
                )?;
                for (name, value) in &node.properties {
                    write!(f, " {name}=\"{}\"", value.escape_debug())?;
                }
            }
            for decision in &self.decisions {
                write!(
                    f,
                    "\n  candidate {} {} {} ({})",
                    decision.node.as_deref().unwrap_or("unmapped"),
                    decision.candidate,
                    decision.outcome,
                    decision.reason
                )?;
            }
        }
        Ok(())
    }
}

fn invalid(error: impl fmt::Display) -> ScandalError {
    ScandalError::InvalidArgument {
        message: error.to_string(),
    }
}