geoarrow_array/
wrap_array.rs

1use std::sync::Arc;
2
3use arrow_array::cast::AsArray;
4use arrow_array::{
5    Array, BinaryArray, BinaryViewArray, FixedSizeListArray, LargeBinaryArray, LargeStringArray,
6    ListArray, StringArray, StringViewArray, StructArray, UnionArray,
7};
8use arrow_schema::DataType;
9use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
10use geoarrow_schema::*;
11
12use crate::GeoArrowArray;
13use crate::array::*;
14
15/// Using a GeoArrow geometry type, wrap the provided storage array as a GeoArrow array.
16///
17/// This is a convenient way to convert from an Arrow array to a GeoArrow array when you have an
18/// extension type. You can also use the `TryFrom` implementations on each GeoArrow array type, but
19/// this may be easier to remember and find.
20pub trait WrapArray<Input> {
21    /// The output GeoArrow array type.
22    type Output: GeoArrowArray;
23
24    /// Wrap the given storage array as an GeoArrow array.
25    ///
26    /// This terminology comes from pyarrow/Arrow C++, where extension types similarly have a
27    /// [`wrap_array`](https://arrow.apache.org/docs/python/generated/pyarrow.ExtensionType.html#pyarrow.ExtensionType.wrap_array)
28    /// method.
29    fn wrap_array(&self, input: Input) -> GeoArrowResult<Self::Output>;
30}
31
32impl WrapArray<&StructArray> for PointType {
33    type Output = PointArray;
34
35    fn wrap_array(&self, input: &StructArray) -> GeoArrowResult<Self::Output> {
36        PointArray::try_from((input, self.clone()))
37    }
38}
39
40impl WrapArray<&FixedSizeListArray> for PointType {
41    type Output = PointArray;
42
43    fn wrap_array(&self, input: &FixedSizeListArray) -> GeoArrowResult<Self::Output> {
44        PointArray::try_from((input, self.clone()))
45    }
46}
47
48impl WrapArray<&dyn Array> for PointType {
49    type Output = PointArray;
50
51    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
52        PointArray::try_from((input, self.clone()))
53    }
54}
55
56impl WrapArray<&ListArray> for LineStringType {
57    type Output = LineStringArray;
58
59    fn wrap_array(&self, input: &ListArray) -> GeoArrowResult<Self::Output> {
60        LineStringArray::try_from((input, self.clone()))
61    }
62}
63
64impl WrapArray<&dyn Array> for LineStringType {
65    type Output = LineStringArray;
66
67    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
68        LineStringArray::try_from((input, self.clone()))
69    }
70}
71
72impl WrapArray<&ListArray> for PolygonType {
73    type Output = PolygonArray;
74
75    fn wrap_array(&self, input: &ListArray) -> GeoArrowResult<Self::Output> {
76        PolygonArray::try_from((input, self.clone()))
77    }
78}
79
80impl WrapArray<&dyn Array> for PolygonType {
81    type Output = PolygonArray;
82
83    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
84        PolygonArray::try_from((input, self.clone()))
85    }
86}
87
88impl WrapArray<&ListArray> for MultiPointType {
89    type Output = MultiPointArray;
90
91    fn wrap_array(&self, input: &ListArray) -> GeoArrowResult<Self::Output> {
92        MultiPointArray::try_from((input, self.clone()))
93    }
94}
95
96impl WrapArray<&dyn Array> for MultiPointType {
97    type Output = MultiPointArray;
98
99    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
100        MultiPointArray::try_from((input, self.clone()))
101    }
102}
103
104impl WrapArray<&ListArray> for MultiLineStringType {
105    type Output = MultiLineStringArray;
106
107    fn wrap_array(&self, input: &ListArray) -> GeoArrowResult<Self::Output> {
108        MultiLineStringArray::try_from((input, self.clone()))
109    }
110}
111
112impl WrapArray<&dyn Array> for MultiLineStringType {
113    type Output = MultiLineStringArray;
114
115    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
116        MultiLineStringArray::try_from((input, self.clone()))
117    }
118}
119
120impl WrapArray<&ListArray> for MultiPolygonType {
121    type Output = MultiPolygonArray;
122
123    fn wrap_array(&self, input: &ListArray) -> GeoArrowResult<Self::Output> {
124        MultiPolygonArray::try_from((input, self.clone()))
125    }
126}
127
128impl WrapArray<&dyn Array> for MultiPolygonType {
129    type Output = MultiPolygonArray;
130
131    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
132        MultiPolygonArray::try_from((input, self.clone()))
133    }
134}
135
136impl WrapArray<&StructArray> for BoxType {
137    type Output = RectArray;
138
139    fn wrap_array(&self, input: &StructArray) -> GeoArrowResult<Self::Output> {
140        RectArray::try_from((input, self.clone()))
141    }
142}
143
144impl WrapArray<&dyn Array> for BoxType {
145    type Output = RectArray;
146
147    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
148        RectArray::try_from((input, self.clone()))
149    }
150}
151
152impl WrapArray<&ListArray> for GeometryCollectionType {
153    type Output = GeometryCollectionArray;
154
155    fn wrap_array(&self, input: &ListArray) -> GeoArrowResult<Self::Output> {
156        GeometryCollectionArray::try_from((input, self.clone()))
157    }
158}
159
160impl WrapArray<&dyn Array> for GeometryCollectionType {
161    type Output = GeometryCollectionArray;
162
163    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
164        GeometryCollectionArray::try_from((input, self.clone()))
165    }
166}
167
168impl WrapArray<&UnionArray> for GeometryType {
169    type Output = GeometryArray;
170
171    fn wrap_array(&self, input: &UnionArray) -> GeoArrowResult<Self::Output> {
172        GeometryArray::try_from((input, self.clone()))
173    }
174}
175
176impl WrapArray<&dyn Array> for GeometryType {
177    type Output = GeometryArray;
178
179    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
180        GeometryArray::try_from((input, self.clone()))
181    }
182}
183
184impl WrapArray<&BinaryViewArray> for WkbType {
185    type Output = WkbViewArray;
186
187    fn wrap_array(&self, input: &BinaryViewArray) -> GeoArrowResult<Self::Output> {
188        Ok(WkbViewArray::from((input.clone(), self.clone())))
189    }
190}
191
192impl WrapArray<&BinaryArray> for WkbType {
193    type Output = WkbArray;
194
195    fn wrap_array(&self, input: &BinaryArray) -> GeoArrowResult<Self::Output> {
196        Ok(WkbArray::from((input.clone(), self.clone())))
197    }
198}
199
200impl WrapArray<&LargeBinaryArray> for WkbType {
201    type Output = LargeWkbArray;
202
203    fn wrap_array(&self, input: &LargeBinaryArray) -> GeoArrowResult<Self::Output> {
204        Ok(LargeWkbArray::from((input.clone(), self.clone())))
205    }
206}
207
208impl WrapArray<&StringViewArray> for WktType {
209    type Output = WktViewArray;
210
211    fn wrap_array(&self, input: &StringViewArray) -> GeoArrowResult<Self::Output> {
212        Ok(WktViewArray::from((input.clone(), self.clone())))
213    }
214}
215
216impl WrapArray<&dyn Array> for WkbType {
217    type Output = Arc<dyn GeoArrowArray>;
218
219    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
220        match input.data_type() {
221            DataType::BinaryView => Ok(Arc::new(WkbViewArray::from((
222                input.as_binary_view().clone(),
223                self.clone(),
224            )))),
225            DataType::Binary => Ok(Arc::new(WkbArray::from((
226                input.as_binary().clone(),
227                self.clone(),
228            )))),
229            DataType::LargeBinary => Ok(Arc::new(LargeWkbArray::from((
230                input.as_binary().clone(),
231                self.clone(),
232            )))),
233            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
234                "Unexpected DataType for WkbType: {dt:?}",
235            ))),
236        }
237    }
238}
239
240impl WrapArray<&StringArray> for WktType {
241    type Output = WktArray;
242
243    fn wrap_array(&self, input: &StringArray) -> GeoArrowResult<Self::Output> {
244        Ok(WktArray::from((input.clone(), self.clone())))
245    }
246}
247
248impl WrapArray<&LargeStringArray> for WktType {
249    type Output = LargeWktArray;
250
251    fn wrap_array(&self, input: &LargeStringArray) -> GeoArrowResult<Self::Output> {
252        Ok(LargeWktArray::from((input.clone(), self.clone())))
253    }
254}
255
256impl WrapArray<&dyn Array> for WktType {
257    type Output = Arc<dyn GeoArrowArray>;
258
259    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
260        match input.data_type() {
261            DataType::Utf8View => Ok(Arc::new(WktViewArray::from((
262                input.as_string_view().clone(),
263                self.clone(),
264            )))),
265            DataType::Utf8 => Ok(Arc::new(WktArray::from((
266                input.as_string().clone(),
267                self.clone(),
268            )))),
269            DataType::LargeUtf8 => Ok(Arc::new(LargeWktArray::from((
270                input.as_string().clone(),
271                self.clone(),
272            )))),
273            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
274                "Unexpected DataType for WktType: {dt:?}",
275            ))),
276        }
277    }
278}
279
280impl WrapArray<&dyn Array> for GeoArrowType {
281    type Output = Arc<dyn GeoArrowArray>;
282
283    fn wrap_array(&self, input: &dyn Array) -> GeoArrowResult<Self::Output> {
284        use GeoArrowType::*;
285
286        let result: Arc<dyn GeoArrowArray> = match self {
287            Point(t) => Arc::new(t.wrap_array(input)?),
288            LineString(t) => Arc::new(t.wrap_array(input)?),
289            Polygon(t) => Arc::new(t.wrap_array(input)?),
290            MultiPoint(t) => Arc::new(t.wrap_array(input)?),
291            MultiLineString(t) => Arc::new(t.wrap_array(input)?),
292            MultiPolygon(t) => Arc::new(t.wrap_array(input)?),
293            GeometryCollection(t) => Arc::new(t.wrap_array(input)?),
294            Rect(t) => Arc::new(t.wrap_array(input)?),
295            Geometry(t) => Arc::new(t.wrap_array(input)?),
296            Wkb(t) => Arc::new(WkbArray::try_from((input, t.clone()))?),
297            LargeWkb(t) => Arc::new(LargeWkbArray::try_from((input, t.clone()))?),
298            WkbView(t) => Arc::new(WkbViewArray::try_from((input, t.clone()))?),
299            Wkt(t) => Arc::new(WktArray::try_from((input, t.clone()))?),
300            LargeWkt(t) => Arc::new(LargeWktArray::try_from((input, t.clone()))?),
301            WktView(t) => Arc::new(WktViewArray::try_from((input, t.clone()))?),
302        };
303        Ok(result)
304    }
305}
geoarrow_array/wrap_array.rs

geoarrow_array/
wrap_array.rs
