Add Dimensions support and another test

pgrx-tests/src/tests/array_borrowed.rs

@@ -72,11 +72,6 @@ fn borrow_optional_array_with_default(
 //     Json(json! { { "values": values } })
 // }
 
-// #[pg_extern]
-// fn borrow_serde_serialize_array_i32_deny_null(values: &FlatArray<'_, i32>) -> Json {
-//     Json(json! { { "values": values.iter_non_null() } })
-// }
-
 #[pg_extern]
 fn borrow_return_text_array() -> Vec<&'static str> {
     vec!["a", "b", "c", "d"]
@@ -115,11 +110,11 @@ fn borrow_display_get_arr_nullbitmap(arr: &FlatArray<'_, i32>) -> String {
     }
 }
 
-// #[pg_extern]
-// fn borrow_get_arr_ndim(arr: &FlatArray<'_, i32>) -> libc::c_int {
-//     // SAFETY: This is a valid FlatArrayType and it's just a field access.
-//     unsafe { RawArray::from_array(arr) }.unwrap().dims().len() as _
-// }
+#[pg_extern]
+fn borrow_get_arr_ndim(arr: &FlatArray<'_, i32>) -> libc::c_int {
+    // SAFETY: This is a valid FlatArrayType and it's just a field access.
+    arr.dims().len() as libc::c_int
+}
 
 // This deliberately iterates the FlatArray.
 // Because FlatArray::iter currently iterates the FlatArray as Datums, this is guaranteed to be "bug-free" regarding size.
@@ -265,13 +260,6 @@ mod tests {
     //     Ok(())
     // }
 
-    // #[pg_test(expected = "array contains NULL")]
-    // fn borrow_test_serde_serialize_array_i32_deny_null() -> Result<Option<Json>, pgrx::spi::Error> {
-    //     Spi::get_one::<Json>(
-    //         "SELECT borrow_serde_serialize_array_i32_deny_null(ARRAY[1, 2, 3, null, 4, 5])",
-    //     )
-    // }
-
     #[pg_test]
     fn borrow_test_return_text_array() {
         let rc = Spi::get_one::<bool>("SELECT ARRAY['a', 'b', 'c', 'd'] = return_text_array();");
@@ -336,19 +324,19 @@ mod tests {
         Ok(())
     }
 
-    // #[pg_test]
-    // fn borrow_test_get_arr_ndim() -> Result<(), pgrx::spi::Error> {
-    //     let ndim = Spi::get_one::<i32>("SELECT borrow_get_arr_ndim(ARRAY[1,2,3,4,5]::int[])")?
-    //         .expect("datum was null");
+    #[pg_test]
+    fn borrow_test_get_arr_ndim() -> Result<(), pgrx::spi::Error> {
+        let ndim = Spi::get_one::<i32>("SELECT borrow_get_arr_ndim(ARRAY[1,2,3,4,5]::int[])")?
+            .expect("datum was null");
 
-    //     assert_eq!(ndim, 1);
+        assert_eq!(ndim, 1);
 
-    //     let ndim = Spi::get_one::<i32>("SELECT borrow_get_arr_ndim('{{1,2,3},{4,5,6}}'::int[])")?
-    //         .expect("datum was null");
+        let ndim = Spi::get_one::<i32>("SELECT borrow_get_arr_ndim('{{1,2,3},{4,5,6}}'::int[])")?
+            .expect("datum was null");
 
-    //     assert_eq!(ndim, 2);
-    //     Ok(())
-    // }
+        assert_eq!(ndim, 2);
+        Ok(())
+    }
 
     #[pg_test]
     fn borrow_test_arr_to_vec() {
@@ -461,7 +449,7 @@ mod tests {
     //     Ok(())
     // }
 
-    // FIXME: needs to be type-subbed to a stringly type
+    // FIXME: needs to be type-subbed to a stringly type and SPI needs to make sense
     // #[pg_test]
     // fn borrow_test_text_array_as_vec_string() -> Result<(), Box<dyn std::error::Error>> {
     //     let a = Spi::get_one::<&FlatArray<'_, String>>(

pgrx/src/array.rs

@@ -21,7 +21,7 @@ use bitvec::ptr::{self as bitptr, BitPtr, BitPtrError, Const, Mut};
 use bitvec::slice::{self as bitslice, BitSlice};
 use core::marker::PhantomData;
 use core::ptr::{self, NonNull};
-use core::{mem, slice};
+use core::{ffi, mem, slice};
 
 mod port;
 
@@ -53,6 +53,17 @@ where
     pub fn count(&self) -> usize {
         self.as_raw().len()
     }
+
+    pub fn dims(&self) -> Dimensions {
+        // SAFETY: Validity of the ptr and ndim field was asserted upon obtaining the FlatArray ref,
+        // so can assume the dims ptr is also valid, allowing making the slice.
+        unsafe {
+            let ptr = self as *const Self as *const pg_sys::ArrayType;
+            let ndim = self.head.ndim as usize;
+            let dims = slice::from_raw_parts(port::ARR_DIMS(ptr.cast_mut()), ndim);
+            Dimensions { dims }
+        }
+    }
 }
 
 impl<'mcx, T> FlatArray<'mcx, T>
@@ -188,6 +199,17 @@ where
     }
 }
 
+#[derive(Clone)]
+pub struct Dimensions<'arr> {
+    dims: &'arr [ffi::c_int],
+}
+
+impl<'arr> Dimensions<'arr> {
+    pub fn len(&self) -> usize {
+        self.dims.len()
+    }
+}
+
 /// Iterator for arrays
 #[derive(Clone)]
 pub struct ArrayIter<'arr, T>