Skip to main content

robin_sparkless_core/
schema.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize)]
4pub enum DataType {
5    String,
6    Integer,
7    Long,
8    Double,
9    Boolean,
10    Date,
11    Timestamp,
12    Binary,
13    Array(Box<DataType>),
14    Map(Box<DataType>, Box<DataType>),
15    Struct(Vec<StructField>),
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct StructField {
20    pub name: String,
21    pub data_type: DataType,
22    pub nullable: bool,
23}
24
25impl StructField {
26    pub fn new(name: String, data_type: DataType, nullable: bool) -> Self {
27        StructField {
28            name,
29            data_type,
30            nullable,
31        }
32    }
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct StructType {
37    fields: Vec<StructField>,
38}
39
40impl StructType {
41    pub fn new(fields: Vec<StructField>) -> Self {
42        StructType { fields }
43    }
44
45    pub fn fields(&self) -> &[StructField] {
46        &self.fields
47    }
48
49    /// Serialize the schema to a JSON string (array of field objects with name, data_type, nullable).
50    /// Useful for bindings that need to expose schema to the host without Polars types.
51    pub fn to_json(&self) -> Result<String, serde_json::Error> {
52        serde_json::to_string(self)
53    }
54
55    /// Serialize the schema to a pretty-printed JSON string.
56    pub fn to_json_pretty(&self) -> Result<String, serde_json::Error> {
57        serde_json::to_string_pretty(self)
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn test_struct_field_new() {
67        let field = StructField::new("age".to_string(), DataType::Integer, true);
68        assert_eq!(field.name, "age");
69        assert!(field.nullable);
70        assert!(matches!(field.data_type, DataType::Integer));
71    }
72
73    #[test]
74    fn test_struct_type_new() {
75        let fields = vec![
76            StructField::new("id".to_string(), DataType::Long, false),
77            StructField::new("name".to_string(), DataType::String, true),
78        ];
79        let schema = StructType::new(fields);
80        assert_eq!(schema.fields().len(), 2);
81        assert_eq!(schema.fields()[0].name, "id");
82        assert_eq!(schema.fields()[1].name, "name");
83    }
84
85    #[test]
86    fn test_struct_type_to_json() {
87        let fields = vec![
88            StructField::new("id".to_string(), DataType::Long, false),
89            StructField::new("name".to_string(), DataType::String, true),
90        ];
91        let schema = StructType::new(fields);
92        let json = schema.to_json().unwrap();
93        assert!(json.contains("\"name\":\"id\""));
94        assert!(json.contains("\"name\":\"name\""));
95        assert!(json.contains("\"data_type\""));
96        assert!(json.contains("\"nullable\""));
97        let _parsed: StructType = serde_json::from_str(&json).unwrap();
98        let pretty = schema.to_json_pretty().unwrap();
99        assert!(pretty.contains('\n'));
100    }
101
102    #[test]
103    fn test_empty_struct_type() {
104        let empty = StructType::new(vec![]);
105        assert!(empty.fields().is_empty());
106    }
107}