Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fewer uses of custom GeometryType concept in MixedArrayBuilder #870

Merged
merged 1 commit into from
Nov 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 0 additions & 40 deletions rust/geoarrow/src/array/mixed/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -94,46 +94,6 @@ pub struct MixedGeometryArray {
pub(crate) slice_offset: usize,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) enum GeometryType {
Point = 1,
LineString = 2,
Polygon = 3,
MultiPoint = 4,
MultiLineString = 5,
MultiPolygon = 6,
GeometryCollection = 7,
}

impl GeometryType {
pub fn default_ordering(&self) -> i8 {
match self {
GeometryType::Point => 1,
GeometryType::LineString => 2,
GeometryType::Polygon => 3,
GeometryType::MultiPoint => 4,
GeometryType::MultiLineString => 5,
GeometryType::MultiPolygon => 6,
GeometryType::GeometryCollection => 7,
}
}
}

impl From<&String> for GeometryType {
fn from(value: &String) -> Self {
match value.as_str() {
"geoarrow.point" => GeometryType::Point,
"geoarrow.linestring" => GeometryType::LineString,
"geoarrow.polygon" => GeometryType::Polygon,
"geoarrow.multipoint" => GeometryType::MultiPoint,
"geoarrow.multilinestring" => GeometryType::MultiLineString,
"geoarrow.multipolygon" => GeometryType::MultiPolygon,
"geoarrow.geometrycollection" => GeometryType::GeometryCollection,
_ => panic!(),
}
}
}

impl MixedGeometryArray {
/// Create a new MixedGeometryArray from parts
///
Expand Down
42 changes: 32 additions & 10 deletions rust/geoarrow/src/array/mixed/builder.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use std::sync::Arc;

use crate::array::metadata::ArrayMetadata;
use crate::array::mixed::array::GeometryType;
use crate::array::mixed::MixedCapacity;
use crate::array::{
CoordType, LineStringBuilder, MixedGeometryArray, MultiLineStringBuilder, MultiPointBuilder,
Expand All @@ -17,7 +16,8 @@ use geo_traits::*;

pub(crate) const DEFAULT_PREFER_MULTI: bool = false;

/// The GeoArrow equivalent to a `Vec<Option<Geometry>>`: a mutable collection of Geometries.
/// The GeoArrow equivalent to a `Vec<Option<Geometry>>`: a mutable collection of Geometries, all
/// of which have the same dimension.
///
/// This currently has the caveat that these geometries must be a _primitive_ geometry type. This
/// does not currently support nested GeometryCollection objects.
Expand All @@ -32,6 +32,11 @@ pub(crate) const DEFAULT_PREFER_MULTI: bool = false;
pub struct MixedGeometryBuilder {
metadata: Arc<ArrayMetadata>,

/// The dimension of this builder.
///
/// All underlying arrays must contain a coordinate buffer of this same dimension.
dim: Dimension,

// Invariant: every item in `types` is `> 0 && < fields.len()`
types: Vec<i8>,

Expand Down Expand Up @@ -96,6 +101,7 @@ impl<'a> MixedGeometryBuilder {
// Don't store array metadata on child arrays
Self {
metadata,
dim,
types: vec![],
points: PointBuilder::with_capacity_and_options(
dim,
Expand Down Expand Up @@ -263,7 +269,10 @@ impl<'a> MixedGeometryBuilder {
#[inline]
pub(crate) fn add_point_type(&mut self) {
self.offsets.push(self.points.len().try_into().unwrap());
self.types.push(GeometryType::Point.default_ordering());
match self.dim {
Dimension::XY => self.types.push(1),
Dimension::XYZ => self.types.push(11),
}
}

/// Add a new LineString to the end of this array.
Expand Down Expand Up @@ -292,7 +301,10 @@ impl<'a> MixedGeometryBuilder {
pub(crate) fn add_line_string_type(&mut self) {
self.offsets
.push(self.line_strings.len().try_into().unwrap());
self.types.push(GeometryType::LineString.default_ordering());
match self.dim {
Dimension::XY => self.types.push(2),
Dimension::XYZ => self.types.push(12),
}
}

/// Add a new Polygon to the end of this array.
Expand All @@ -317,7 +329,10 @@ impl<'a> MixedGeometryBuilder {
#[inline]
pub(crate) fn add_polygon_type(&mut self) {
self.offsets.push(self.polygons.len().try_into().unwrap());
self.types.push(GeometryType::Polygon.default_ordering());
match self.dim {
Dimension::XY => self.types.push(3),
Dimension::XYZ => self.types.push(13),
}
}

/// Add a new MultiPoint to the end of this array.
Expand All @@ -338,7 +353,10 @@ impl<'a> MixedGeometryBuilder {
pub(crate) fn add_multi_point_type(&mut self) {
self.offsets
.push(self.multi_points.len().try_into().unwrap());
self.types.push(GeometryType::MultiPoint.default_ordering());
match self.dim {
Dimension::XY => self.types.push(4),
Dimension::XYZ => self.types.push(14),
}
}

/// Add a new MultiLineString to the end of this array.
Expand All @@ -359,8 +377,10 @@ impl<'a> MixedGeometryBuilder {
pub(crate) fn add_multi_line_string_type(&mut self) {
self.offsets
.push(self.multi_line_strings.len().try_into().unwrap());
self.types
.push(GeometryType::MultiLineString.default_ordering());
match self.dim {
Dimension::XY => self.types.push(5),
Dimension::XYZ => self.types.push(15),
}
}

/// Add a new MultiPolygon to the end of this array.
Expand All @@ -381,8 +401,10 @@ impl<'a> MixedGeometryBuilder {
pub(crate) fn add_multi_polygon_type(&mut self) {
self.offsets
.push(self.multi_polygons.len().try_into().unwrap());
self.types
.push(GeometryType::MultiPolygon.default_ordering());
match self.dim {
Dimension::XY => self.types.push(6),
Dimension::XYZ => self.types.push(16),
}
}

#[inline]
Expand Down
13 changes: 11 additions & 2 deletions rust/geoarrow/src/io/geozero/array/mixed.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use std::sync::Arc;

use crate::array::metadata::ArrayMetadata;
use crate::array::mixed::array::GeometryType;
use crate::array::{CoordType, MixedGeometryArray, MixedGeometryBuilder};
use crate::datatypes::Dimension;
use crate::io::geozero::scalar::process_geometry;
Expand Down Expand Up @@ -59,6 +58,17 @@ impl<T: GeozeroGeometry> ToMixedArray for T {
}
}

/// The current geometry type in which we're pushing coordinates.
#[derive(Debug, Clone, Copy, PartialEq)]
enum GeometryType {
Point = 1,
LineString = 2,
Polygon = 3,
MultiPoint = 4,
MultiLineString = 5,
MultiPolygon = 6,
}

/// A streaming builder for GeoArrow MixedGeometryArray.
///
/// This is useful in conjunction with [`geozero`] APIs because its coordinate stream requires the
Expand Down Expand Up @@ -140,7 +150,6 @@ impl GeomProcessor for MixedGeometryStreamBuilder {
GeometryType::MultiPoint => self.builder.multi_points.xy(x, y, idx),
GeometryType::MultiLineString => self.builder.multi_line_strings.xy(x, y, idx),
GeometryType::MultiPolygon => self.builder.multi_polygons.xy(x, y, idx),
GeometryType::GeometryCollection => todo!(),
}
}

Expand Down
Loading