Skip to content

Commit a613a30

Browse files
kylebarronkylebarron
authored
WKB cleanup (#787)
1 parent bf73645 commit a613a30

File tree

13 files changed

+121
-111
lines changed

13 files changed

+121
-111
lines changed

‎src/algorithm/native/type_id.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,10 +119,9 @@ impl<O: OffsetSizeTrait> TypeIds for MixedGeometryArray<O, 2> {
119119
impl<O: OffsetSizeTrait> TypeIds for WKBArray<O> {
120120
fn get_type_ids(&self) -> Int16Array {
121121
let mut output_array = Int16Builder::with_capacity(self.len());
122-
123122
self.iter().for_each(|maybe_wkb| {
124123
output_array.append_option(maybe_wkb.map(|wkb| {
125-
let type_id = u32::from(wkb.get_wkb_geometry_type());
124+
let type_id = u32::from(wkb.wkb_type().unwrap());
126125
type_id.try_into().unwrap()
127126
}))
128127
});
@@ -133,7 +132,7 @@ impl<O: OffsetSizeTrait> TypeIds for WKBArray<O> {
133132
fn get_unique_type_ids(&self) -> HashSet<i16> {
134133
let mut values = HashSet::new();
135134
self.iter().flatten().for_each(|wkb| {
136-
let type_id = u32::from(wkb.get_wkb_geometry_type());
135+
let type_id = u32::from(wkb.wkb_type().unwrap());
137136
values.insert(type_id.try_into().unwrap());
138137
});
139138

‎src/io/wkb/common.rs‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,65 @@
1+
use std::io::Cursor;
2+
3+
use arrow_array::OffsetSizeTrait;
4+
use byteorder::{BigEndian, LittleEndian, ReadBytesExt};
15
use num_enum::{IntoPrimitive, TryFromPrimitive};
26

7+
use crate::error::{GeoArrowError, Result};
8+
use crate::scalar::WKB;
9+
10+
/// The various WKB types supported by this crate
311
#[derive(Clone, Copy, Debug, PartialEq, TryFromPrimitive, IntoPrimitive)]
412
#[repr(u32)]
513
pub enum WKBType {
14+
/// A WKB Point
615
Point = 1,
16+
/// A WKB LineString
717
LineString = 2,
18+
/// A WKB Polygon
819
Polygon = 3,
20+
/// A WKB MultiPoint
921
MultiPoint = 4,
22+
/// A WKB MultiLineString
1023
MultiLineString = 5,
24+
/// A WKB MultiPolygon
1125
MultiPolygon = 6,
26+
/// A WKB GeometryCollection
1227
GeometryCollection = 7,
28+
/// A WKB PointZ
1329
PointZ = 1001,
30+
/// A WKB LineStringZ
1431
LineStringZ = 1002,
32+
/// A WKB PolygonZ
1533
PolygonZ = 1003,
34+
/// A WKB MultiPointZ
1635
MultiPointZ = 1004,
36+
/// A WKB MultiLineStringZ
1737
MultiLineStringZ = 1005,
38+
/// A WKB MultiPolygonZ
1839
MultiPolygonZ = 1006,
40+
/// A WKB GeometryCollectionZ
1941
GeometryCollectionZ = 1007,
2042
}
43+
44+
impl WKBType {
45+
/// Construct from a byte slice representing a WKB geometry
46+
pub fn from_buffer(buf: &[u8]) -> Result<Self> {
47+
let mut reader = Cursor::new(buf);
48+
let byte_order = reader.read_u8().unwrap();
49+
let geometry_type = match byte_order {
50+
0 => reader.read_u32::<BigEndian>().unwrap(),
51+
1 => reader.read_u32::<LittleEndian>().unwrap(),
52+
_ => panic!("Unexpected byte order."),
53+
};
54+
Self::try_from_primitive(geometry_type)
55+
.map_err(|err| GeoArrowError::General(err.to_string()))
56+
}
57+
}
58+
59+
impl<'a, O: OffsetSizeTrait> TryFrom<WKB<'a, O>> for WKBType {
60+
type Error = GeoArrowError;
61+
62+
fn try_from(value: WKB<'a, O>) -> std::result::Result<Self, Self::Error> {
63+
Self::from_buffer(value.as_ref())
64+
}
65+
}

‎src/io/wkb/reader/coord.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ pub struct WKBCoord<'a> {
3636
}
3737

3838
impl<'a> WKBCoord<'a> {
39-
pub fn new(buf: &'a [u8], byte_order: Endianness, offset: u64, dim: Dimension) -> Self {
39+
pub(crate) fn new(buf: &'a [u8], byte_order: Endianness, offset: u64, dim: Dimension) -> Self {
4040
Self {
4141
buf,
4242
byte_order,

‎src/io/wkb/reader/geometry.rs‎

Lines changed: 40 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
11
use std::io::Cursor;
22

33
use arrow_array::OffsetSizeTrait;
4-
use byteorder::{BigEndian, LittleEndian, ReadBytesExt};
4+
use byteorder::ReadBytesExt;
55

66
use crate::datatypes::Dimension;
7-
use crate::geo_traits::{
8-
GeometryCollectionTrait, GeometryTrait, LineStringTrait, MultiLineStringTrait, MultiPointTrait,
9-
MultiPolygonTrait, PointTrait, PolygonTrait,
10-
};
7+
use crate::error::Result;
8+
use crate::geo_traits::GeometryTrait;
119
use crate::io::wkb::common::WKBType;
1210
use crate::io::wkb::reader::geometry_collection::WKBGeometryCollection;
1311
use crate::io::wkb::reader::rect::WKBRect;
@@ -18,97 +16,66 @@ use crate::io::wkb::reader::{
1816
use crate::scalar::WKB;
1917

2018
impl<'a, O: OffsetSizeTrait> WKB<'a, O> {
19+
/// Convert this WKB scalar to a [WKBGeometry]
20+
///
21+
/// This "prepares" the WKB input for constant-time coordinate access.
2122
pub fn to_wkb_object(&'a self) -> WKBGeometry<'a> {
22-
let buf = self.arr.value(self.geom_index);
23+
let buf = self.as_slice();
2324
let mut reader = Cursor::new(buf);
2425
let byte_order = reader.read_u8().unwrap();
25-
let geometry_type_u32 = match byte_order {
26-
0 => reader.read_u32::<BigEndian>().unwrap(),
27-
1 => reader.read_u32::<LittleEndian>().unwrap(),
28-
_ => panic!("Unexpected byte order."),
29-
};
30-
let geometry_type = WKBType::try_from(geometry_type_u32).unwrap();
26+
let wkb_type = self.wkb_type().unwrap();
27+
28+
use Dimension::*;
3129

32-
match geometry_type {
33-
WKBType::Point => {
34-
WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, Dimension::XY))
30+
match wkb_type {
31+
WKBType::Point => WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, XY)),
32+
WKBType::LineString => {
33+
WKBGeometry::LineString(WKBLineString::new(buf, byte_order.into(), 0, XY))
3534
}
36-
WKBType::LineString => WKBGeometry::LineString(WKBLineString::new(
37-
buf,
38-
byte_order.into(),
39-
0,
40-
Dimension::XY,
41-
)),
4235
WKBType::Polygon => {
43-
WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, Dimension::XY))
36+
WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, XY))
4437
}
4538
WKBType::MultiPoint => {
46-
WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), Dimension::XY))
39+
WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), XY))
40+
}
41+
WKBType::MultiLineString => {
42+
WKBGeometry::MultiLineString(WKBMultiLineString::new(buf, byte_order.into(), XY))
43+
}
44+
WKBType::MultiPolygon => {
45+
WKBGeometry::MultiPolygon(WKBMultiPolygon::new(buf, byte_order.into(), XY))
4746
}
48-
WKBType::MultiLineString => WKBGeometry::MultiLineString(WKBMultiLineString::new(
49-
buf,
50-
byte_order.into(),
51-
Dimension::XY,
52-
)),
53-
WKBType::MultiPolygon => WKBGeometry::MultiPolygon(WKBMultiPolygon::new(
54-
buf,
55-
byte_order.into(),
56-
Dimension::XY,
57-
)),
5847
WKBType::GeometryCollection => WKBGeometry::GeometryCollection(
59-
WKBGeometryCollection::new(buf, byte_order.into(), Dimension::XY),
48+
WKBGeometryCollection::new(buf, byte_order.into(), XY),
6049
),
61-
WKBType::PointZ => {
62-
WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, Dimension::XYZ))
50+
WKBType::PointZ => WKBGeometry::Point(WKBPoint::new(buf, byte_order.into(), 0, XYZ)),
51+
WKBType::LineStringZ => {
52+
WKBGeometry::LineString(WKBLineString::new(buf, byte_order.into(), 0, XYZ))
6353
}
64-
WKBType::LineStringZ => WKBGeometry::LineString(WKBLineString::new(
65-
buf,
66-
byte_order.into(),
67-
0,
68-
Dimension::XYZ,
69-
)),
7054
WKBType::PolygonZ => {
71-
WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, Dimension::XYZ))
55+
WKBGeometry::Polygon(WKBPolygon::new(buf, byte_order.into(), 0, XYZ))
7256
}
7357
WKBType::MultiPointZ => {
74-
WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), Dimension::XYZ))
58+
WKBGeometry::MultiPoint(WKBMultiPoint::new(buf, byte_order.into(), XYZ))
59+
}
60+
WKBType::MultiLineStringZ => {
61+
WKBGeometry::MultiLineString(WKBMultiLineString::new(buf, byte_order.into(), XYZ))
62+
}
63+
WKBType::MultiPolygonZ => {
64+
WKBGeometry::MultiPolygon(WKBMultiPolygon::new(buf, byte_order.into(), XYZ))
7565
}
76-
WKBType::MultiLineStringZ => WKBGeometry::MultiLineString(WKBMultiLineString::new(
77-
buf,
78-
byte_order.into(),
79-
Dimension::XYZ,
80-
)),
81-
WKBType::MultiPolygonZ => WKBGeometry::MultiPolygon(WKBMultiPolygon::new(
82-
buf,
83-
byte_order.into(),
84-
Dimension::XYZ,
85-
)),
8666
WKBType::GeometryCollectionZ => WKBGeometry::GeometryCollection(
87-
WKBGeometryCollection::new(buf, byte_order.into(), Dimension::XYZ),
67+
WKBGeometryCollection::new(buf, byte_order.into(), XYZ),
8868
),
8969
}
9070
}
9171

92-
pub fn get_wkb_geometry_type(&'a self) -> WKBType {
93-
let buf = self.arr.value(self.geom_index);
94-
let mut reader = Cursor::new(buf);
95-
let byte_order = reader.read_u8().unwrap();
96-
let geometry_type = match byte_order {
97-
0 => reader.read_u32::<BigEndian>().unwrap(),
98-
1 => reader.read_u32::<LittleEndian>().unwrap(),
99-
_ => panic!("Unexpected byte order."),
100-
};
101-
geometry_type.try_into().unwrap()
102-
}
103-
104-
pub fn to_wkb_line_string(&'a self) -> WKBLineString<'a> {
105-
match self.to_wkb_object() {
106-
WKBGeometry::LineString(geom) => geom,
107-
_ => panic!(),
108-
}
72+
/// Access the [WKBType] of this WKB object.
73+
pub fn wkb_type(&'a self) -> Result<WKBType> {
74+
WKBType::from_buffer(self.as_ref())
10975
}
11076
}
11177

78+
/// Endianness
11279
#[derive(Debug, Clone, Copy)]
11380
pub enum Endianness {
11481
BigEndian,
@@ -226,15 +193,6 @@ impl<'a> WKBGeometry<'a> {
226193
}
227194
}
228195

229-
impl<'a> From<WKBGeometry<'a>> for WKBLineString<'a> {
230-
fn from(value: WKBGeometry<'a>) -> Self {
231-
match value {
232-
WKBGeometry::LineString(geom) => geom,
233-
_ => panic!(),
234-
}
235-
}
236-
}
237-
238196
impl<'a> GeometryTrait for WKBGeometry<'a> {
239197
type T = f64;
240198
type Point<'b> = WKBPoint<'a> where Self: 'b;
@@ -247,15 +205,7 @@ impl<'a> GeometryTrait for WKBGeometry<'a> {
247205
type Rect<'b> = WKBRect<'a> where Self: 'b;
248206

249207
fn dim(&self) -> usize {
250-
match self {
251-
WKBGeometry::Point(g) => PointTrait::dim(g),
252-
WKBGeometry::LineString(g) => LineStringTrait::dim(g),
253-
WKBGeometry::Polygon(g) => PolygonTrait::dim(g),
254-
WKBGeometry::MultiPoint(g) => MultiPointTrait::dim(g),
255-
WKBGeometry::MultiLineString(g) => MultiLineStringTrait::dim(g),
256-
WKBGeometry::MultiPolygon(g) => MultiPolygonTrait::dim(g),
257-
WKBGeometry::GeometryCollection(g) => GeometryCollectionTrait::dim(g),
258-
}
208+
self.dimension().size()
259209
}
260210

261211
fn as_type(

‎src/io/wkb/reader/linearring.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ use crate::io::wkb::reader::geometry::Endianness;
99

1010
/// A linear ring in a WKB buffer.
1111
///
12+
/// This has been preprocessed, so access to any internal coordinate is `O(1)`.
13+
///
1214
/// See page 65 of <https://portal.ogc.org/files/?artifact_id=25355>.
1315
#[derive(Debug, Clone, Copy)]
1416
pub struct WKBLinearRing<'a> {

‎src/io/wkb/reader/linestring.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ use crate::io::wkb::reader::geometry::Endianness;
1010

1111
const HEADER_BYTES: u64 = 5;
1212

13+
/// A WKB LineString
14+
///
15+
/// This has been preprocessed, so access to any internal coordinate is `O(1)`.
1316
#[derive(Debug, Clone, Copy)]
1417
pub struct WKBLineString<'a> {
1518
buf: &'a [u8],

‎src/io/wkb/reader/multilinestring.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ use crate::io::wkb::reader::linestring::WKBLineString;
1010

1111
const HEADER_BYTES: u64 = 5;
1212

13+
/// A WKB MultiLineString
14+
///
15+
/// This has been preprocessed, so access to any internal coordinate is `O(1)`.
1316
#[derive(Debug, Clone)]
1417
pub struct WKBMultiLineString<'a> {
1518
/// A WKBLineString object for each of the internal line strings
@@ -19,7 +22,7 @@ pub struct WKBMultiLineString<'a> {
1922
}
2023

2124
impl<'a> WKBMultiLineString<'a> {
22-
pub fn new(buf: &'a [u8], byte_order: Endianness, dim: Dimension) -> Self {
25+
pub(crate) fn new(buf: &'a [u8], byte_order: Endianness, dim: Dimension) -> Self {
2326
let mut reader = Cursor::new(buf);
2427
reader.set_position(HEADER_BYTES);
2528
let num_line_strings = match byte_order {

‎src/io/wkb/reader/multipoint.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ use crate::geo_traits::MultiPointTrait;
88
use crate::io::wkb::reader::geometry::Endianness;
99
use crate::io::wkb::reader::point::WKBPoint;
1010

11+
/// A WKB MultiPoint
12+
///
13+
/// This has been preprocessed, so access to any internal coordinate is `O(1)`.
1114
#[derive(Debug, Clone, Copy)]
1215
pub struct WKBMultiPoint<'a> {
1316
buf: &'a [u8],
@@ -19,7 +22,7 @@ pub struct WKBMultiPoint<'a> {
1922
}
2023

2124
impl<'a> WKBMultiPoint<'a> {
22-
pub fn new(buf: &'a [u8], byte_order: Endianness, dim: Dimension) -> Self {
25+
pub(crate) fn new(buf: &'a [u8], byte_order: Endianness, dim: Dimension) -> Self {
2326
// TODO: assert WKB type?
2427
let mut reader = Cursor::new(buf);
2528
// Set reader to after 1-byte byteOrder and 4-byte wkbType

‎src/io/wkb/reader/multipolygon.rs‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,9 @@ use crate::io::wkb::reader::polygon::WKBPolygon;
1010

1111
const HEADER_BYTES: u64 = 5;
1212

13+
/// A WKB MultiPolygon
1314
#[derive(Debug, Clone)]
1415
pub struct WKBMultiPolygon<'a> {
15-
// buf: &'a [u8],
16-
// byte_order: Endianness,
17-
18-
// /// The number of polygons in this MultiPolygon
19-
// num_polygons: usize,
20-
21-
// /// The offset in the buffer where each WKBPolygon object begins
22-
// ///
23-
// /// The length of this vec must match the number of polygons
24-
// // polygon_offsets: Vec<usize>,
2516
/// A WKBPolygon object for each of the internal line strings
2617
wkb_polygons: Vec<WKBPolygon<'a>>,
2718

@@ -30,7 +21,7 @@ pub struct WKBMultiPolygon<'a> {
3021
}
3122

3223
impl<'a> WKBMultiPolygon<'a> {
33-
pub fn new(buf: &'a [u8], byte_order: Endianness, dim: Dimension) -> Self {
24+
pub(crate) fn new(buf: &'a [u8], byte_order: Endianness, dim: Dimension) -> Self {
3425
let mut reader = Cursor::new(buf);
3526
reader.set_position(HEADER_BYTES);
3627
let num_polygons = match byte_order {

‎src/io/wkb/reader/point.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ use crate::geo_traits::{CoordTrait, MultiPointTrait, PointTrait};
44
use crate::io::wkb::reader::coord::WKBCoord;
55
use crate::io::wkb::reader::geometry::Endianness;
66

7-
/// A 2D Point in WKB
7+
/// A WKB Point.
8+
///
9+
/// This has been preprocessed, so access to any internal coordinate is `O(1)`.
810
///
911
/// See page 66 of <https://portal.ogc.org/files/?artifact_id=25355>.
1012
#[derive(Debug, Clone, Copy)]

0 commit comments

Comments
 (0)