From b510864410929d1e8f0cc47230f2d040ea2f9b3d Mon Sep 17 00:00:00 2001 From: Peter Hamilton Date: Wed, 11 Oct 2023 19:08:31 -0700 Subject: [PATCH] Add wrappers for postgres geometric types. For simple `Copy` types, just alias them and add impls for FromDatum and ToDatum. For more complex variable length types, create owned structs with impls for FromDatum and ToDatum as well as SqlTranslatable. --- README.md | 9 +- pgrx-tests/src/tests/geo_tests.rs | 95 +++++++-- pgrx/src/datum/geo.rs | 318 +++++++++++++++++++++++++++++- pgrx/src/datum/mod.rs | 2 +- 4 files changed, 402 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 1e284e97b..4e6dd7bcf 100644 --- a/README.md +++ b/README.md @@ -185,8 +185,13 @@ cargo pgrx init | `timestamp with time zone` | `pgrx::TimestampWithTimeZone` | | `anyarray` | `pgrx::AnyArray` | | `anyelement` | `pgrx::AnyElement` | -| `box` | `pgrx::pg_sys::BOX` | -| `point` | `pgrx::pg_sys::Point` | +| `box` | `pgrx::geo::Box` | +| `circle` | `pgrx::geo::Circle` | +| `line` | `pgrx::geo::Line` | +| `lseg` | `pgrx::geo::LineSegment` | +| `path` | `pgrx::geo::Path` | +| `point` | `pgrx::geo::Point` | +| `polygon` | `pgrx::geo::Polygon` | | `tid` | `pgrx::pg_sys::ItemPointerData` | | `cstring` | `&core::ffi::CStr` | | `inet` | `pgrx::Inet(String)` -- TODO: needs better support | diff --git a/pgrx-tests/src/tests/geo_tests.rs b/pgrx-tests/src/tests/geo_tests.rs index a8833b7a8..f41daea26 100644 --- a/pgrx-tests/src/tests/geo_tests.rs +++ b/pgrx-tests/src/tests/geo_tests.rs @@ -13,24 +13,97 @@ mod tests { #[allow(unused_imports)] use crate as pgrx_tests; + use pgrx::datum::geo::*; use pgrx::prelude::*; #[pg_test] - fn test_point_into_datum() -> spi::Result<()> { - let p = - Spi::get_one::("SELECT '42, 99'::point")?.expect("SPI result was null"); - assert_eq!(p.x, 42.0); - assert_eq!(p.y, 99.0); + fn test_point_datum() -> spi::Result<()> { + let p = Spi::get_one::("SELECT '42, 99'::point")?.expect("SPI result was null"); + assert_eq!(p, Point { x: 42.0, y: 99.0 }); + let p2 = Spi::get_one_with_args::( + "SELECT $1", + vec![(Point::type_oid().into(), p.into_datum())], + )? + .expect("SPI result was null"); + assert_eq!(p, p2); Ok(()) } #[pg_test] - fn test_box_into_datum() -> spi::Result<()> { - let b = Spi::get_one::("SELECT '1,2,3,4'::box")?.expect("SPI result was null"); - assert_eq!(b.high.x, 3.0); - assert_eq!(b.high.y, 4.0); - assert_eq!(b.low.x, 1.0); - assert_eq!(b.low.y, 2.0); + fn test_box_datum() -> spi::Result<()> { + let b = Spi::get_one::("SELECT '1,2,3,4'::box")?.expect("SPI result was null"); + assert_eq!(b, Box { high: Point { x: 3.0, y: 4.0 }, low: Point { x: 1.0, y: 2.0 } }); + let b2 = Spi::get_one_with_args::( + "SELECT $1", + vec![(Box::type_oid().into(), b.into_datum())], + )? + .expect("SPI result was null"); + assert_eq!(b, b2); + Ok(()) + } + + #[pg_test] + fn test_lseg_datum() -> spi::Result<()> { + let l = Spi::get_one::("SELECT '(1,2),(3,4)'::lseg")? + .expect("SPI result was null"); + assert_eq!(l.p, [Point { x: 1.0, y: 2.0 }, Point { x: 3.0, y: 4.0 }]); + let l2 = Spi::get_one_with_args::( + "SELECT $1", + vec![(LineSegment::type_oid().into(), l.into_datum())], + )? + .expect("SPI result was null"); + assert_eq!(l.p, l2.p); + Ok(()) + } + + #[pg_test] + fn test_path_datum() -> spi::Result<()> { + // Closed path + let p = Spi::get_one::("SELECT '((1,2),(3,4))'::path")?.expect("SPI result was null"); + assert_eq!(p.points(), [Point { x: 1.0, y: 2.0 }, Point { x: 3.0, y: 4.0 }]); + assert_eq!(p.closed(), true); + let p2 = Spi::get_one_with_args::( + "SELECT $1", + vec![(Path::type_oid().into(), p.clone().into_datum())], + )? + .expect("SPI result was null"); + assert_eq!(p.points(), p2.points()); + assert_eq!(p.closed(), p2.closed()); + + // Open path + let p = Spi::get_one::("SELECT '[(1,2),(3,4)]'::path")?.expect("SPI result was null"); + assert_eq!(p.points(), [Point { x: 1.0, y: 2.0 }, Point { x: 3.0, y: 4.0 }]); + assert_eq!(p.closed(), false); + let p2 = Spi::get_one_with_args::( + "SELECT $1", + vec![(Path::type_oid().into(), p.clone().into_datum())], + )? + .expect("SPI result was null"); + assert_eq!(p.points(), p2.points()); + assert_eq!(p.closed(), p2.closed()); + + Ok(()) + } + + #[pg_test] + fn test_polygon_datum() -> spi::Result<()> { + let p = Spi::get_one::("SELECT '((1,2),(3,4),(0,5))'::polygon")? + .expect("SPI result was null"); + assert_eq!( + p.points(), + [Point { x: 1.0, y: 2.0 }, Point { x: 3.0, y: 4.0 }, Point { x: 0.0, y: 5.0 }] + ); + assert_eq!( + p.boundbox(), + Box { high: Point { x: 3.0, y: 5.0 }, low: Point { x: 0.0, y: 2.0 } } + ); + let p2 = Spi::get_one_with_args::( + "SELECT $1", + vec![(Polygon::type_oid().into(), p.clone().into_datum())], + )? + .expect("SPI result was null"); + assert_eq!(p.points(), p2.points()); + assert_eq!(p.boundbox(), p2.boundbox()); Ok(()) } } diff --git a/pgrx/src/datum/geo.rs b/pgrx/src/datum/geo.rs index 85a017549..84038039c 100644 --- a/pgrx/src/datum/geo.rs +++ b/pgrx/src/datum/geo.rs @@ -7,9 +7,30 @@ //LICENSE All rights reserved. //LICENSE //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. -use crate::{pg_sys, FromDatum, IntoDatum, PgMemoryContexts}; -impl FromDatum for pg_sys::BOX { +//! This module contains implementations to make it easier to work with Postgres' geometric types +//! in Rust. +//! +//! Simple types that support `Copy` are re-exported from the `pg_sys` crate and have `IntoDatum` +//! and `FromDatum` trait implementations. Variable Length types have zero-copy and owned structs +//! that have `IntoDatum` and `FromDatum` trait implementations. +//! +//! See the [Postgres docs](https://www.postgresql.org/docs/current/datatype-geometric.html) for +//! more information on the geometric types. +use std::{mem, ptr}; + +use pgrx_pg_sys::pfree; +use pgrx_sql_entity_graph::metadata::{ + ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable, +}; + +use crate::{pg_sys, set_varsize_4b, FromDatum, IntoDatum, PgMemoryContexts}; + +// Copy types + +pub type Box = pg_sys::BOX; + +impl FromDatum for Box { unsafe fn from_polymorphic_datum( datum: pg_sys::Datum, is_null: bool, @@ -21,17 +42,17 @@ impl FromDatum for pg_sys::BOX { if is_null { None } else { - let the_box = datum.cast_mut_ptr::(); + let the_box = datum.cast_mut_ptr::(); Some(the_box.read()) } } } -impl IntoDatum for pg_sys::BOX { +impl IntoDatum for Box { fn into_datum(mut self) -> Option { unsafe { let ptr = PgMemoryContexts::CurrentMemoryContext - .copy_ptr_into(&mut self, std::mem::size_of::()); + .copy_ptr_into(&mut self, std::mem::size_of::()); Some(ptr.into()) } } @@ -41,7 +62,77 @@ impl IntoDatum for pg_sys::BOX { } } -impl FromDatum for pg_sys::Point { +pub type Line = pg_sys::LINE; + +impl FromDatum for Line { + unsafe fn from_polymorphic_datum( + datum: pg_sys::Datum, + is_null: bool, + _: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + let line = datum.cast_mut_ptr::(); + Some(line.read()) + } + } +} + +impl IntoDatum for Line { + fn into_datum(mut self) -> Option { + unsafe { + let ptr = PgMemoryContexts::CurrentMemoryContext + .copy_ptr_into(&mut self, std::mem::size_of::()); + Some(ptr.into()) + } + } + + fn type_oid() -> pg_sys::Oid { + pg_sys::LINEOID + } +} + +pub type LineSegment = pg_sys::LSEG; + +impl FromDatum for LineSegment { + unsafe fn from_polymorphic_datum( + datum: pg_sys::Datum, + is_null: bool, + _: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + let lseg = datum.cast_mut_ptr::(); + Some(lseg.read()) + } + } +} + +impl IntoDatum for LineSegment { + fn into_datum(mut self) -> Option { + unsafe { + let ptr = PgMemoryContexts::CurrentMemoryContext + .copy_ptr_into(&mut self, std::mem::size_of::()); + Some(ptr.into()) + } + } + + fn type_oid() -> pg_sys::Oid { + pg_sys::LSEGOID + } +} + +pub type Point = pg_sys::Point; + +impl FromDatum for Point { unsafe fn from_polymorphic_datum( datum: pg_sys::Datum, is_null: bool, @@ -59,11 +150,11 @@ impl FromDatum for pg_sys::Point { } } -impl IntoDatum for pg_sys::Point { +impl IntoDatum for Point { fn into_datum(mut self) -> Option { unsafe { let copy = PgMemoryContexts::CurrentMemoryContext - .copy_ptr_into(&mut self, std::mem::size_of::()); + .copy_ptr_into(&mut self, std::mem::size_of::()); Some(copy.into()) } } @@ -72,3 +163,214 @@ impl IntoDatum for pg_sys::Point { pg_sys::POINTOID } } + +pub type Circle = pg_sys::CIRCLE; + +impl FromDatum for Circle { + unsafe fn from_polymorphic_datum( + datum: pg_sys::Datum, + is_null: bool, + _: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + let circle = datum.cast_mut_ptr::(); + Some(circle.read()) + } + } +} + +impl IntoDatum for Circle { + fn into_datum(mut self) -> Option { + unsafe { + let ptr = PgMemoryContexts::CurrentMemoryContext + .copy_ptr_into(&mut self, std::mem::size_of::()); + Some(ptr.into()) + } + } + + fn type_oid() -> pg_sys::Oid { + pg_sys::CIRCLEOID + } +} + +// Variable Length Types + +/// An owned Postgres `path` type. +#[derive(Debug, Clone, Default)] +pub struct Path { + points: Vec, + closed: bool, +} + +impl Path { + pub fn new(points: Vec, closed: bool) -> Self { + Self { points, closed } + } + + pub fn points(&self) -> &[Point] { + &self.points + } + + pub fn closed(&self) -> bool { + self.closed + } +} + +impl IntoDatum for Path { + fn into_datum(self) -> Option { + let num_points = self.points.len(); + let reserve_size = num_points * mem::size_of::(); + // 16 bytes for the header (4 for varsize, 4 for npts, 4 for closed, 4 for padding) + let total_size = reserve_size + 16; + + unsafe { + let path = PgMemoryContexts::CurrentMemoryContext + .palloc(total_size as usize) + .cast::(); + set_varsize_4b(path.cast(), total_size as i32); + (*path).npts = num_points as i32; + (*path).closed = self.closed as i32; + let points = (*path).p.as_mut_slice(num_points); + for (i, point) in self.points.iter().enumerate() { + points[i] = *point; + } + Some(pg_sys::Datum::from(path)) + } + } + + fn type_oid() -> pg_sys::Oid { + pg_sys::PATHOID + } +} + +impl FromDatum for Path { + unsafe fn from_polymorphic_datum( + datum: pg_sys::Datum, + is_null: bool, + _: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + let data = pg_sys::pg_detoast_datum(datum.cast_mut_ptr()).cast::(); + let closed = (*data).closed != 0; + let points = ptr::addr_of!((*data).p); + let points = (*points).as_slice((*data).npts as usize).to_vec(); + if data != datum.cast_mut_ptr() { + pfree(data.cast()); + } + Some(Path { points, closed }) + } + } +} + +unsafe impl SqlTranslatable for Path { + fn argument_sql() -> Result { + Ok(SqlMapping::literal("path")) + } + + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::literal("path"))) + } +} + +/// An owned Postgres `polygon` type. +#[derive(Debug, Clone, Default)] +pub struct Polygon { + points: Vec, + boundbox: Box, +} + +unsafe impl SqlTranslatable for Polygon { + fn argument_sql() -> Result { + Ok(SqlMapping::literal("polygon")) + } + + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::literal("polygon"))) + } +} + +impl Polygon { + pub fn new(points: Vec) -> Self { + // We determine our boundbox by finding the min/max x/y values of all the points + // Fold over the points, updating the min/max as we go + let Some(boundbox) = points.iter().fold(None, |acc, point| { + let boundbox = acc.unwrap_or_else(|| Box { high: *point, low: *point }); + Some(Box { + high: Point { x: boundbox.high.x.max(point.x), y: boundbox.high.y.max(point.y) }, + low: Point { x: boundbox.low.x.min(point.x), y: boundbox.low.y.min(point.y) }, + }) + }) else { + return Self::default(); + }; + Self { points, boundbox } + } + + pub fn points(&self) -> &[Point] { + &self.points + } + + pub fn boundbox(&self) -> Box { + self.boundbox + } +} + +impl IntoDatum for Polygon { + fn into_datum(self) -> Option { + let num_points = self.points.len(); + let reserve_size = num_points * mem::size_of::(); + // 40 bytes for the header (4 for varsize, 4 for npts, 32 for boundbox) + let total_size = reserve_size + 40; + + unsafe { + let polygon = PgMemoryContexts::CurrentMemoryContext + .palloc(total_size as usize) + .cast::(); + set_varsize_4b(polygon.cast(), total_size as i32); + (*polygon).npts = num_points as i32; + let points = (*polygon).p.as_mut_slice(num_points); + for (i, point) in self.points.iter().enumerate() { + points[i] = *point; + } + (*polygon).boundbox = self.boundbox; + Some(pg_sys::Datum::from(polygon)) + } + } + + fn type_oid() -> pg_sys::Oid { + pg_sys::POLYGONOID + } +} + +impl FromDatum for Polygon { + unsafe fn from_polymorphic_datum( + datum: pg_sys::Datum, + is_null: bool, + _: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + let data = pg_sys::pg_detoast_datum(datum.cast_mut_ptr()).cast::(); + let points = ptr::addr_of!((*data).p); + let points = (*points).as_slice((*data).npts as usize).to_vec(); + let boundbox = (*data).boundbox; + if data != datum.cast_mut_ptr() { + pfree(data.cast()); + } + Some(Polygon { points, boundbox }) + } + } +} diff --git a/pgrx/src/datum/mod.rs b/pgrx/src/datum/mod.rs index 2a9ef2e3c..6e222e7cb 100644 --- a/pgrx/src/datum/mod.rs +++ b/pgrx/src/datum/mod.rs @@ -15,7 +15,7 @@ mod array; mod date; pub mod datetime_support; mod from; -mod geo; +pub mod geo; mod inet; mod internal; mod interval;