sedona_testing/
create.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17use std::{str::FromStr, sync::Arc};
18
19use arrow_array::{ArrayRef, BinaryArray, BinaryViewArray, StringViewArray, StructArray};
20use arrow_schema::{DataType, Field};
21use datafusion_common::ScalarValue;
22use datafusion_expr::ColumnarValue;
23use sedona_schema::datatypes::SedonaType;
24use wkb::{writer::WriteOptions, Endianness};
25use wkt::Wkt;
26
27/// Create a [`ColumnarValue`] array from a sequence of WKT literals
28///
29/// Panics on invalid WKT or unsupported data type.
30pub fn create_array_value(wkt_values: &[Option<&str>], data_type: &SedonaType) -> ColumnarValue {
31    ColumnarValue::Array(create_array_storage(wkt_values, data_type))
32}
33
34/// Create a [`ColumnarValue`] scalar from a WKT literal
35///
36/// Panics on invalid WKT or unsupported data type.
37pub fn create_scalar_value(wkt_value: Option<&str>, data_type: &SedonaType) -> ColumnarValue {
38    ColumnarValue::Scalar(create_scalar_storage(wkt_value, data_type))
39}
40
41/// Create a [`ScalarValue`] from a WKT literal
42///
43/// Panics on invalid WKT or unsupported data type.
44pub fn create_scalar(wkt_value: Option<&str>, data_type: &SedonaType) -> ScalarValue {
45    create_scalar_storage(wkt_value, data_type)
46}
47
48/// Create an [`ArrayRef`] from a sequence of WKT literals
49///
50/// Panics on invalid WKT or unsupported data type.
51pub fn create_array(wkt_values: &[Option<&str>], data_type: &SedonaType) -> ArrayRef {
52    create_array_storage(wkt_values, data_type)
53}
54
55/// Create the storage [`ArrayRef`] from a sequence of WKT literals
56///
57/// Panics on invalid WKT or unsupported data type. Supports Item CRS
58/// types; however, sets the CRS to Null
59pub fn create_array_storage(wkt_values: &[Option<&str>], data_type: &SedonaType) -> ArrayRef {
60    match data_type {
61        SedonaType::Wkb(_, _) => Arc::new(make_wkb_array::<BinaryArray>(wkt_values)),
62        SedonaType::WkbView(_, _) => Arc::new(make_wkb_array::<BinaryViewArray>(wkt_values)),
63        SedonaType::Arrow(DataType::Struct(fields))
64            if fields.iter().map(|f| f.name()).collect::<Vec<_>>() == vec!["item", "crs"] =>
65        {
66            let item_type = SedonaType::from_storage_field(&fields[0]).unwrap();
67            create_array_item_crs(wkt_values, (0..wkt_values.len()).map(|_| None), &item_type)
68        }
69        _ => panic!("create_array_storage not implemented for {data_type:?}"),
70    }
71}
72
73/// Create the storage [`ArrayRef`] from a sequence of WKT literals
74///
75/// Panics on invalid WKT or unsupported data type.
76pub fn create_array_item_crs<'a>(
77    wkt_values: &[Option<&str>],
78    crs: impl IntoIterator<Item = Option<&'a str>>,
79    item_type: &SedonaType,
80) -> ArrayRef {
81    let out_fields = vec![
82        item_type.to_storage_field("item", true).unwrap(),
83        Field::new("crs", DataType::Utf8View, true),
84    ];
85
86    let item_array = create_array_storage(wkt_values, item_type);
87    let crs_array = Arc::new(crs.into_iter().collect::<StringViewArray>());
88    let nulls = item_array.nulls().cloned();
89    Arc::new(StructArray::new(
90        out_fields.into(),
91        vec![item_array, crs_array],
92        nulls,
93    ))
94}
95
96/// Create the storage [`ScalarValue`] from a WKT literal
97///
98/// Panics on invalid WKT or unsupported data type. Item CRS values
99/// are created with a Null CRS: use [create_scalar_item_crs] to explicitly
100/// create Item CRS scalars with a specific CRS.
101pub fn create_scalar_storage(wkt_value: Option<&str>, data_type: &SedonaType) -> ScalarValue {
102    match data_type {
103        SedonaType::Wkb(_, _) => ScalarValue::Binary(wkt_value.map(make_wkb)),
104        SedonaType::WkbView(_, _) => ScalarValue::BinaryView(wkt_value.map(make_wkb)),
105        SedonaType::Arrow(DataType::Struct(fields))
106            if fields.iter().map(|f| f.name()).collect::<Vec<_>>() == vec!["item", "crs"] =>
107        {
108            let item_type = SedonaType::from_storage_field(&fields[0]).unwrap();
109            create_scalar_item_crs(wkt_value, None, &item_type)
110        }
111        _ => panic!("create_scalar_storage not implemented for {data_type:?}"),
112    }
113}
114
115/// Create a [`ScalarValue`] of an item_crs array from a WKT literal
116///
117/// Panics on invalid WKT or unsupported data type.
118pub fn create_scalar_item_crs(
119    wkt_value: Option<&str>,
120    crs: Option<&str>,
121    item_type: &SedonaType,
122) -> ScalarValue {
123    let out_fields = vec![
124        item_type.to_storage_field("item", true).unwrap(),
125        Field::new("crs", DataType::Utf8View, true),
126    ];
127
128    let storage_item = create_scalar_storage(wkt_value, item_type)
129        .to_array()
130        .unwrap();
131    let storage_crs = ScalarValue::Utf8View(crs.map(|item| item.to_string()))
132        .to_array()
133        .unwrap();
134    let nulls = storage_item.nulls().cloned();
135    let item_crs_array =
136        StructArray::try_new(out_fields.into(), vec![storage_item, storage_crs], nulls).unwrap();
137
138    ScalarValue::Struct(Arc::new(item_crs_array))
139}
140
141fn make_wkb_array<T>(wkt_values: &[Option<&str>]) -> T
142where
143    T: FromIterator<Option<Vec<u8>>>,
144{
145    wkt_values
146        .iter()
147        .map(|maybe_wkt| maybe_wkt.map(make_wkb))
148        .collect()
149}
150
151/// Create a WKB from a WKT string.
152pub fn make_wkb(wkt_value: &str) -> Vec<u8> {
153    let geom = Wkt::<f64>::from_str(wkt_value).unwrap();
154    let mut out: Vec<u8> = vec![];
155    wkb::writer::write_geometry(
156        &mut out,
157        &geom,
158        &WriteOptions {
159            endianness: Endianness::LittleEndian,
160        },
161    )
162    .unwrap();
163    out
164}
165
166#[cfg(test)]
167mod tests {
168    use arrow_schema::DataType;
169    use datafusion_common::cast::as_binary_array;
170    use sedona_schema::datatypes::{WKB_GEOMETRY, WKB_VIEW_GEOMETRY};
171
172    use super::*;
173
174    #[test]
175    fn scalars() {
176        let wkb_scalar = create_scalar_storage(Some("POINT (0 1)"), &WKB_GEOMETRY);
177        assert_eq!(&wkb_scalar.data_type(), WKB_GEOMETRY.storage_type());
178        assert!(create_scalar_storage(None, &WKB_GEOMETRY).is_null());
179
180        let wkb_view_scalar = create_scalar_storage(Some("POINT (0 1)"), &WKB_VIEW_GEOMETRY);
181        assert_eq!(
182            &wkb_view_scalar.data_type(),
183            WKB_VIEW_GEOMETRY.storage_type()
184        );
185        assert!(create_scalar_storage(None, &WKB_VIEW_GEOMETRY).is_null());
186    }
187
188    #[test]
189    #[should_panic(expected = "create_scalar_storage not implemented")]
190    fn scalar_storage_invalid() {
191        create_scalar_storage(Some("POINT (0 1)"), &SedonaType::Arrow(DataType::Null));
192    }
193
194    #[test]
195    fn arrays() {
196        let wkb_array = create_array_storage(
197            &[Some("POINT (0 1)"), None, Some("POINT (1 2)")],
198            &WKB_GEOMETRY,
199        );
200        assert_eq!(wkb_array.data_type(), WKB_GEOMETRY.storage_type());
201        assert_eq!(wkb_array.len(), 3);
202        let wkb_binary_array = as_binary_array(&wkb_array).unwrap();
203        assert_eq!(
204            wkb_binary_array
205                .iter()
206                .map(|maybe_item| maybe_item.is_some())
207                .collect::<Vec<bool>>(),
208            vec![true, false, true]
209        );
210
211        let wkb_array = create_array_storage(
212            &[Some("POINT (0 1)"), None, Some("POINT (1 2)")],
213            &WKB_VIEW_GEOMETRY,
214        );
215        assert_eq!(wkb_array.data_type(), WKB_VIEW_GEOMETRY.storage_type());
216    }
217
218    #[test]
219    #[should_panic(expected = "create_array_storage not implemented")]
220    fn array_storage_invalid() {
221        create_array_storage(&[Some("POINT (0 1)")], &SedonaType::Arrow(DataType::Null));
222    }
223}
sedona_testing/create.rs

sedona_testing/
create.rs
