From 3c875c8aaf55215bd013e641cf391cba18a1bcba Mon Sep 17 00:00:00 2001 From: "Eric B. Ridge" Date: Thu, 20 Apr 2023 13:21:16 -0400 Subject: [PATCH] impl `Sum` and `Default` for `AnyNumeric` In doing so, I realized that the `FromDatum` implementations for both `AnyNumeric` and `Vec/Vec>` also needed to overload the `from_datum_in_memory_context()` function in order to properly return Arrays from Spi. So I'm mixing a little peanut butter with my jelly here, but now there's some tests around this that use the new `impl Sum for AnyNumeric`. --- pgrx-tests/src/tests/numeric_tests.rs | 27 +++++++++++ pgrx/src/datum/array.rs | 63 ++++++++++++++++++++----- pgrx/src/datum/numeric.rs | 17 +++++++ pgrx/src/datum/numeric_support/datum.rs | 22 ++++++++- 4 files changed, 115 insertions(+), 14 deletions(-) diff --git a/pgrx-tests/src/tests/numeric_tests.rs b/pgrx-tests/src/tests/numeric_tests.rs index 76957531f..56961a643 100644 --- a/pgrx-tests/src/tests/numeric_tests.rs +++ b/pgrx-tests/src/tests/numeric_tests.rs @@ -33,6 +33,14 @@ mod tests { AnyNumeric::try_from(std::u64::MAX).unwrap() } + #[pg_test] + fn select_a_numeric() -> Result<(), Box> { + let result = Spi::get_one::("SELECT 42::numeric")?.expect("SPI returned null"); + let expected: AnyNumeric = 42.into(); + assert_eq!(expected, result); + Ok(()) + } + #[pg_test] fn test_return_an_i32_numeric() { let result = Spi::get_one::("SELECT 32::numeric = tests.return_an_i32_numeric();"); @@ -197,4 +205,23 @@ mod tests { vec![(1, AnyNumeric::from(1)), (2, AnyNumeric::from(2)), (3, AnyNumeric::from(3)),] ) } + + #[pg_test] + fn test_anynumeric_sum() -> Result<(), Box> { + let numbers = Spi::get_one::>("SELECT ARRAY[1.0, 2.0, 3.0]::numeric[]")? + .expect("SPI result was null"); + let expected: AnyNumeric = 6.into(); + assert_eq!(expected, numbers.into_iter().sum()); + Ok(()) + } + + #[pg_test] + fn test_option_anynumeric_sum() -> Result<(), Box> { + let numbers = + Spi::get_one::>>("SELECT ARRAY[1.0, 2.0, 3.0]::numeric[]")? + .expect("SPI result was null"); + let expected: AnyNumeric = 6.into(); + assert_eq!(expected, numbers.into_iter().map(|n| n.unwrap_or_default()).sum()); + Ok(()) + } } diff --git a/pgrx/src/datum/array.rs b/pgrx/src/datum/array.rs index 9a27dc7cb..8813fef48 100644 --- a/pgrx/src/datum/array.rs +++ b/pgrx/src/datum/array.rs @@ -471,6 +471,26 @@ impl<'a, T: FromDatum> FromDatum for Array<'a, T> { Some(Array::deconstruct_from(ptr, raw, layout)) } } + + unsafe fn from_datum_in_memory_context( + mut memory_context: PgMemoryContexts, + datum: pg_sys::Datum, + is_null: bool, + typoid: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + memory_context.switch_to(|_| { + // copy the Datum into this MemoryContext, and then instantiate the Array wrapper + let copy = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr()); + Array::::from_polymorphic_datum(pg_sys::Datum::from(copy), false, typoid) + }) + } + } } impl FromDatum for Vec { @@ -483,15 +503,23 @@ impl FromDatum for Vec { if is_null { None } else { - let array = Array::::from_polymorphic_datum(datum, is_null, typoid).unwrap(); - let mut v = Vec::with_capacity(array.len()); - - for element in array.iter() { - v.push(element.expect("array element was NULL")) - } - Some(v) + Array::::from_polymorphic_datum(datum, is_null, typoid) + .map(|array| array.iter_deny_null().collect::>()) } } + + unsafe fn from_datum_in_memory_context( + memory_context: PgMemoryContexts, + datum: pg_sys::Datum, + is_null: bool, + typoid: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + Array::::from_datum_in_memory_context(memory_context, datum, is_null, typoid) + .map(|array| array.iter_deny_null().collect::>()) + } } impl FromDatum for Vec> { @@ -501,12 +529,21 @@ impl FromDatum for Vec> { is_null: bool, typoid: pg_sys::Oid, ) -> Option>> { - if is_null || datum.is_null() { - None - } else { - let array = Array::::from_polymorphic_datum(datum, is_null, typoid).unwrap(); - Some(array.iter().collect::>()) - } + Array::::from_polymorphic_datum(datum, is_null, typoid) + .map(|array| array.iter().collect::>()) + } + + unsafe fn from_datum_in_memory_context( + memory_context: PgMemoryContexts, + datum: pg_sys::Datum, + is_null: bool, + typoid: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + Array::::from_datum_in_memory_context(memory_context, datum, is_null, typoid) + .map(|array| array.iter().collect::>()) } } diff --git a/pgrx/src/datum/numeric.rs b/pgrx/src/datum/numeric.rs index e6f50e710..a3607b8df 100644 --- a/pgrx/src/datum/numeric.rs +++ b/pgrx/src/datum/numeric.rs @@ -10,6 +10,7 @@ Use of this source code is governed by the MIT license that can be found in the use core::ffi::CStr; use core::fmt::{Debug, Display, Formatter}; use std::fmt; +use std::iter::Sum; use crate::numeric_support::convert::from_primitive_helper; pub use crate::numeric_support::error::Error; @@ -87,6 +88,22 @@ impl Display for Numeric { } } +impl Default for AnyNumeric { + fn default() -> Self { + 0.into() + } +} + +impl Sum for AnyNumeric { + fn sum>(mut iter: I) -> Self { + let mut sum = iter.next().unwrap_or_default(); + for n in iter { + sum += n; + } + sum + } +} + #[inline(always)] pub(crate) const fn make_typmod(precision: u32, scale: u32) -> i32 { let (precision, scale) = (precision as i32, scale as i32); diff --git a/pgrx/src/datum/numeric_support/datum.rs b/pgrx/src/datum/numeric_support/datum.rs index 31cbc1cec..101ab553f 100644 --- a/pgrx/src/datum/numeric_support/datum.rs +++ b/pgrx/src/datum/numeric_support/datum.rs @@ -1,4 +1,4 @@ -use crate::{pg_sys, AnyNumeric, FromDatum, IntoDatum, Numeric}; +use crate::{pg_sys, AnyNumeric, FromDatum, IntoDatum, Numeric, PgMemoryContexts}; impl FromDatum for AnyNumeric { #[inline] @@ -21,6 +21,26 @@ impl FromDatum for AnyNumeric { Some(AnyNumeric { inner: numeric, need_pfree: is_copy }) } } + + unsafe fn from_datum_in_memory_context( + mut memory_context: PgMemoryContexts, + datum: pg_sys::Datum, + is_null: bool, + _typoid: pg_sys::Oid, + ) -> Option + where + Self: Sized, + { + if is_null { + None + } else { + memory_context.switch_to(|_| { + // copy the Datum into this MemoryContext and then create the AnyNumeric over that + let copy = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr()); + Some(AnyNumeric { inner: copy.cast(), need_pfree: true }) + }) + } + } } impl IntoDatum for AnyNumeric {