1use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10
11#[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 TimestampMs,
25}
26
27impl FvType {
28 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#[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 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#[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 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 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 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 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}