quadbin_schema/
datatype.rs

1use std::sync::Arc;
2
3use arrow_schema::extension::ExtensionType;
4use arrow_schema::{DataType, Field};
5
6use crate::Metadata;
7
8use crate::QuadbinType;
9use crate::error::{QuadbinArrowError, QuadbinArrowResult};
10
11#[derive(Debug, Clone, PartialEq, Eq, Hash)]
12pub enum QuadbinArrowType {
13    QuadbinU64(QuadbinType),
14}
15
16impl From<QuadbinArrowType> for DataType {
17    fn from(value: QuadbinArrowType) -> Self {
18        value.to_data_type()
19    }
20}
21
22impl QuadbinArrowType {
23    /// Returns the [Metadata] contained within this type.
24    pub fn metadata(&self) -> &Arc<Metadata> {
25        use QuadbinArrowType::*;
26        match self {
27            QuadbinU64(t) => t.metadata(),
28        }
29    }
30
31    pub fn to_data_type(&self) -> DataType {
32        use QuadbinArrowType::*;
33        match self {
34            QuadbinU64(_t) => DataType::UInt64,
35        }
36    }
37
38    pub fn to_field<N: Into<String>>(&self, name: N, nullable: bool) -> Field {
39        use QuadbinArrowType::*;
40        match self {
41            QuadbinU64(t) => {
42                Field::new(name, self.to_data_type(), nullable).with_extension_type(t.clone())
43            }
44        }
45    }
46
47    /// Applies the provided [Metadata] onto self.
48    pub fn with_metadata(self, meta: Arc<Metadata>) -> QuadbinArrowType {
49        use QuadbinArrowType::*;
50        match self {
51            QuadbinU64(t) => QuadbinU64(t.with_metadata(meta)),
52        }
53    }
54
55    pub fn from_extension_field(field: &Field) -> QuadbinArrowResult<Self> {
56        let extension_name = field.extension_type_name().ok_or(QuadbinArrowError::InvalidGeoArrow(
57                "Expected GeoArrow extension metadata, but found none, and `require_geoarrow_metadata` is `true`.".to_string(),
58            ))?;
59
60        use QuadbinArrowType::*;
61        let data_type = match extension_name {
62            QuadbinType::NAME => QuadbinU64(field.try_extension_type()?),
63            name => {
64                return Err(QuadbinArrowError::InvalidGeoArrow(format!(
65                    "Expected a Quadbin extension name, got an Arrow extension type with name: '{name}'.",
66                )));
67            }
68        };
69        Ok(data_type)
70    }
71
72    pub fn from_arrow_field(field: &Field) -> QuadbinArrowResult<Self> {
73        use QuadbinArrowType::*;
74        if let Ok(geo_type) = Self::from_extension_field(field) {
75            Ok(geo_type)
76        } else {
77            let metadata = Arc::new(Metadata::try_from(field)?);
78            let data_type = match field.data_type() {
79                DataType::UInt64 => QuadbinU64(QuadbinType::new(metadata)),
80                _ => return Err(QuadbinArrowError::InvalidGeoArrow("Only FixedSizeList, Struct, Binary, LargeBinary, BinaryView, String, LargeString, and StringView arrays are unambigously typed for a GeoArrow type and can be used without extension metadata.\nEnsure your array input has GeoArrow metadata.".to_string())),
81            };
82
83            Ok(data_type)
84        }
85    }
86}
87
88impl TryFrom<&Field> for QuadbinArrowType {
89    type Error = QuadbinArrowError;
90
91    fn try_from(field: &Field) -> QuadbinArrowResult<Self> {
92        Self::from_extension_field(field)
93    }
94}
quadbin_schema/datatype.rs

quadbin_schema/
datatype.rs
