diff --git a/pgrx-examples/aggregate/src/lib.rs b/pgrx-examples/aggregate/src/lib.rs index 264191eef..e996a84bd 100644 --- a/pgrx-examples/aggregate/src/lib.rs +++ b/pgrx-examples/aggregate/src/lib.rs @@ -120,7 +120,7 @@ impl Aggregate for IntegerAvgState { // unimplemented!() // } - // fn deserial(current: Self::State, _buf: Vec, _internal: PgBox, _fcinfo: pgrx::pg_sys::FunctionCallInfo) -> PgBox { + // fn deserial(_buf: Vec, _internal: pgrx::Internal, _fcinfo: pgrx::pg_sys::FunctionCallInfo) -> pgrx::Internal { // unimplemented!() // } diff --git a/pgrx-sql-entity-graph/src/aggregate/mod.rs b/pgrx-sql-entity-graph/src/aggregate/mod.rs index 13d2f78eb..092b2f555 100644 --- a/pgrx-sql-entity-graph/src/aggregate/mod.rs +++ b/pgrx-sql-entity-graph/src/aggregate/mod.rs @@ -453,11 +453,11 @@ impl PgAggregate { pg_externs.push(parse_quote! { #[allow(non_snake_case, clippy::too_many_arguments)] #pg_extern_attr - fn #fn_name(this: #type_state_without_self, buf: Vec, internal: ::pgrx::pgbox::PgBox<#type_state_without_self>, fcinfo: ::pgrx::pg_sys::FunctionCallInfo) -> ::pgrx::pgbox::PgBox<#type_state_without_self> { + fn #fn_name(buf: Vec, internal: ::pgrx::datum::Internal, fcinfo: ::pgrx::pg_sys::FunctionCallInfo) -> ::pgrx::datum::Internal { unsafe { <#target_path as ::pgrx::aggregate::Aggregate::<#generic_type>>::in_memory_context( fcinfo, - move |_context| <#target_path as ::pgrx::aggregate::Aggregate::<#generic_type>>::deserial(this, buf, internal, fcinfo) + move |_context| <#target_path as ::pgrx::aggregate::Aggregate::<#generic_type>>::deserial(buf, internal, fcinfo) ) } } @@ -465,7 +465,7 @@ impl PgAggregate { Some(fn_name) } else { item_impl.items.push(parse_quote! { - fn deserial(current: #type_state_without_self, _buf: Vec, _internal: ::pgrx::pgbox::PgBox<#type_state_without_self>, _fcinfo: ::pgrx::pg_sys::FunctionCallInfo) -> ::pgrx::pgbox::PgBox<#type_state_without_self> { + fn deserial(_buf: Vec, _internal: ::pgrx::datum::Internal, _fcinfo: ::pgrx::pg_sys::FunctionCallInfo) -> ::pgrx::datum::Internal { unimplemented!("Call to deserial on an aggregate which does not support it.") } }); @@ -1266,7 +1266,7 @@ mod tests { let tokens: ItemImpl = parse_quote! { #[pg_aggregate] impl Aggregate for DemoAgg { - type State = PgVarlena; + type State = Internal; type Args = i32; type OrderBy = i32; type MovingState = i32; @@ -1294,7 +1294,7 @@ mod tests { todo!() } - fn deserial(current: Self::State, _buf: Vec, _internal: PgBox) -> PgBox { + fn deserial(_buf: Vec, _internal: Internal) -> Internal { todo!() } @@ -1325,6 +1325,40 @@ mod tests { Ok(()) } + #[test] + fn agg_deserial_has_postgres_signature() -> Result<()> { + let tokens: ItemImpl = parse_quote! { + #[pg_aggregate] + impl Aggregate for DemoAgg { + type State = Internal; + type Args = i32; + + fn state(current: Self::State, arg: Self::Args) -> Self::State { + todo!() + } + + fn deserial(buf: Vec, internal: Internal) -> Internal { + todo!() + } + } + }; + + let agg = PgAggregate::new(tokens)?; + let deserial = agg + .0 + .pg_externs + .iter() + .find(|func| func.sig.ident == "demo_agg_demo_name_deserial") + .expect("deserial extern should be generated"); + + assert_eq!(deserial.sig.inputs.len(), 3); + let signature = deserial.sig.to_token_stream().to_string(); + assert!(signature.contains("buf : Vec < u8 >")); + assert!(signature.contains("internal : :: pgrx :: datum :: Internal")); + assert!(signature.ends_with("-> :: pgrx :: datum :: Internal")); + Ok(()) + } + #[test] fn agg_missing_required() -> Result<()> { // This is not valid as it is missing required types/consts. diff --git a/pgrx-unit-tests/src/tests/aggregate_tests.rs b/pgrx-unit-tests/src/tests/aggregate_tests.rs index a04f157d6..7d91306be 100644 --- a/pgrx-unit-tests/src/tests/aggregate_tests.rs +++ b/pgrx-unit-tests/src/tests/aggregate_tests.rs @@ -156,6 +156,59 @@ impl Aggregate for DemoPercentileDisc { } } +#[derive(AggregateName)] +struct ParallelInternalSum; + +#[pg_aggregate] +impl Aggregate for ParallelInternalSum { + type Args = i32; + type State = Internal; + type Finalize = i64; + + const PARALLEL: Option = Some(ParallelOption::Safe); + + fn state( + mut current: Self::State, + arg: Self::Args, + _fcinfo: pg_sys::FunctionCallInfo, + ) -> Self::State { + unsafe { + *current.get_or_insert_default::() += i64::from(arg); + } + current + } + + fn combine( + mut current: Self::State, + other: Self::State, + _fcinfo: pg_sys::FunctionCallInfo, + ) -> Self::State { + let other = unsafe { other.get::().copied().unwrap_or_default() }; + unsafe { + *current.get_or_insert_default::() += other; + } + current + } + + fn finalize( + current: Self::State, + _direct_args: Self::OrderedSetArgs, + _fcinfo: pg_sys::FunctionCallInfo, + ) -> Self::Finalize { + unsafe { current.get::().copied().unwrap_or_default() } + } + + fn serial(current: Self::State, _fcinfo: pg_sys::FunctionCallInfo) -> Vec { + let value = unsafe { current.get::().copied().unwrap_or_default() }; + value.to_le_bytes().to_vec() + } + + fn deserial(buf: Vec, _internal: Internal, _fcinfo: pg_sys::FunctionCallInfo) -> Internal { + let value = i64::from_le_bytes(buf.try_into().expect("serialized i64 state")); + Internal::new(value) + } +} + #[pgrx::pg_schema] mod demo_schema { use pgrx::PostgresType; @@ -337,6 +390,21 @@ mod tests { assert_eq!(retval, Ok(Some(5))); } + #[pg_test] + fn aggregate_parallel_internal_deserial() { + let signature = Spi::get_one::( + "SELECT oidvectortypes(proargtypes) FROM pg_proc \ + WHERE proname = 'parallel_internal_sum_parallel_internal_sum_deserial'", + ); + assert_eq!(signature, Ok(Some("bytea, internal".into()))); + + let value = Spi::get_one::( + "SELECT ParallelInternalSum(value) \ + FROM UNNEST(ARRAY [1, 1, 2]) AS value", + ); + assert_eq!(value, Ok(Some(4))); + } + #[pg_test] fn aggregate_demo_custom_state() { let retval = Spi::get_one::( diff --git a/pgrx/src/aggregate.rs b/pgrx/src/aggregate.rs index ce765b238..58cf75a13 100644 --- a/pgrx/src/aggregate.rs +++ b/pgrx/src/aggregate.rs @@ -272,10 +272,10 @@ CREATE AGGREGATE DemoSum ( */ +use crate::datum::Internal; use crate::error; use crate::memcxt::PgMemoryContexts; use crate::pg_sys::{AggCheckCallContext, CurrentMemoryContext, FunctionCallInfo, MemoryContext}; -use crate::pgbox::PgBox; pub use pgrx_sql_entity_graph::{FinalizeModify, ParallelOption}; @@ -416,13 +416,12 @@ where /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub. fn serial(current: Self::State, fcinfo: FunctionCallInfo) -> Vec; + /// Deserializes a `bytea` value into an `internal` aggregate state. + /// + /// PostgreSQL requires the second `internal` argument for type safety but does not use it. + /// /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub. - fn deserial( - current: Self::State, - _buf: Vec, - _internal: PgBox, - fcinfo: FunctionCallInfo, - ) -> PgBox; + fn deserial(_buf: Vec, _internal: Internal, fcinfo: FunctionCallInfo) -> Internal; /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub. fn moving_state(