spiral-sdk 0.1.7

//! The SPQL type registry's types as Arrow extension fields.
//!
//! Each type is a struct, such as [`Blob`] or [`Duration`], whose fields are its type parameters,
//! and implements [`SpqlType`] with its catalog facts and its Arrow storage. A column of a type
//! is its storage, tagged with the field metadata `ARROW:extension:name` and
//! `ARROW:extension:metadata`. Writers name the type's current Arrow name; readers also accept
//! the old ids in its `ALIASES`.

use std::collections::BTreeMap;

use arrow_schema::{DataType, Field, TimeUnit};

mod registry;

pub use registry::*;

/// The Arrow field metadata key that names a field's extension type.
pub const EXTENSION_NAME_KEY: &str = "ARROW:extension:name";
/// The Arrow field metadata key that holds a field's extension metadata.
pub const EXTENSION_METADATA_KEY: &str = "ARROW:extension:metadata";

/// A type of the SPQL type registry. A value of the implementing struct is one instance of the
/// type, with its type parameters.
pub trait SpqlType: Sized {
    /// The type id that writers name.
    const ID: &'static str;
    /// The PascalCase type name.
    const NAME: &'static str;
    /// The Arrow extension name the type is exported under.
    const ARROW_NAME: &'static str;
    /// Old ids that still read as this type.
    const ALIASES: &'static [&'static str];
    const FAMILY: &'static str;
    const KIND: Kind;
    const STABILITY: Stability;
    /// The catalog version that introduced the type.
    const SINCE: &'static str;
    /// The id of the type that replaces a deprecated one.
    const DEPRECATED_BY: Option<&'static str>;
    /// The property bridges, which the extension metadata carries.
    const BRIDGES: Bridges;

    /// The Arrow storage of a column of this type.
    fn storage_type(&self) -> DataType;

    /// The type parameters, as the extension metadata carries them.
    fn params(&self) -> Vec<(&'static str, &'static str)> {
        Vec::new()
    }

    /// The type with the parameters that its extension metadata carries.
    fn from_params(params: &mut ExtensionParams) -> Result<Self, TypeError>;

    /// The canonical extension metadata: the bridges and parameters as one JSON object with
    /// sorted keys and no whitespace.
    fn extension_metadata(&self) -> String {
        let mut object = BTreeMap::from([
            ("order", Self::BRIDGES.order.as_str()),
            ("eq", Self::BRIDGES.eq.as_str()),
        ]);
        object.extend(self.params());
        serde_json::to_string(&object).unwrap_or_else(|_| unreachable!("string maps serialize"))
    }

    /// A field named `name` of this type, tagged with its extension name and metadata.
    fn field(&self, name: impl Into<String>, nullable: bool) -> Field {
        Field::new(name, self.storage_type(), nullable).with_metadata(
            [
                (EXTENSION_NAME_KEY.to_owned(), Self::ARROW_NAME.to_owned()),
                (EXTENSION_METADATA_KEY.to_owned(), self.extension_metadata()),
            ]
            .into(),
        )
    }

    /// Whether `field` carries this type, under its current extension name or an alias.
    fn matches(field: &Field) -> bool {
        field
            .metadata()
            .get(EXTENSION_NAME_KEY)
            .is_some_and(|name| {
                name == Self::ARROW_NAME
                    || name == Self::ID
                    || Self::ALIASES.contains(&name.as_str())
            })
    }

    /// The type that `field` carries, with its parameters, or `None` when it carries another.
    fn from_field(field: &Field) -> Result<Option<Self>, TypeError> {
        if !Self::matches(field) {
            return Ok(None);
        }
        let metadata = field
            .metadata()
            .get(EXTENSION_METADATA_KEY)
            .map_or("", String::as_str);
        let mut params = ExtensionParams::read(Self::ID, metadata)?;
        let type_ = Self::from_params(&mut params)?;
        params.finish()?;
        Ok(Some(type_))
    }
}

/// What a type is in the catalog.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Kind {
    Component,
    Archetype,
    Encoded,
    Reference,
    Index,
    Decoded,
}

/// How settled a type is.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Stability {
    Experimental,
    Stable,
    Deprecated,
}

/// Whether a reader without a type may use its storage for one property of its values.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Bridge {
    /// The storage does not carry the property.
    None,
    /// The storage's own ordering or equality is the type's.
    Storage,
    /// Values are equal when they name the same object.
    Identity,
}

impl Bridge {
    /// The bridge as the extension metadata spells it.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::None => "none",
            Self::Storage => "storage",
            Self::Identity => "identity",
        }
    }
}

/// The `order` and `eq` property bridges.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Bridges {
    pub order: Bridge,
    pub eq: Bridge,
}

/// A type parameter, which a type's extension metadata carries as a string.
pub trait TypeParam: Sized {
    /// The parameter's metadata string.
    fn to_param(&self) -> &'static str;

    /// The parameter its metadata string names.
    fn from_param(value: &str) -> Option<Self>;
}

/// The unit of a `spql.Duration`. Arrow carries a duration in one of these units.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum DurationUnit {
    Seconds,
    Milliseconds,
    Microseconds,
    Nanoseconds,
}

impl DurationUnit {
    /// The Arrow time unit of a duration column.
    pub fn arrow_unit(self) -> TimeUnit {
        match self {
            Self::Seconds => TimeUnit::Second,
            Self::Milliseconds => TimeUnit::Millisecond,
            Self::Microseconds => TimeUnit::Microsecond,
            Self::Nanoseconds => TimeUnit::Nanosecond,
        }
    }
}

impl TypeParam for DurationUnit {
    fn to_param(&self) -> &'static str {
        match self {
            Self::Seconds => "s",
            Self::Milliseconds => "ms",
            Self::Microseconds => "us",
            Self::Nanoseconds => "ns",
        }
    }

    fn from_param(value: &str) -> Option<Self> {
        Some(match value {
            "s" => Self::Seconds,
            "ms" => Self::Milliseconds,
            "us" => Self::Microseconds,
            "ns" => Self::Nanoseconds,
            _ => return None,
        })
    }
}

/// The calendar frequency of a `spql.Period`.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum PeriodFrequency {
    Day,
    Week,
    MonthStart,
    MonthEnd,
    QuarterStart,
    QuarterEnd,
    YearStart,
    YearEnd,
}

impl TypeParam for PeriodFrequency {
    fn to_param(&self) -> &'static str {
        match self {
            Self::Day => "D",
            Self::Week => "W",
            Self::MonthStart => "MS",
            Self::MonthEnd => "ME",
            Self::QuarterStart => "QS",
            Self::QuarterEnd => "QE",
            Self::YearStart => "YS",
            Self::YearEnd => "YE",
        }
    }

    fn from_param(value: &str) -> Option<Self> {
        Some(match value {
            "D" => Self::Day,
            "W" => Self::Week,
            "MS" => Self::MonthStart,
            "ME" => Self::MonthEnd,
            "QS" => Self::QuarterStart,
            "QE" => Self::QuarterEnd,
            "YS" => Self::YearStart,
            "YE" => Self::YearEnd,
            _ => return None,
        })
    }
}

/// The parameters of a field's extension metadata, which [`SpqlType::from_params`] takes in turn.
#[derive(Debug)]
pub struct ExtensionParams {
    id: &'static str,
    params: BTreeMap<String, String>,
}

impl ExtensionParams {
    /// The parameters in the metadata of a field of type `id`. Metadata written before the
    /// catalog recorded bridges is empty.
    fn read(id: &'static str, metadata: &str) -> Result<Self, TypeError> {
        let mut params = if metadata.is_empty() {
            BTreeMap::new()
        } else {
            serde_json::from_str::<BTreeMap<String, String>>(metadata)
                .map_err(|source| TypeError::Metadata { id, source })?
        };
        params.remove("order");
        params.remove("eq");
        Ok(Self { id, params })
    }

    /// Takes the parameter `name`.
    pub fn take<P: TypeParam>(&mut self, name: &'static str) -> Result<P, TypeError> {
        let id = self.id;
        let value = self
            .params
            .remove(name)
            .ok_or(TypeError::MissingParam { id, name })?;
        P::from_param(&value).ok_or(TypeError::InvalidParam { id, name, value })
    }

    fn finish(self) -> Result<(), TypeError> {
        match self.params.into_keys().next() {
            Some(name) => Err(TypeError::UnknownParam { id: self.id, name }),
            None => Ok(()),
        }
    }
}

/// Extension metadata that does not describe its type.
#[derive(Debug, thiserror::Error)]
pub enum TypeError {
    #[error("{id} extension metadata is not a JSON object of strings")]
    Metadata {
        id: &'static str,
        source: serde_json::Error,
    },
    #[error("{id} metadata must name its `{name}`")]
    MissingParam {
        id: &'static str,
        name: &'static str,
    },
    #[error("{id} parameter `{name}` cannot be `{value}`")]
    InvalidParam {
        id: &'static str,
        name: &'static str,
        value: String,
    },
    #[error("{id} has no parameter `{name}`")]
    UnknownParam { id: &'static str, name: String },
}

#[cfg(test)]
mod tests {
    use std::fs::File;

    use arrow_ipc::reader::StreamReader;
    use arrow_schema::{DataType, Field, TimeUnit};

    use super::{
        Blob, Duration, DurationUnit, EXTENSION_METADATA_KEY, EXTENSION_NAME_KEY, EncodedMedia,
        SpqlType, TypeError,
    };

    #[test]
    fn fields_carry_the_extension_name_and_canonical_metadata() {
        let field = Blob.field("video", false);
        assert_eq!(
            field.metadata()[EXTENSION_METADATA_KEY],
            r#"{"eq":"identity","order":"none"}"#
        );
        assert!(Blob::matches(&field));
        assert!(!EncodedMedia::matches(&field));

        let DataType::Struct(fields) = EncodedMedia.storage_type() else {
            panic!("EncodedMedia has struct storage");
        };
        assert!(Blob::matches(&fields[0]));
    }

    /// A `spql.EncodedMedia` column that spiral-scandal-runtime's functions tests wrote.
    #[test]
    fn storage_is_what_the_engine_writes() {
        let path = concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../../typescript/sdk/fixtures/rust/encoded_media.arrows"
        );
        let reader = StreamReader::try_new(File::open(path).unwrap(), None).unwrap();
        let written = reader.schema().field(0).clone();
        assert_eq!(
            EncodedMedia.field(written.name(), written.is_nullable()),
            written
        );
    }

    #[test]
    fn aliases_read_as_the_current_type() {
        let old = Field::new("video", Blob.storage_type(), true)
            .with_metadata([(EXTENSION_NAME_KEY.to_owned(), "spql.io.Blob".to_owned())].into());
        assert_eq!(Blob::from_field(&old).unwrap(), Some(Blob));
        assert_eq!(EncodedMedia::from_field(&old).unwrap(), None);
    }

    #[test]
    fn parameters_are_struct_fields() {
        let duration = Duration {
            unit: DurationUnit::Milliseconds,
        };
        let field = duration.field("elapsed", true);
        assert_eq!(
            field.data_type(),
            &DataType::Duration(TimeUnit::Millisecond)
        );
        assert_eq!(
            field.metadata()[EXTENSION_METADATA_KEY],
            r#"{"eq":"storage","order":"storage","unit":"ms"}"#
        );
        assert_eq!(Duration::from_field(&field).unwrap(), Some(duration));

        let untagged = Field::new("elapsed", DataType::Duration(TimeUnit::Millisecond), true)
            .with_metadata([(EXTENSION_NAME_KEY.to_owned(), "spql.Duration".to_owned())].into());
        assert!(matches!(
            Duration::from_field(&untagged),
            Err(TypeError::MissingParam { name: "unit", .. })
        ));
    }
}