Adding support for contrib/intarray style fast-allocated arrays

Author: mhov

pgrx-tests/src/tests/array_tests.rs

@@ -183,6 +183,11 @@ fn validate_cstring_array<'a>(
     Ok(true)
 }
 
+#[pg_extern]
+fn int_array_roundtrip(arr: Array<i32>) -> Array<i32> {
+    arr
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgrx::pg_schema]
 mod tests {
@@ -506,4 +511,85 @@ mod tests {
 
         Ok(())
     }
+
+    #[pg_test]
+    fn test_int_array_roundtrip_test() -> Result<(), Box<dyn std::error::Error>> {
+        let a = Spi::get_one::<Vec<i32>>("SELECT int_array_roundtrip(ARRAY[1, 2, 3, 4, 5])")?;
+
+        assert_eq!(a, Some(vec![1, 2, 3, 4, 5]));
+
+        Ok(())
+    }
+
+    #[pg_test]
+    fn test_array_new_from_slice() -> Result<(), Box<dyn std::error::Error>> {
+        let a = Spi::get_one::<Array<i8>>("SELECT ARRAY[1, 2, 3, 4, 5]::\"char\"[]")?
+            .expect("spi result was NULL");
+        let b = Array::<i8>::new_from_slice(&[1, 2, 3, 4, 5]).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, b.as_slice()?);
+
+        let a = Spi::get_one::<Array<i16>>("SELECT ARRAY[1, 2, 3, 4, 5]::smallint[]")?
+            .expect("spi result was NULL");
+        let b = Array::<i16>::new_from_slice(&[1, 2, 3, 4, 5]).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, b.as_slice()?);
+
+        let a = Spi::get_one::<Array<i32>>("SELECT ARRAY[1, 2, 3, 4, 5]::integer[]")?
+            .expect("spi result was NULL");
+        let b = Array::<i32>::new_from_slice(&[1, 2, 3, 4, 5]).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, b.as_slice()?);
+
+        let a = Spi::get_one::<Array<i64>>("SELECT ARRAY[1, 2, 3, 4, 5]::bigint[]")?
+            .expect("spi result was NULL");
+        let b = Array::<i64>::new_from_slice(&[1, 2, 3, 4, 5]).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, b.as_slice()?);
+
+        let a = Spi::get_one::<Array<f32>>("SELECT ARRAY[1.0, 2.0, 3.0, 4.0, 5.0]::float4[]")?
+            .expect("spi result was NULL");
+        let b = Array::<f32>::new_from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0])
+            .expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, b.as_slice()?);
+
+        let a = Spi::get_one::<Array<f64>>("SELECT ARRAY[1.0, 2.0, 3.0, 4.0, 5.0]::float8[]")?
+            .expect("spi result was NULL");
+        let b = Array::<f64>::new_from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0])
+            .expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, b.as_slice()?);
+
+        Ok(())
+    }
+
+    #[pg_test]
+    fn test_new_array_with_len() -> Result<(), Box<dyn std::error::Error>> {
+        let a = Array::<i8>::new_with_len(5).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, &[0, 0, 0, 0, 0]);
+
+        let a = Array::<i16>::new_with_len(5).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, &[0, 0, 0, 0, 0]);
+
+        let a = Array::<i32>::new_with_len(5).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, &[0, 0, 0, 0, 0]);
+
+        let a = Array::<i64>::new_with_len(5).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, &[0, 0, 0, 0, 0]);
+
+        let a = Array::<f32>::new_with_len(5).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, &[0.0, 0.0, 0.0, 0.0, 0.0]);
+
+        let a = Array::<f64>::new_with_len(5).expect("failed to create array");
+
+        assert_eq!(a.as_slice()?, &[0.0, 0.0, 0.0, 0.0, 0.0]);
+
+        Ok(())
+    }
 }

pgrx/src/array.rs

@@ -8,13 +8,14 @@
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
 #![allow(clippy::precedence)]
-use crate::datum::Array;
+use crate::datum::{Array, IntoDatum, UnboxDatum};
 use crate::toast::{Toast, Toasty};
-use crate::{layout, pg_sys, varlena};
+use crate::{layout, pg_sys, set_varsize_4b, varlena, PgMemoryContexts};
 use bitvec::ptr::{self as bitptr, BitPtr, BitPtrError, Mut};
 use bitvec::slice::BitSlice;
 use core::ptr::{self, NonNull};
 use core::slice;
+use pgrx_pg_sys::ArrayType;
 
 mod port;
 
@@ -373,10 +374,65 @@ impl RawArray {
         let ptr = self.ptr.as_ptr().cast::<u8>();
         ptr.wrapping_add(unsafe { varlena::varsize_any(ptr.cast()) })
     }
+
+    /// Slightly faster than new_array_type_with_len(0)
+    pub fn new_empty_array_type<T>() -> Result<RawArray, ArrayAllocError>
+    where
+        T: IntoDatum,
+        T: UnboxDatum,
+        T: Sized,
+    {
+        unsafe {
+            let array_type = pg_sys::construct_empty_array(T::type_oid());
+            let array_type =
+                NonNull::new(array_type).ok_or(ArrayAllocError::MemoryAllocationFailed)?;
+            Ok(RawArray::from_ptr(array_type))
+        }
+    }
+
+    /// Rustified version of new_intArrayType(int num) from https://github.com/postgres/postgres/blob/master/contrib/intarray/_int_tool.c#L219
+    pub fn new_array_type_with_len<T>(len: usize) -> Result<RawArray, ArrayAllocError>
+    where
+        T: IntoDatum,
+        T: UnboxDatum,
+        T: Sized,
+    {
+        if len == 0 {
+            return Self::new_empty_array_type::<T>();
+        }
+        let elem_size = std::mem::size_of::<T>();
+        let nbytes: usize = port::ARR_OVERHEAD_NONULLS(1) + elem_size * len;
+
+        unsafe {
+            let array_type = PgMemoryContexts::For(pg_sys::CurrentMemoryContext).palloc0(nbytes)
+                as *mut ArrayType;
+            if array_type.is_null() {
+                return Err(ArrayAllocError::MemoryAllocationFailed);
+            }
+            set_varsize_4b(array_type as *mut pg_sys::varlena, nbytes as i32);
+            (*array_type).ndim = 1;
+            (*array_type).dataoffset = 0; /* marker for no null bitmap */
+            (*array_type).elemtype = T::type_oid();
+
+            let ndims = port::ARR_DIMS(array_type);
+            *ndims = len as i32; // equivalent of ARR_DIMS(r)[0] = num;
+            let arr_lbound = port::ARR_LBOUND(array_type);
+            *arr_lbound = 1;
+
+            let array_type = NonNull::new_unchecked(array_type);
+            Ok(RawArray::from_ptr(array_type))
+        }
+    }
 }
 
 impl Toasty for RawArray {
     unsafe fn drop_toast(&mut self) {
         unsafe { pg_sys::pfree(self.ptr.as_ptr().cast()) }
     }
 }
+
+#[derive(thiserror::Error, Debug, Copy, Clone, Eq, PartialEq)]
+pub enum ArrayAllocError {
+    #[error("Failed to allocate memory for Array")]
+    MemoryAllocationFailed,
+}

pgrx/src/array/port.rs

@@ -126,3 +126,22 @@ pub(super) unsafe fn ARR_DATA_PTR(a: *mut pg_sys::ArrayType) -> *mut u8 {
 
     unsafe { a.cast::<u8>().add(ARR_DATA_OFFSET(a)) }
 }
+
+/// Returns a pointer to the lower bounds of the array.
+/// # Safety
+/// Does a field access, but doesn't deref out of bounds of ArrayType.  The caller asserts that
+/// `a` is a properly allocated [`pg_sys::ArrayType`]
+///
+/// [`pg_sys::ArrayType`] is typically allocated past its size, and its somewhere in that region
+/// that the returned pointer points, so don't attempt to `pfree` it.
+#[inline(always)]
+pub(super) unsafe fn ARR_LBOUND(a: *mut pg_sys::ArrayType) -> *mut i32 {
+    //  #define ARR_LBOUND(a) \
+    //        ((int *) (((char *) (a)) + sizeof(ArrayType) + \
+    //                  sizeof(int) * ARR_NDIM(a)))
+
+    a.cast::<u8>()
+        .add(std::mem::size_of::<pg_sys::ArrayType>())
+        .add(std::mem::size_of::<i32>() * ((*a).ndim as usize))
+        .cast::<i32>()
+}

pgrx/src/callconv.rs

@@ -20,9 +20,9 @@ use crate::datum::{Range, RangeSubType};
 use crate::heap_tuple::PgHeapTuple;
 use crate::layout::PassBy;
 use crate::nullable::Nullable;
-use crate::pg_sys;
 use crate::pgbox::*;
 use crate::rel::PgRelation;
+use crate::{pg_sys, Array};
 use crate::{PgBox, PgMemoryContexts};
 
 use core::marker::PhantomData;
@@ -600,6 +600,19 @@ where
     }
 }
 
+unsafe impl<'mcx, T: UnboxDatum> BoxRet for Array<'mcx, T>
+where
+    T: IntoDatum,
+{
+    #[inline]
+    unsafe fn box_into<'fcx>(self, fcinfo: &mut FcInfo<'fcx>) -> Datum<'fcx> {
+        match self.into_datum() {
+            Some(datum) => unsafe { fcinfo.return_raw_datum(datum) },
+            None => fcinfo.return_null(),
+        }
+    }
+}
+
 unsafe impl<T: Copy> BoxRet for PgVarlena<T> {
     unsafe fn box_into<'fcx>(self, fcinfo: &mut FcInfo<'fcx>) -> Datum<'fcx> {
         match self.into_datum() {