lance_namespace/
schema.rs

1//! Schema conversion utilities for Lance Namespace.
2//!
3//! This module provides functions to convert between JsonArrow schema representations
4//! and Arrow schema types.
5
6use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7use lance_namespace_reqwest_client::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema};
8
9use crate::namespace::{NamespaceError, Result};
10
11/// Convert JsonArrowSchema to Arrow Schema
12pub fn convert_json_arrow_schema(json_schema: &JsonArrowSchema) -> Result<ArrowSchema> {
13    let fields: Result<Vec<Field>> = json_schema
14        .fields
15        .iter()
16        .map(|f| convert_json_arrow_field(f))
17        .collect();
18
19    let metadata = json_schema
20        .metadata
21        .as_ref()
22        .map(|m| m.clone())
23        .unwrap_or_default();
24
25    Ok(ArrowSchema::new_with_metadata(fields?, metadata))
26}
27
28/// Convert JsonArrowField to Arrow Field
29pub fn convert_json_arrow_field(json_field: &JsonArrowField) -> Result<Field> {
30    let data_type = convert_json_arrow_type(&json_field.r#type)?;
31    let nullable = json_field.nullable;
32
33    Ok(Field::new(&json_field.name, data_type, nullable))
34}
35
36/// Convert JsonArrowDataType to Arrow DataType
37pub fn convert_json_arrow_type(json_type: &JsonArrowDataType) -> Result<DataType> {
38    let type_name = json_type.r#type.to_lowercase();
39
40    match type_name.as_str() {
41        "null" => Ok(DataType::Null),
42        "bool" | "boolean" => Ok(DataType::Boolean),
43        "int8" => Ok(DataType::Int8),
44        "uint8" => Ok(DataType::UInt8),
45        "int16" => Ok(DataType::Int16),
46        "uint16" => Ok(DataType::UInt16),
47        "int32" => Ok(DataType::Int32),
48        "uint32" => Ok(DataType::UInt32),
49        "int64" => Ok(DataType::Int64),
50        "uint64" => Ok(DataType::UInt64),
51        "float32" => Ok(DataType::Float32),
52        "float64" => Ok(DataType::Float64),
53        "utf8" => Ok(DataType::Utf8),
54        "binary" => Ok(DataType::Binary),
55        _ => Err(NamespaceError::Other(format!(
56            "Unsupported Arrow type: {}",
57            type_name
58        ))),
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use std::collections::HashMap;
66
67    #[test]
68    fn test_convert_basic_types() {
69        // Test int32
70        let int_type = JsonArrowDataType::new("int32".to_string());
71        let result = convert_json_arrow_type(&int_type).unwrap();
72        assert_eq!(result, DataType::Int32);
73
74        // Test utf8
75        let string_type = JsonArrowDataType::new("utf8".to_string());
76        let result = convert_json_arrow_type(&string_type).unwrap();
77        assert_eq!(result, DataType::Utf8);
78
79        // Test float64
80        let float_type = JsonArrowDataType::new("float64".to_string());
81        let result = convert_json_arrow_type(&float_type).unwrap();
82        assert_eq!(result, DataType::Float64);
83
84        // Test binary
85        let binary_type = JsonArrowDataType::new("binary".to_string());
86        let result = convert_json_arrow_type(&binary_type).unwrap();
87        assert_eq!(result, DataType::Binary);
88    }
89
90    #[test]
91    fn test_convert_field() {
92        let int_type = JsonArrowDataType::new("int32".to_string());
93        let field = JsonArrowField {
94            name: "test_field".to_string(),
95            r#type: Box::new(int_type),
96            nullable: false,
97            metadata: None,
98        };
99
100        let result = convert_json_arrow_field(&field).unwrap();
101        assert_eq!(result.name(), "test_field");
102        assert_eq!(result.data_type(), &DataType::Int32);
103        assert!(!result.is_nullable());
104    }
105
106    #[test]
107    fn test_convert_schema() {
108        let int_type = JsonArrowDataType::new("int32".to_string());
109        let string_type = JsonArrowDataType::new("utf8".to_string());
110
111        let id_field = JsonArrowField {
112            name: "id".to_string(),
113            r#type: Box::new(int_type),
114            nullable: false,
115            metadata: None,
116        };
117
118        let name_field = JsonArrowField {
119            name: "name".to_string(),
120            r#type: Box::new(string_type),
121            nullable: true,
122            metadata: None,
123        };
124
125        let mut metadata = HashMap::new();
126        metadata.insert("key".to_string(), "value".to_string());
127
128        let schema = JsonArrowSchema {
129            fields: vec![id_field, name_field],
130            metadata: Some(metadata.clone()),
131        };
132
133        let result = convert_json_arrow_schema(&schema).unwrap();
134        assert_eq!(result.fields().len(), 2);
135        assert_eq!(result.field(0).name(), "id");
136        assert_eq!(result.field(1).name(), "name");
137        assert_eq!(result.metadata(), &metadata);
138    }
139
140    #[test]
141    fn test_unsupported_type() {
142        let unsupported_type = JsonArrowDataType::new("unsupported".to_string());
143        let result = convert_json_arrow_type(&unsupported_type);
144        assert!(result.is_err());
145        assert!(result
146            .unwrap_err()
147            .to_string()
148            .contains("Unsupported Arrow type"));
149    }
150}
lance_namespace/schema.rs

lance_namespace/
schema.rs
