impl Sum and Default for AnyNumeric (#1117)

In doing so, I realized that the `FromDatum` implementations for both `AnyNumeric` and `Vec<T>/Vec<Option<T>>` also needed to overload the `from_datum_in_memory_context()` function in order to properly return Arrays from Spi.
pgcentralfoundation · Apr 24, 2023 · 76aec0c · 76aec0c
1 parent abf0b4b
commit 76aec0c
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 14 deletions.
diff --git 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<dyn std::error::Error>> {
+        let result = Spi::get_one::<AnyNumeric>("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::<bool>("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<dyn std::error::Error>> {
+        let numbers = Spi::get_one::<Vec<AnyNumeric>>("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<dyn std::error::Error>> {
+        let numbers =
+            Spi::get_one::<Vec<Option<AnyNumeric>>>("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
@@ -480,6 +480,26 @@ impl<'a, T: FromDatum> FromDatum for Array<'a, T> {
             Some(Array::deconstruct_from(raw))
         }
     }
+
+    unsafe fn from_datum_in_memory_context(
+        mut memory_context: PgMemoryContexts,
+        datum: pg_sys::Datum,
+        is_null: bool,
+        typoid: pg_sys::Oid,
+    ) -> Option<Self>
+    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::<T>::from_polymorphic_datum(pg_sys::Datum::from(copy), false, typoid)
+            })
+        }
+    }
 }
 
 impl<T: FromDatum> FromDatum for Vec<T> {
@@ -492,15 +512,23 @@ impl<T: FromDatum> FromDatum for Vec<T> {
         if is_null {
             None
         } else {
-            let array = Array::<T>::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::<T>::from_polymorphic_datum(datum, is_null, typoid)
+                .map(|array| array.iter_deny_null().collect::<Vec<_>>())
         }
     }
+
+    unsafe fn from_datum_in_memory_context(
+        memory_context: PgMemoryContexts,
+        datum: pg_sys::Datum,
+        is_null: bool,
+        typoid: pg_sys::Oid,
+    ) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        Array::<T>::from_datum_in_memory_context(memory_context, datum, is_null, typoid)
+            .map(|array| array.iter_deny_null().collect::<Vec<_>>())
+    }
 }
 
 impl<T: FromDatum> FromDatum for Vec<Option<T>> {
@@ -510,12 +538,21 @@ impl<T: FromDatum> FromDatum for Vec<Option<T>> {
         is_null: bool,
         typoid: pg_sys::Oid,
     ) -> Option<Vec<Option<T>>> {
-        if is_null || datum.is_null() {
-            None
-        } else {
-            let array = Array::<T>::from_polymorphic_datum(datum, is_null, typoid).unwrap();
-            Some(array.iter().collect::<Vec<_>>())
-        }
+        Array::<T>::from_polymorphic_datum(datum, is_null, typoid)
+            .map(|array| array.iter().collect::<Vec<_>>())
+    }
+
+    unsafe fn from_datum_in_memory_context(
+        memory_context: PgMemoryContexts,
+        datum: pg_sys::Datum,
+        is_null: bool,
+        typoid: pg_sys::Oid,
+    ) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        Array::<T>::from_datum_in_memory_context(memory_context, datum, is_null, typoid)
+            .map(|array| array.iter().collect::<Vec<_>>())
     }
 }
 

diff --git 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<const P: u32, const S: u32> Display for Numeric<P, S> {
     }
 }
 
+impl Default for AnyNumeric {
+    fn default() -> Self {
+        0.into()
+    }
+}
+
+impl Sum for AnyNumeric {
+    fn sum<I: Iterator<Item = Self>>(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
@@ -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<Self>
+    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 {