fv-compute 0.2.0

//! The signature type system: a transform is a typed function over Arrow data.
//! Column types are **base types today**; value-types layer on as richer signatures later.
//! `FvType` is the language-neutral name in a `transform.toml`; `to_arrow()` is the binding to
//! the engine's Arrow line, so the contract stays authorable without an Arrow dependency in the
//! manifest itself.

use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

/// Base column types the contract admits. Kept deliberately small; extended with
/// value-types (currency, geo, classification) on the same seam later.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FvType {
    Bool,
    Int32,
    Int64,
    Float32,
    Float64,
    Utf8,
    Date32,
    /// UTC timestamp, millisecond precision (the plane's temporal default).
    TimestampMs,
}

impl FvType {
    /// Bind a contract type to the engine's Arrow `DataType`.
    pub fn to_arrow(self) -> DataType {
        match self {
            FvType::Bool => DataType::Boolean,
            FvType::Int32 => DataType::Int32,
            FvType::Int64 => DataType::Int64,
            FvType::Float32 => DataType::Float32,
            FvType::Float64 => DataType::Float64,
            FvType::Utf8 => DataType::Utf8,
            FvType::Date32 => DataType::Date32,
            FvType::TimestampMs => DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
        }
    }
}

/// One column in a signature: a name + its base type.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ColumnSpec {
    pub name: String,
    #[serde(rename = "type")]
    pub dtype: FvType,
}

impl ColumnSpec {
    /// A column with a name and a base type.
    pub fn new(name: impl Into<String>, dtype: FvType) -> Self {
        Self {
            name: name.into(),
            dtype,
        }
    }
}

impl From<(&str, FvType)> for ColumnSpec {
    fn from((name, dtype): (&str, FvType)) -> Self {
        Self::new(name, dtype)
    }
}

/// A named tabular signature (one input dataset, or the single output). `name` is optional for
/// the output; required-and-distinct for multi-input transforms (validated in the manifest).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SchemaSpec {
    #[serde(default)]
    pub name: Option<String>,
    pub columns: Vec<ColumnSpec>,
}

impl SchemaSpec {
    /// An unnamed signature from `(name, type)` pairs or [`ColumnSpec`]s.
    ///
    /// ```
    /// use fv_compute::{FvType, SchemaSpec};
    /// let sig = SchemaSpec::new([("id", FvType::Int64), ("amount", FvType::Float64)]);
    /// assert!(sig.has_column("amount"));
    /// ```
    pub fn new(columns: impl IntoIterator<Item = impl Into<ColumnSpec>>) -> Self {
        Self {
            name: None,
            columns: columns.into_iter().map(Into::into).collect(),
        }
    }

    /// A named signature (required when a transform has more than one input).
    pub fn named(name: impl Into<String>, columns: impl IntoIterator<Item = impl Into<ColumnSpec>>) -> Self {
        Self {
            name: Some(name.into()),
            ..Self::new(columns)
        }
    }

    /// Materialize as an Arrow `Schema`. All fields nullable=true (the plane is null-tolerant;
    /// tighter nullability is a refinement).
    pub fn to_arrow_schema(&self) -> Schema {
        Schema::new(
            self.columns
                .iter()
                .map(|c| Field::new(&c.name, c.dtype.to_arrow(), true))
                .collect::<Vec<_>>(),
        )
    }

    pub fn to_arrow_schema_ref(&self) -> Arc<Schema> {
        Arc::new(self.to_arrow_schema())
    }

    /// Does this signature contain a column of the given name?
    pub fn has_column(&self, name: &str) -> bool {
        self.columns.iter().any(|c| c.name == name)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fvtype_maps_to_expected_arrow() {
        assert_eq!(FvType::Int64.to_arrow(), DataType::Int64);
        assert_eq!(FvType::Float64.to_arrow(), DataType::Float64);
        assert_eq!(FvType::Utf8.to_arrow(), DataType::Utf8);
        assert_eq!(
            FvType::TimestampMs.to_arrow(),
            DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into()))
        );
    }

    #[test]
    fn schema_spec_builds_arrow_schema() {
        let s = SchemaSpec {
            name: Some("in".into()),
            columns: vec![
                ColumnSpec {
                    name: "id".into(),
                    dtype: FvType::Int64,
                },
                ColumnSpec {
                    name: "amount".into(),
                    dtype: FvType::Float64,
                },
            ],
        };
        let arrow = s.to_arrow_schema();
        assert_eq!(arrow.fields().len(), 2);
        assert_eq!(arrow.field(0).name(), "id");
        assert_eq!(arrow.field(1).data_type(), &DataType::Float64);
        assert!(arrow.field(0).is_nullable());
        assert!(s.has_column("amount") && !s.has_column("nope"));
    }
}