spiral-sdk 0.1.9

//! Exact table scan requests and results. The scandal contract defines the wire types; this
//! module adds their behavior and the runtime's own scan options.

use std::ops::Range;

use serde::{Deserialize, Serialize};
use spiral_contract::safelong::SafeUnsignedLong;
use uuid::Uuid;

pub use super::api::{
    FilterClause, FilterSupport, FragmentReceipt, PlannedSplit, PreparedWriteReceipt, RowRange,
    Scan, ScanInfo, ScanMetricsSnapshot, ScanPlan, Split, SplitTake, StorageCounters, StorageScope,
    TableVersion, WindowCounters,
};
use super::emit::Emit;
use super::error::ScandalError;
use super::expression::{Expression, NamedExpression};

/// The most positions one `TakeRows` request takes.
pub const MAX_TAKE_POSITIONS: usize = 65_536;
/// The most splits one `TakeRows` request takes from.
pub const MAX_TAKES: usize = 256;

impl Default for StorageCounters {
    fn default() -> Self {
        Self {
            scope: StorageScope::Complete,
            bytes: SafeUnsignedLong(0),
            requests: SafeUnsignedLong(0),
            read_time_ns: SafeUnsignedLong(0),
        }
    }
}

impl Default for WindowCounters {
    fn default() -> Self {
        Self {
            visited: SafeUnsignedLong(0),
            pruned_statistics: SafeUnsignedLong(0),
        }
    }
}

impl Default for ScanMetricsSnapshot {
    fn default() -> Self {
        Self {
            storage: None,
            windows: None,
            fragments_considered: SafeUnsignedLong(0),
            fragments_pruned: SafeUnsignedLong(0),
            files_considered: SafeUnsignedLong(0),
            files_pruned: SafeUnsignedLong(0),
            output_rows: SafeUnsignedLong(0),
        }
    }
}

impl ScanMetricsSnapshot {
    /// Add another part of the same execution, such as one split's reads, to these counters.
    pub fn add(&mut self, other: &Self) {
        self.storage = match (self.storage.take(), other.storage.clone()) {
            (Some(mut storage), Some(other)) => {
                storage.bytes.0 += other.bytes.0;
                storage.requests.0 += other.requests.0;
                storage.read_time_ns.0 += other.read_time_ns.0;
                Some(storage)
            }
            (storage, other) => storage.or(other),
        };
        self.windows = match (self.windows.take(), other.windows.clone()) {
            (Some(mut windows), Some(other)) => {
                windows.visited.0 += other.visited.0;
                windows.pruned_statistics.0 += other.pruned_statistics.0;
                Some(windows)
            }
            (windows, other) => windows.or(other),
        };
        self.fragments_considered.0 += other.fragments_considered.0;
        self.fragments_pruned.0 += other.fragments_pruned.0;
        self.files_considered.0 += other.files_considered.0;
        self.files_pruned.0 += other.files_pruned.0;
        self.output_rows.0 += other.output_rows.0;
    }
}

/// Exact scan options. A split token always names the same table version.
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ScanOptions {
    /// Collect native execution counters. Disabled by default.
    #[serde(default, skip_serializing_if = "is_false")]
    pub collect_metrics: bool,
    /// Allow unordered execution lanes to share scan work dynamically.
    /// Otherwise read in stored order through one lane.
    #[serde(default)]
    pub partitioned: bool,
    /// Maximum execution lanes (1–256), independent of the fragment count.
    /// Defaults to four for unordered scans. Limits and explicit splits use one lane.
    pub max_streams: Option<usize>,
    /// Top-level columns, or all columns when omitted.
    pub columns: Option<Vec<String>>,
    /// Named row-local output expressions. Mutually exclusive with `columns`.
    pub projection: Option<Vec<NamedExpression>>,
    /// Exact predicate; unsupported expressions fail before streaming starts.
    pub filter: Option<Expression>,
    /// Maximum returned rows, in stored order.
    #[serde(default, with = "super::safelong::option_u64")]
    pub limit: Option<u64>,
    /// Maximum rows in each delivered batch.
    pub batch_rows: Option<usize>,
    /// Opaque split descriptor tied to the requested immutable version.
    pub split_token: Option<String>,
    /// Half-open stored positions in this exact version, before filtering or deletion removal.
    /// Positions concatenate fragments in stored order and are not stable across rewrites.
    /// Incompatible with split tokens; execution is ordered.
    pub row_range: Option<ScanRowRange>,
    /// A top-k or aggregates of the output rows in place of the rows. A scan with an emit
    /// returns one merged result, reading up to `max_streams` splits at once, and takes no
    /// `limit`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub emit: Option<Emit>,
    /// Add a last output column of this name holding each row's stable row ID, a `u64` that
    /// survives rewrites and compaction. Invalid with an `emit`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub row_id_column: Option<String>,
    /// Keep only rows whose row IDs fall in these half-open ranges. Planning prunes files
    /// whose row IDs cannot match, and reads enforce the ranges exactly.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        with = "super::safelong::option_ranges"
    )]
    pub row_ids: Option<Vec<Range<u64>>>,
}

/// Half-open stored positions within one immutable Table version.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ScanRowRange {
    /// Inclusive stored position, before filtering and deletion removal.
    #[serde(with = "super::safelong::u64")]
    pub start: u64,
    /// Exclusive stored position.
    #[serde(with = "super::safelong::u64")]
    pub end: u64,
}

impl ScanRowRange {
    /// Reject a reversed range; empty ranges are valid.
    pub fn validate(self) -> Result<(), ScandalError> {
        if self.start > self.end {
            return Err(ScandalError::InvalidArgument {
                message: "rowRange start exceeds end".into(),
            });
        }
        Ok(())
    }
}

fn is_false(value: &bool) -> bool {
    !value
}

impl ScanOptions {
    /// Validate the concurrency budget and return the effective execution lane count.
    pub fn stream_count(&self) -> Result<usize, ScandalError> {
        if let Some(ranges) = &self.row_ids
            && ranges.iter().any(|range| range.start > range.end)
        {
            return Err(ScandalError::InvalidArgument {
                message: "a rowIds range starts after its end".into(),
            });
        }
        if let Some(range) = self.row_range {
            range.validate()?;
            if self.split_token.is_some() {
                return Err(ScandalError::InvalidArgument {
                    message: "rowRange and splitToken are mutually exclusive".into(),
                });
            }
        }
        let count = self.max_streams.unwrap_or(4);
        if !(1..=256).contains(&count) {
            return Err(ScandalError::InvalidArgument {
                message: "maxStreams must be between 1 and 256".into(),
            });
        }
        Ok(
            if self.partitioned
                && self.limit.is_none()
                && self.split_token.is_none()
                && self.row_range.is_none()
            {
                count
            } else {
                1
            },
        )
    }
}

/// Independently readable split of one immutable table version.
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TableSplit {
    /// Descriptor for scheduling; does not grant access to data.
    pub token: String,
    #[serde(default, with = "super::safelong::option_u64")]
    pub estimated_bytes: Option<u64>,
    #[serde(default, with = "super::safelong::option_u64")]
    pub estimated_rows: Option<u64>,
}

impl ScanInfo {
    /// Reject predicates that cannot be enforced exactly before any rows are read.
    pub fn require_exact_filter(&self) -> Result<(), ScandalError> {
        require_exact_filter(&self.filter_clauses)
    }
}

/// Reject predicate clauses that require residual evaluation outside the scan.
pub fn require_exact_filter(clauses: &[FilterClause]) -> Result<(), ScandalError> {
    if let Some(clause) = clauses
        .iter()
        .find(|clause| clause.support != FilterSupport::Exact)
    {
        return Err(ScandalError::InvalidArgument {
            message: clause
                .reason
                .clone()
                .unwrap_or_else(|| "scan predicate cannot be enforced exactly".to_owned()),
        });
    }
    Ok(())
}

impl Scan {
    /// Read every column of `table` without a filter or limit.
    pub fn new(table: TableVersion) -> Self {
        Self {
            table,
            columns: None,
            projection: None,
            filter: None,
            limit: None,
            batch_rows: None,
            collect_metrics: None,
            emit: None,
            row_range: None,
            row_id_column: None,
            row_ids: None,
        }
    }

    /// Whether the scan reports execution counters.
    pub fn collects_metrics(&self) -> bool {
        self.collect_metrics.unwrap_or(false)
    }
}

impl RowRange {
    /// Stored rows `start..end`.
    pub fn new(rows: Range<u64>) -> Self {
        Self {
            start: SafeUnsignedLong(rows.start),
            end: SafeUnsignedLong(rows.end),
        }
    }

    /// These stored rows as a range.
    pub fn range(&self) -> Range<u64> {
        self.start.0..self.end.0
    }
}

/// Registered immutable table data awaiting Catalog staging and publication.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TableWriteResult {
    /// Resource whose storage namespace contains the written data.
    pub resource_id: Uuid,
    /// Exact immutable version registered with the resource service.
    pub version_id: Uuid,
}

#[cfg(test)]
mod tests {
    use super::{ScanOptions, ScanRowRange, TableSplit};

    #[test]
    fn default_options_and_lossless_split_estimates_round_trip() {
        let options: ScanOptions = serde_json::from_str("{}").unwrap();
        assert!(!options.partitioned);
        assert!(!options.collect_metrics);
        assert!(
            serde_json::to_value(&options)
                .unwrap()
                .get("collectMetrics")
                .is_none()
        );
        assert_eq!(options.stream_count().unwrap(), 1);
        let mut parallel = ScanOptions {
            partitioned: true,
            max_streams: Some(3),
            ..Default::default()
        };
        assert_eq!(parallel.stream_count().unwrap(), 3);
        parallel.limit = Some(1);
        assert_eq!(parallel.stream_count().unwrap(), 1);
        parallel.max_streams = Some(0);
        assert!(parallel.stream_count().is_err());
        parallel.max_streams = Some(257);
        assert!(parallel.stream_count().is_err());
        assert!(options.filter.is_none());
        let split = TableSplit {
            token: "exact".to_owned(),
            estimated_rows: Some(u64::MAX),
            estimated_bytes: Some(12),
        };
        let json = serde_json::to_value(split).unwrap();
        assert_eq!(json["estimatedRows"], u64::MAX.to_string());
        assert_eq!(json["estimatedBytes"], 12);
        let decoded: TableSplit = serde_json::from_value(json).unwrap();
        assert_eq!(decoded.estimated_rows, Some(u64::MAX));
    }
    #[test]
    fn row_ranges_encode_losslessly_and_force_ordered_execution() {
        let options = ScanOptions {
            partitioned: true,
            max_streams: Some(4),
            row_range: Some(ScanRowRange {
                start: 3,
                end: u64::MAX,
            }),
            ..Default::default()
        };
        let json = serde_json::to_value(&options).unwrap();
        assert_eq!(
            json["rowRange"],
            serde_json::json!({"start": 3, "end": "18446744073709551615"})
        );
        let decoded: ScanOptions = serde_json::from_value(json).unwrap();
        assert_eq!(decoded.row_range, options.row_range);
        assert_eq!(decoded.stream_count().unwrap(), 1);
        assert!(
            ScanOptions {
                split_token: Some("split".into()),
                ..options.clone()
            }
            .stream_count()
            .is_err()
        );
        assert!(
            ScanOptions {
                row_range: Some(ScanRowRange { start: 4, end: 3 }),
                ..options
            }
            .stream_count()
            .is_err()
        );
    }
}