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};
20use datafusion_common::ScalarValue;
21use datafusion_expr::ColumnarValue;
22use sedona_schema::datatypes::SedonaType;
23use wkb::{writer::WriteOptions, Endianness};
24use wkt::Wkt;
25
26/// Create a [`ColumnarValue`] array from a sequence of WKT literals
27///
28/// Panics on invalid WKT or unsupported data type.
29pub fn create_array_value(wkt_values: &[Option<&str>], data_type: &SedonaType) -> ColumnarValue {
30    ColumnarValue::Array(create_array_storage(wkt_values, data_type))
31}
32
33/// Create a [`ColumnarValue`] scalar from a WKT literal
34///
35/// Panics on invalid WKT or unsupported data type.
36pub fn create_scalar_value(wkt_value: Option<&str>, data_type: &SedonaType) -> ColumnarValue {
37    ColumnarValue::Scalar(create_scalar_storage(wkt_value, data_type))
38}
39
40/// Create a [`ScalarValue`] from a WKT literal
41///
42/// Panics on invalid WKT or unsupported data type.
43pub fn create_scalar(wkt_value: Option<&str>, data_type: &SedonaType) -> ScalarValue {
44    create_scalar_storage(wkt_value, data_type)
45}
46
47/// Create an [`ArrayRef`] from a sequence of WKT literals
48///
49/// Panics on invalid WKT or unsupported data type.
50pub fn create_array(wkt_values: &[Option<&str>], data_type: &SedonaType) -> ArrayRef {
51    create_array_storage(wkt_values, data_type)
52}
53
54/// Create the storage [`ArrayRef`] from a sequence of WKT literals
55///
56/// Panics on invalid WKT or unsupported data type.
57pub fn create_array_storage(wkt_values: &[Option<&str>], data_type: &SedonaType) -> ArrayRef {
58    match data_type {
59        SedonaType::Wkb(_, _) => Arc::new(make_wkb_array::<BinaryArray>(wkt_values)),
60        SedonaType::WkbView(_, _) => Arc::new(make_wkb_array::<BinaryViewArray>(wkt_values)),
61        _ => panic!("create_array_storage not implemented for {data_type:?}"),
62    }
63}
64
65/// Create the storage [`ScalarValue`] from a WKT literal
66///
67/// Panics on invalid WKT or unsupported data type.
68pub fn create_scalar_storage(wkt_value: Option<&str>, data_type: &SedonaType) -> ScalarValue {
69    match data_type {
70        SedonaType::Wkb(_, _) => ScalarValue::Binary(wkt_value.map(make_wkb)),
71        SedonaType::WkbView(_, _) => ScalarValue::BinaryView(wkt_value.map(make_wkb)),
72        _ => panic!("create_scalar_storage not implemented for {data_type:?}"),
73    }
74}
75
76fn make_wkb_array<T>(wkt_values: &[Option<&str>]) -> T
77where
78    T: FromIterator<Option<Vec<u8>>>,
79{
80    wkt_values
81        .iter()
82        .map(|maybe_wkt| maybe_wkt.map(make_wkb))
83        .collect()
84}
85
86/// Create a WKB from a WKT string.
87pub fn make_wkb(wkt_value: &str) -> Vec<u8> {
88    let geom = Wkt::<f64>::from_str(wkt_value).unwrap();
89    let mut out: Vec<u8> = vec![];
90    wkb::writer::write_geometry(
91        &mut out,
92        &geom,
93        &WriteOptions {
94            endianness: Endianness::LittleEndian,
95        },
96    )
97    .unwrap();
98    out
99}
100
101#[cfg(test)]
102mod tests {
103    use arrow_schema::DataType;
104    use datafusion_common::cast::as_binary_array;
105    use sedona_schema::datatypes::{WKB_GEOMETRY, WKB_VIEW_GEOMETRY};
106
107    use super::*;
108
109    #[test]
110    fn scalars() {
111        let wkb_scalar = create_scalar_storage(Some("POINT (0 1)"), &WKB_GEOMETRY);
112        assert_eq!(&wkb_scalar.data_type(), WKB_GEOMETRY.storage_type());
113        assert!(create_scalar_storage(None, &WKB_GEOMETRY).is_null());
114
115        let wkb_view_scalar = create_scalar_storage(Some("POINT (0 1)"), &WKB_VIEW_GEOMETRY);
116        assert_eq!(
117            &wkb_view_scalar.data_type(),
118            WKB_VIEW_GEOMETRY.storage_type()
119        );
120        assert!(create_scalar_storage(None, &WKB_VIEW_GEOMETRY).is_null());
121    }
122
123    #[test]
124    #[should_panic(expected = "create_scalar_storage not implemented")]
125    fn scalar_storage_invalid() {
126        create_scalar_storage(Some("POINT (0 1)"), &SedonaType::Arrow(DataType::Null));
127    }
128
129    #[test]
130    fn arrays() {
131        let wkb_array = create_array_storage(
132            &[Some("POINT (0 1)"), None, Some("POINT (1 2)")],
133            &WKB_GEOMETRY,
134        );
135        assert_eq!(wkb_array.data_type(), WKB_GEOMETRY.storage_type());
136        assert_eq!(wkb_array.len(), 3);
137        let wkb_binary_array = as_binary_array(&wkb_array).unwrap();
138        assert_eq!(
139            wkb_binary_array
140                .iter()
141                .map(|maybe_item| maybe_item.is_some())
142                .collect::<Vec<bool>>(),
143            vec![true, false, true]
144        );
145
146        let wkb_array = create_array_storage(
147            &[Some("POINT (0 1)"), None, Some("POINT (1 2)")],
148            &WKB_VIEW_GEOMETRY,
149        );
150        assert_eq!(wkb_array.data_type(), WKB_VIEW_GEOMETRY.storage_type());
151    }
152
153    #[test]
154    #[should_panic(expected = "create_array_storage not implemented")]
155    fn array_storage_invalid() {
156        create_array_storage(&[Some("POINT (0 1)")], &SedonaType::Arrow(DataType::Null));
157    }
158}
sedona_testing/create.rs

sedona_testing/
create.rs
