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#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
52pub struct QuadbinType {
53 metadata: Arc<Metadata>,
54}
55
56impl QuadbinType {
57 pub fn new(metadata: Arc<Metadata>) -> Self {
59 Self { metadata }
60 }
61
62 pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
64 Self { metadata }
65 }
66
67 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}