Skip to content

Commit c46b6bf

Browse files
committed
restructure files
1 parent fd739ab commit c46b6bf

File tree

7 files changed

+107
-32
lines changed

7 files changed

+107
-32
lines changed

rust/sedona-functions/src/st_to_sedona_spark.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ impl SedonaScalarKernel for STGeomToSedonaSpark {
1717
fn return_type(&self, args: &[SedonaType]) -> datafusion_common::Result<Option<SedonaType>> {
1818
let matcher = ArgMatcher::new(
1919
vec![ArgMatcher::is_geometry()],
20-
SedonaType::Arrow(DataType::BinaryView),
20+
SedonaType::Arrow(DataType::Binary),
2121
);
2222

2323
matcher.match_args(args)

rust/sedona-serde/src/deserialize.rs

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use crate::linestring::{parse_linestring, parse_multilinestring};
19-
use crate::point::{parse_multipoint, parse_point, write_empty_point};
20-
use crate::polygon::{parse_multipolygon, parse_polygon, write_empty_polygon};
18+
use crate::linestring::{deserialize_linestring, deserialize_multilinestring};
19+
use crate::point::{deserialize_multipoint, deserialize_point, deserialize_empty_point};
20+
use crate::polygon::{deserialize_multipolygon, deserialize_polygon, deserialize_empty_polygon};
2121
use crate::wkb::write_wkb_byte_order_marker;
2222
use arrow_array::builder::BinaryBuilder;
2323
use byteorder::{ByteOrder, LittleEndian, ReadBytesExt, WriteBytesExt};
@@ -36,10 +36,10 @@ pub fn deserialize(builder: &mut BinaryBuilder, bytes: &[u8]) -> datafusion_comm
3636

3737
let mut reader = Cursor::new(bytes);
3838

39-
parse_geometry::<LittleEndian, LittleEndian>(builder, &mut reader, bytes)
39+
deserialize_geometry::<LittleEndian, LittleEndian>(builder, &mut reader, bytes)
4040
}
4141

42-
pub fn parse_geometry<IN: ByteOrder, OUT: ByteOrder>(
42+
pub fn deserialize_geometry<IN: ByteOrder, OUT: ByteOrder>(
4343
builder: &mut BinaryBuilder,
4444
cursor: &mut Cursor<&[u8]>,
4545
bytes: &[u8],
@@ -50,6 +50,12 @@ pub fn parse_geometry<IN: ByteOrder, OUT: ByteOrder>(
5050

5151
let dimension = get_dimension((preamble_byte) >> 1);
5252

53+
if dimension != Dimension::XY {
54+
return Err(DataFusionError::Execution(
55+
"Only 2D geometries (XY) are supported".to_string(),
56+
));
57+
}
58+
5359
let _has_srid = (preamble_byte & 0x01) != 0;
5460

5561
cursor.set_position(cursor.position() + 3); // Skip 3 bytes
@@ -58,42 +64,42 @@ pub fn parse_geometry<IN: ByteOrder, OUT: ByteOrder>(
5864
1 => {
5965
let number_of_coordinates = cursor.read_u32::<IN>()?;
6066
if number_of_coordinates == 0 {
61-
write_empty_point::<OUT>(builder, dimension)?;
67+
deserialize_empty_point::<OUT>(builder, dimension)?;
6268
return Ok(());
6369
}
6470

65-
parse_point::<OUT>(builder, cursor, dimension)?;
71+
deserialize_point::<OUT>(builder, cursor, dimension)?;
6672
}
6773
2 => {
68-
parse_linestring::<IN, OUT>(builder, cursor, dimension)?;
74+
deserialize_linestring::<IN, OUT>(builder, cursor, dimension)?;
6975
}
7076
3 => {
7177
let mut meta_data_reader = Cursor::new(bytes);
7278

7379
let number_of_points = cursor.read_u32::<IN>()?;
7480
if number_of_points == 0 {
75-
write_empty_polygon::<OUT>(builder, dimension)?;
81+
deserialize_empty_polygon::<OUT>(builder, dimension)?;
7682

7783
return Ok(());
7884
}
7985

8086
let metadata_start_position = number_of_points * 8 * 2;
8187
meta_data_reader.set_position(cursor.position() + (metadata_start_position) as u64);
8288

83-
parse_polygon::<IN, OUT>(builder, cursor, &mut meta_data_reader, dimension)?;
89+
deserialize_polygon::<IN, OUT>(builder, cursor, &mut meta_data_reader, dimension)?;
8490
cursor.set_position(meta_data_reader.position());
8591
}
8692
4 => {
87-
parse_multipoint::<IN, OUT>(builder, cursor, dimension)?;
93+
deserialize_multipoint::<IN, OUT>(builder, cursor, dimension)?;
8894
}
8995
5 => {
9096
let mut meta_data_reader = Cursor::new(bytes);
91-
parse_multilinestring::<IN, OUT>(builder, cursor, &mut meta_data_reader, dimension)?;
97+
deserialize_multilinestring::<IN, OUT>(builder, cursor, &mut meta_data_reader, dimension)?;
9298
cursor.set_position(meta_data_reader.position());
9399
}
94100
6 => {
95101
let mut meta_data_reader = Cursor::new(bytes);
96-
parse_multipolygon::<IN, OUT>(builder, cursor, &mut meta_data_reader, dimension)?;
102+
deserialize_multipolygon::<IN, OUT>(builder, cursor, &mut meta_data_reader, dimension)?;
97103
cursor.set_position(meta_data_reader.position());
98104
}
99105
7 => {
@@ -104,7 +110,7 @@ pub fn parse_geometry<IN: ByteOrder, OUT: ByteOrder>(
104110
builder.write_u32::<OUT>(number_of_geometries)?;
105111

106112
for _i in 0..number_of_geometries {
107-
parse_geometry::<IN, OUT>(builder, cursor, bytes)?;
113+
deserialize_geometry::<IN, OUT>(builder, cursor, bytes)?;
108114
}
109115
}
110116
_ => {

rust/sedona-serde/src/linestring.rs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ fn get_linestring_marker(dimension: Dimension) -> u32 {
3030
}
3131
}
3232

33-
pub fn parse_linestring<IN: ByteOrder, OUT: ByteOrder>(
33+
pub fn deserialize_linestring<IN: ByteOrder, OUT: ByteOrder>(
3434
builder: &mut BinaryBuilder,
3535
cursor: &mut Cursor<&[u8]>,
3636
dimension: Dimension,
@@ -53,7 +53,7 @@ pub fn parse_linestring<IN: ByteOrder, OUT: ByteOrder>(
5353
Ok(())
5454
}
5555

56-
pub fn parse_multilinestring<IN: ByteOrder, OUT: ByteOrder>(
56+
pub fn deserialize_multilinestring<IN: ByteOrder, OUT: ByteOrder>(
5757
builder: &mut BinaryBuilder,
5858
cursor: &mut Cursor<&[u8]>,
5959
metadata_reader: &mut Cursor<&[u8]>,
@@ -145,7 +145,6 @@ pub fn serialize_multilinestring<OUT: ByteOrder>(
145145

146146
let _number_of_points = cursor.read_u32::<OUT>()?;
147147
total_number_of_points += _number_of_points;
148-
// number_of_points+= _number_of_points;
149148
metadata_cursor.write_u32::<OUT>(_number_of_points)?;
150149

151150
for _ in 0.._number_of_points * 2 {

rust/sedona-serde/src/point.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ fn get_byte_type_for_point(dimension: Dimension) -> u32 {
3333
}
3434
}
3535

36-
pub fn write_empty_point<OUT: ByteOrder>(
36+
pub fn deserialize_empty_point<OUT: ByteOrder>(
3737
builder: &mut BinaryBuilder,
3838
dimension: Dimension,
3939
) -> Result<()> {
@@ -46,7 +46,7 @@ pub fn write_empty_point<OUT: ByteOrder>(
4646
Ok(())
4747
}
4848

49-
pub fn parse_point<OUT: ByteOrder>(
49+
pub fn deserialize_point<OUT: ByteOrder>(
5050
builder: &mut BinaryBuilder,
5151
cursor: &mut Cursor<&[u8]>,
5252
dimension: Dimension,
@@ -65,7 +65,7 @@ pub fn parse_point<OUT: ByteOrder>(
6565
Ok(())
6666
}
6767

68-
pub fn parse_multipoint<IN: ByteOrder, OUT: ByteOrder>(
68+
pub fn deserialize_multipoint<IN: ByteOrder, OUT: ByteOrder>(
6969
builder: &mut BinaryBuilder,
7070
cursor: &mut Cursor<&[u8]>,
7171
dimension: Dimension,
@@ -90,7 +90,7 @@ pub fn parse_multipoint<IN: ByteOrder, OUT: ByteOrder>(
9090
builder.write_u32::<OUT>(number_of_points)?;
9191

9292
for _ in 0..number_of_points {
93-
parse_point::<OUT>(builder, cursor, dimension)?;
93+
deserialize_point::<OUT>(builder, cursor, dimension)?;
9494
}
9595

9696
Ok(())

rust/sedona-serde/src/polygon.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ pub(crate) fn get_polygon_marker(dimension: Dimension) -> u32 {
3030
}
3131
}
3232

33-
pub fn parse_polygon<IN: ByteOrder, OUT: ByteOrder>(
33+
pub fn deserialize_polygon<IN: ByteOrder, OUT: ByteOrder>(
3434
builder: &mut BinaryBuilder,
3535
cursor: &mut Cursor<&[u8]>,
3636
metadata_reader: &mut Cursor<&[u8]>,
@@ -57,7 +57,7 @@ pub fn parse_polygon<IN: ByteOrder, OUT: ByteOrder>(
5757
Ok(())
5858
}
5959

60-
pub(crate) fn parse_multipolygon<IN: ByteOrder, OUT: ByteOrder>(
60+
pub(crate) fn deserialize_multipolygon<IN: ByteOrder, OUT: ByteOrder>(
6161
builder: &mut BinaryBuilder,
6262
cursor: &mut Cursor<&[u8]>,
6363
metadata_reader: &mut Cursor<&[u8]>,
@@ -80,13 +80,13 @@ pub(crate) fn parse_multipolygon<IN: ByteOrder, OUT: ByteOrder>(
8080
builder.write_u32::<OUT>(number_of_geometries)?;
8181

8282
for _ in 0..number_of_geometries {
83-
parse_polygon::<IN, OUT>(builder, cursor, metadata_reader, dimension)?;
83+
deserialize_polygon::<IN, OUT>(builder, cursor, metadata_reader, dimension)?;
8484
}
8585

8686
Ok(())
8787
}
8888

89-
pub(crate) fn write_empty_polygon<OUT: ByteOrder>(
89+
pub(crate) fn deserialize_empty_polygon<OUT: ByteOrder>(
9090
builder: &mut BinaryBuilder,
9191
dimension: Dimension,
9292
) -> datafusion_common::Result<()> {

rust/sedona-serde/src/serialize.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,12 +59,12 @@ pub fn write_geometry<IN: ByteOrder, OUT: ByteOrder>(
5959
}
6060

6161
match wkb_byte {
62-
1 => return serialize_point::<LittleEndian>(builder, cursor),
63-
2 => return serialize_linestring::<LittleEndian>(builder, cursor),
64-
3 => return serialize_polygon::<LittleEndian>(builder, cursor),
65-
4 => return serialize_multipoint::<LittleEndian>(builder, cursor),
66-
5 => return serialize_multilinestring::<LittleEndian>(builder, cursor),
67-
6 => return serialize_multipolygon::<LittleEndian>(builder, cursor),
62+
1 => return serialize_point::<OUT>(builder, cursor),
63+
2 => return serialize_linestring::<OUT>(builder, cursor),
64+
3 => return serialize_polygon::<OUT>(builder, cursor),
65+
4 => return serialize_multipoint::<OUT>(builder, cursor),
66+
5 => return serialize_multilinestring::<OUT>(builder, cursor),
67+
6 => return serialize_multipolygon::<OUT>(builder, cursor),
6868
7 => {
6969
let number_of_geometries = cursor.read_u32::<IN>()?;
7070
builder.write_u32::<OUT>(number_of_geometries)?;

rust/sedona/src/context.rs

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -758,4 +758,74 @@ mod tests {
758758
.await
759759
.expect("should succeed because aws and gcs options were stripped");
760760
}
761+
762+
#[tokio::test]
763+
async fn test_sedona_spark_serde() -> Result<()> {
764+
let ctx = SedonaContext::new();
765+
766+
let geometry_data = ctx.sql(
767+
"SELECT
768+
ST_AsText(
769+
ST_GeomFromSedonaSpark(X'1200000001000000000000000000F03F000000000000F03F', 'EPSG:4326')
770+
) AS geom"
771+
)
772+
.await?
773+
.collect()
774+
.await?;
775+
776+
assert_batches_eq!(
777+
[
778+
"+------------+",
779+
"| geom |",
780+
"+------------+",
781+
"| POINT(1 1) |",
782+
"+------------+",
783+
],
784+
&geometry_data
785+
);
786+
787+
let srid_value = ctx.sql(
788+
"SELECT
789+
ST_SRID(
790+
ST_GeomFromSedonaSpark(X'1200000001000000000000000000F03F000000000000F03F', 'EPSG:4326')
791+
) AS srid"
792+
)
793+
.await?
794+
.collect()
795+
.await?;
796+
797+
assert_batches_eq!(
798+
[
799+
"+------+",
800+
"| srid |",
801+
"+------+",
802+
"| 4326 |",
803+
"+------+",
804+
],
805+
&srid_value
806+
);
807+
808+
let from_sedona_spark_and_reverse = ctx.sql(
809+
"SELECT
810+
ST_GeomToSedonaSpark(
811+
ST_GeomFromSedonaSpark(X'1200000001000000000000000000F03F000000000000F03F', 'EPSG:4326')
812+
) AS sedona_bytes"
813+
)
814+
.await?
815+
.collect()
816+
.await?;
817+
818+
assert_batches_eq!(
819+
[
820+
"+--------------------------------------------------+",
821+
"| sedona_bytes |",
822+
"+--------------------------------------------------+",
823+
"| 130010e601000000000000000000f03f000000000000f03f |",
824+
"+--------------------------------------------------+",
825+
],
826+
&from_sedona_spark_and_reverse
827+
);
828+
829+
Ok(())
830+
}
761831
}

0 commit comments

Comments
 (0)