Skip to main content

quadbin_schema/
type.rs

1use std::sync::Arc;
2
3use crate::Metadata;
4use arrow_schema::extension::ExtensionType;
5use arrow_schema::{ArrowError, DataType, Field};
6
7use crate::error::{QuadbinArrowError, QuadbinArrowResult};
8
9// macro_rules! define_basic_type {
10//     (
11//         $(#[$($attrss:meta)*])*
12//         $struct_name:ident
13//     ) => {
14//         $(#[$($attrss)*])*
15//         #[derive(Debug, Clone, PartialEq, Eq, Hash)]
16//         pub struct $struct_name {
17//             metadata: Arc<Metadata>,
18//         }
19
20//         impl $struct_name {
21//             /// Construct a new type from parts.
22//             pub fn new( metadata: Arc<Metadata>) -> Self {
23//                 Self {
24//                     metadata,
25//                 }
26//             }
27
28//             /// Change the underlying [`Metadata`]
29//             pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
30//                 Self { metadata, ..self }
31//             }
32
33//             /// Retrieve the underlying [`Metadata`]
34//             pub fn metadata(&self) -> &Arc<Metadata> {
35//                 &self.metadata
36//             }
37
38//             /// Convert this type to a [`Field`], retaining extension metadata.
39//             pub fn to_field<N: Into<String>>(&self, name: N, nullable: bool) -> Field {
40//                 Field::new(name, self.data_type(), nullable).with_extension_type(self.clone())
41//             }
42
43//             /// Extract into components
44//             pub fn into_inner(self) -> Arc<Metadata> {
45//                 ( self.metadata)
46//             }
47//         }
48//     };
49// }
50
51#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
52pub struct QuadbinType {
53    metadata: Arc<Metadata>,
54}
55
56impl QuadbinType {
57    /// Construct a new type from parts.
58    pub fn new(metadata: Arc<Metadata>) -> Self {
59        Self { metadata }
60    }
61
62    /// Change the underlying [`Metadata`]
63    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
64        Self { metadata }
65    }
66
67    /// Retrieve the underlying [`Metadata`]
68    pub fn metadata(&self) -> &Arc<Metadata> {
69        &self.metadata
70    }
71
72    pub fn to_data_type(&self) -> DataType {
73        DataType::UInt64
74    }
75
76    pub fn to_field<N: Into<String>>(&self, name: N, nullable: bool) -> Field {
77        Field::new(name, self.to_data_type(), nullable).with_extension_type(self.clone())
78    }
79
80    pub fn from_extension_field(field: &Field) -> QuadbinArrowResult<Self> {
81        let extension_name = field.extension_type_name().ok_or(QuadbinArrowError::InvalidGeoArrow(
82                "Expected rasterarrow extension metadata, but found none, and `require_rasterarrow_metadata` is `true`.".to_string(),
83            ))?;
84
85        let data_type = match extension_name {
86            QuadbinType::NAME => match field.data_type() {
87                DataType::UInt64 => field.try_extension_type()?,
88                _ => {
89                    return Err(QuadbinArrowError::InvalidGeoArrow(format!(
90                        "Expected binary type for a field with extension name 'geoarrow.wkb', got '{}'",
91                        field.data_type()
92                    )));
93                }
94            },
95            name => {
96                return Err(QuadbinArrowError::InvalidGeoArrow(format!(
97                    "Expected a GeoArrow extension name, got an Arrow extension type with name: '{name}'.",
98                )));
99            }
100        };
101        Ok(data_type)
102    }
103
104    pub fn from_arrow_field(field: &Field) -> QuadbinArrowResult<Self> {
105        if let Ok(geo_type) = Self::from_extension_field(field) {
106            Ok(geo_type)
107        } else {
108            let metadata = Arc::new(Metadata::try_from(field)?);
109            let data_type = match field.data_type() {
110                DataType::UInt64 => QuadbinType::new(metadata),
111                _ => 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())),
112             };
113            Ok(data_type)
114        }
115    }
116}
117
118impl ExtensionType for QuadbinType {
119    const NAME: &'static str = "rasterarrow.quadbin";
120
121    type Metadata = Arc<Metadata>;
122
123    fn metadata(&self) -> &Self::Metadata {
124        self.metadata()
125    }
126
127    fn serialize_metadata(&self) -> Option<String> {
128        self.metadata.serialize()
129    }
130
131    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
132        Ok(Arc::new(Metadata::deserialize(metadata)?))
133    }
134
135    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
136        match data_type {
137            DataType::UInt64 => Ok(()),
138            dt => Err(ArrowError::SchemaError(format!(
139                "Unexpected data type {dt}"
140            ))),
141        }
142    }
143
144    fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError> {
145        let raster = Self { metadata };
146        raster.supports_data_type(data_type)?;
147        Ok(raster)
148    }
149}
150
151impl From<QuadbinType> for DataType {
152    fn from(value: QuadbinType) -> Self {
153        value.to_data_type()
154    }
155}
156
157impl TryFrom<&Field> for QuadbinType {
158    type Error = QuadbinArrowError;
159
160    fn try_from(field: &Field) -> QuadbinArrowResult<Self> {
161        Self::from_extension_field(field)
162    }
163}