lance_namespace/
schema.rs

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

lance_namespace/
schema.rs
