Skip to main content

fv_compute/
types.rs

1//! The signature type system: a transform is a typed function over Arrow data.
2//! Column types are **base types today**; value-types layer on as richer signatures later.
3//! `FvType` is the language-neutral name in a `transform.toml`; `to_arrow()` is the binding to
4//! the engine's Arrow line, so the contract stays authorable without an Arrow dependency in the
5//! manifest itself.
6
7use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10
11/// Base column types the contract admits. Kept deliberately small; extended with
12/// value-types (currency, geo, classification) on the same seam later.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum FvType {
16    Bool,
17    Int32,
18    Int64,
19    Float32,
20    Float64,
21    Utf8,
22    Date32,
23    /// UTC timestamp, millisecond precision (the plane's temporal default).
24    TimestampMs,
25}
26
27impl FvType {
28    /// Bind a contract type to the engine's Arrow `DataType`.
29    pub fn to_arrow(self) -> DataType {
30        match self {
31            FvType::Bool => DataType::Boolean,
32            FvType::Int32 => DataType::Int32,
33            FvType::Int64 => DataType::Int64,
34            FvType::Float32 => DataType::Float32,
35            FvType::Float64 => DataType::Float64,
36            FvType::Utf8 => DataType::Utf8,
37            FvType::Date32 => DataType::Date32,
38            FvType::TimestampMs => DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
39        }
40    }
41}
42
43/// One column in a signature: a name + its base type.
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
45pub struct ColumnSpec {
46    pub name: String,
47    #[serde(rename = "type")]
48    pub dtype: FvType,
49}
50
51impl ColumnSpec {
52    /// A column with a name and a base type.
53    pub fn new(name: impl Into<String>, dtype: FvType) -> Self {
54        Self {
55            name: name.into(),
56            dtype,
57        }
58    }
59}
60
61impl From<(&str, FvType)> for ColumnSpec {
62    fn from((name, dtype): (&str, FvType)) -> Self {
63        Self::new(name, dtype)
64    }
65}
66
67/// A named tabular signature (one input dataset, or the single output). `name` is optional for
68/// the output; required-and-distinct for multi-input transforms (validated in the manifest).
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
70pub struct SchemaSpec {
71    #[serde(default)]
72    pub name: Option<String>,
73    pub columns: Vec<ColumnSpec>,
74}
75
76impl SchemaSpec {
77    /// An unnamed signature from `(name, type)` pairs or [`ColumnSpec`]s.
78    ///
79    /// ```
80    /// use fv_compute::{FvType, SchemaSpec};
81    /// let sig = SchemaSpec::new([("id", FvType::Int64), ("amount", FvType::Float64)]);
82    /// assert!(sig.has_column("amount"));
83    /// ```
84    pub fn new(columns: impl IntoIterator<Item = impl Into<ColumnSpec>>) -> Self {
85        Self {
86            name: None,
87            columns: columns.into_iter().map(Into::into).collect(),
88        }
89    }
90
91    /// A named signature (required when a transform has more than one input).
92    pub fn named(name: impl Into<String>, columns: impl IntoIterator<Item = impl Into<ColumnSpec>>) -> Self {
93        Self {
94            name: Some(name.into()),
95            ..Self::new(columns)
96        }
97    }
98
99    /// Materialize as an Arrow `Schema`. All fields nullable=true (the plane is null-tolerant;
100    /// tighter nullability is a refinement).
101    pub fn to_arrow_schema(&self) -> Schema {
102        Schema::new(
103            self.columns
104                .iter()
105                .map(|c| Field::new(&c.name, c.dtype.to_arrow(), true))
106                .collect::<Vec<_>>(),
107        )
108    }
109
110    pub fn to_arrow_schema_ref(&self) -> Arc<Schema> {
111        Arc::new(self.to_arrow_schema())
112    }
113
114    /// Does this signature contain a column of the given name?
115    pub fn has_column(&self, name: &str) -> bool {
116        self.columns.iter().any(|c| c.name == name)
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn fvtype_maps_to_expected_arrow() {
126        assert_eq!(FvType::Int64.to_arrow(), DataType::Int64);
127        assert_eq!(FvType::Float64.to_arrow(), DataType::Float64);
128        assert_eq!(FvType::Utf8.to_arrow(), DataType::Utf8);
129        assert_eq!(
130            FvType::TimestampMs.to_arrow(),
131            DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into()))
132        );
133    }
134
135    #[test]
136    fn schema_spec_builds_arrow_schema() {
137        let s = SchemaSpec {
138            name: Some("in".into()),
139            columns: vec![
140                ColumnSpec {
141                    name: "id".into(),
142                    dtype: FvType::Int64,
143                },
144                ColumnSpec {
145                    name: "amount".into(),
146                    dtype: FvType::Float64,
147                },
148            ],
149        };
150        let arrow = s.to_arrow_schema();
151        assert_eq!(arrow.fields().len(), 2);
152        assert_eq!(arrow.field(0).name(), "id");
153        assert_eq!(arrow.field(1).data_type(), &DataType::Float64);
154        assert!(arrow.field(0).is_nullable());
155        assert!(s.has_column("amount") && !s.has_column("nope"));
156    }
157}