storage-kit 0.0.1

Portable database and transactional state-store Kit.
Documentation
use std::{collections::HashMap, sync::Arc};

use arrow_schema::{DataType, Field, Schema, SchemaRef};

use crate::{CollectionSchema, FieldSchema, FieldType};

use super::vector::{
    LANCE_UTF8_FIELD_TYPE_METADATA_KEY, metadata_field_type_tag, vector_element_type,
};

pub(crate) fn schema_to_arrow(schema: &CollectionSchema) -> SchemaRef {
    Arc::new(Schema::new(
        schema
            .fields
            .iter()
            .map(field_schema_to_arrow)
            .collect::<Vec<_>>(),
    ))
}

fn field_schema_to_arrow(field: &FieldSchema) -> Field {
    let data_type = field_type_to_arrow_data_type(&field.field_type);

    let mut metadata: HashMap<String, String> = field.metadata.clone().into_iter().collect();
    if let Some(tag) = metadata_field_type_tag(&field.field_type) {
        metadata.insert(LANCE_UTF8_FIELD_TYPE_METADATA_KEY.into(), tag.into());
    }

    Field::new(&field.name, data_type, field.nullable).with_metadata(metadata)
}

pub(crate) fn field_type_to_arrow_data_type(field_type: &FieldType) -> DataType {
    match field_type {
        FieldType::Bool => DataType::Boolean,
        FieldType::I64 => DataType::Int64,
        FieldType::U64 => DataType::UInt64,
        FieldType::F64 => DataType::Float64,
        FieldType::String => DataType::Utf8,
        FieldType::Json => DataType::Utf8,
        FieldType::Binary => DataType::Binary,
        FieldType::DenseVector { dim, element } => DataType::FixedSizeList(
            Arc::new(Field::new("item", vector_element_type(element), true)),
            *dim as i32,
        ),
    }
}