Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SRF functions need to "fetch" (call FromDatum) each argument while … #784

Merged
merged 2 commits into from
Oct 19, 2022
Merged

SRF functions need to "fetch" (call FromDatum) each argument while … #784

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b390e84
SRF functions need to "fetch" (call `FromDatum`) each argument while …
eeeebbbbrrrr Oct 19, 2022
715b4e1
only need to staticize the lifetime on the resolved type
eeeebbbbrrrr Oct 19, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions pgx-examples/strings/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,14 @@ fn split_set<'a>(input: &'a str, pattern: &'a str) -> SetOfIterator<'a, &'a str>
SetOfIterator::new(input.split_terminator(pattern).into_iter())
}

#[pg_extern]
fn split_table<'a>(
input: &'a str,
pattern: &'a str,
) -> TableIterator<'a, (name!(i, i32), name!(s, &'a str))> {
TableIterator::new(input.split_terminator(pattern).enumerate().map(|(i, s)| (i as i32, s)))
}

#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
Expand Down
49 changes: 49 additions & 0 deletions pgx-tests/src/tests/srf_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,19 @@ fn return_none_setof_iterator() -> Option<SetOfIterator<'static, i32>> {
}
}

#[pg_extern]
fn split_set_with_borrow<'a>(input: &'a str, pattern: &'a str) -> SetOfIterator<'a, &'a str> {
SetOfIterator::new(input.split_terminator(pattern))
}

#[pg_extern]
fn split_table_with_borrow<'a>(
input: &'a str,
pattern: &'a str,
) -> TableIterator<'a, (name!(i, i32), name!(s, &'a str))> {
TableIterator::new(input.split_terminator(pattern).enumerate().map(|(i, s)| (i as i32, s)))
}

#[cfg(any(test, feature = "pg_test"))]
#[pgx::pg_schema]
mod tests {
Expand Down Expand Up @@ -161,4 +174,40 @@ mod tests {

assert_eq!(cnt, Some(0))
}

#[pg_test]
fn test_srf_setof_datum_detoasting_with_borrow() {
let cnt = Spi::connect(|mut client| {
// build up a table with one large column that Postgres will be forced to TOAST
client.update("CREATE TABLE test_srf_datum_detoasting AS SELECT array_to_string(array_agg(g),' ') s FROM (SELECT 'a' g FROM generate_series(1, 1000000)) x;", None, None);

// and make sure we can use the DETOASTED value with our SRF function
let table = client.select(
"SELECT split_set_with_borrow(s, ' ') FROM test_srf_datum_detoasting",
None,
None,
);

Ok(Some(table.len() as i64))
});
assert_eq!(cnt, Some(1000000))
}

#[pg_test]
fn test_srf_table_datum_detoasting_with_borrow() {
let cnt = Spi::connect(|mut client| {
// build up a table with one large column that Postgres will be forced to TOAST
client.update("CREATE TABLE test_srf_datum_detoasting AS SELECT array_to_string(array_agg(g),' ') s FROM (SELECT 'a' g FROM generate_series(1, 1000000)) x;", None, None);

// and make sure we can use the DETOASTED value with our SRF function
let table = client.select(
"SELECT split_table_with_borrow(s, ' ') FROM test_srf_datum_detoasting",
None,
None,
);

Ok(Some(table.len() as i64))
});
assert_eq!(cnt, Some(1000000))
}
}
20 changes: 18 additions & 2 deletions pgx-utils/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ All rights reserved.
Use of this source code is governed by the MIT license that can be found in the LICENSE file.
*/

use crate::sql_entity_graph::PositioningRef;
use proc_macro2::{TokenStream, TokenTree};
use crate::sql_entity_graph::{NameMacro, PositioningRef};
use proc_macro2::{Ident, TokenStream, TokenTree};
use quote::{format_ident, quote, ToTokens, TokenStreamExt};
use std::collections::HashSet;
use syn::{GenericArgument, PathArguments, Type, TypeParamBound};
Expand Down Expand Up @@ -408,6 +408,22 @@ pub fn staticize_lifetimes(value: &mut syn::Type) {
}
}

syn::Type::Macro(type_macro) => {
let mac = &type_macro.mac;
let archetype = mac.path.segments.last().unwrap();
match archetype.ident.to_string().as_str() {
"name" => {
let mut out: NameMacro = mac.parse_body().expect("name!() in wrong format");
staticize_lifetimes(&mut out.used_ty.resolved_ty);

// rewrite the name!() macro so that it has a static lifetime, if any
let ident = syn::parse_str::<Ident>(&out.ident).unwrap();
let ty = out.used_ty.resolved_ty;
type_macro.mac = syn::parse_quote! {name!(#ident, #ty)};
Comment on lines +415 to +422
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

#[pg_extern]
fn split_table_with_borrow<'a>(
    input: &'a str,
    pattern: &'a str,
) -> TableIterator<'a, (name!(i, i32), name!(s, &'a str))>

So in this case, the return lifetimes we are expecting will be

) -> TableIterator<'static, (name!(i, i32), name!(s, &'static str))>

?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm. I think this can be done in a better way in the future, but I can't find the solution immediately and we're staticizing other lifetimes, so I think this is acceptable for the current moment.

}
_ => {}
}
}
_ => {}
}
}
Expand Down
51 changes: 30 additions & 21 deletions pgx-utils/src/sql_entity_graph/pg_extern/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -440,22 +440,15 @@ impl PgExtern {
optional,
} => {
let result_ident = syn::Ident::new("result", self.func.sig.span());
let funcctx_ident = syn::Ident::new("funcctx", self.func.sig.span());
let retval_ty_resolved = retval_ty.original_ty;
let result_handler = if optional {
// don't need unsafe annotations because of the larger unsafe block coming up
quote_spanned! { self.func.sig.span() =>
let #result_ident = match pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| { #func_name(#(#arg_pats),*) }) {
Some(result) => result,
None => {
pgx::srf_return_done(#fcinfo_ident, &mut funcctx);
return pgx::pg_return_null(#fcinfo_ident)
}
};
#func_name(#(#arg_pats),*)
}
} else {
quote_spanned! { self.func.sig.span() =>
let #result_ident = pgx::PgMemoryContexts::For(#funcctx_ident.multi_call_memory_ctx).switch_to(|_| { #func_name(#(#arg_pats),*) });
Some(#func_name(#(#arg_pats),*))
}
};

Expand All @@ -478,8 +471,19 @@ impl PgExtern {
funcctx.user_fctx = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).palloc_struct::<IteratorHolder<#retval_ty_resolved>>() as *mut ::core::ffi::c_void;
iterator_holder = pgx::PgBox::from_pg(funcctx.user_fctx as *mut IteratorHolder<#retval_ty_resolved>);

#( #arg_fetches )*
#result_handler
// function arguments need to be "fetched" while in the function call's
// multi-call-memory-context to ensure that any detoasted datums will
// live long enough for the SRF to use them over each call
Comment on lines +474 to +476
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is what needs to also be documented in a place outside the context of this function's comments. Especially so because this is something that will always need to remain true: we always must satisfy this sort of lifetime-binding condition, and it doesn't matter whether the code is burned and rewritten a million times: it remains a constraint. It should be easier to reference this sort of information than reading thousands of lines of implementations line by line, because it determines how fast someone can begin to understand the concepts required to write correct code for implementing PGX.

let #result_ident = match pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| {
#( #arg_fetches )*
#result_handler
}) {
Some(result) => result,
None => {
pgx::srf_return_done(#fcinfo_ident, &mut funcctx);
return pgx::pg_return_null(#fcinfo_ident)
}
};

iterator_holder.iter = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).leak_trivial_alloc(result);
}
Expand Down Expand Up @@ -546,17 +550,11 @@ impl PgExtern {
let result_handler = if optional {
// don't need unsafe annotations because of the larger unsafe block coming up
quote_spanned! { self.func.sig.span() =>
let #result_ident = match pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| { #func_name(#(#arg_pats),*) }) {
Some(result) => result,
None => {
pgx::srf_return_done(#fcinfo_ident, &mut funcctx);
return pgx::pg_return_null(#fcinfo_ident)
}
};
#func_name(#(#arg_pats),*)
}
} else {
quote_spanned! { self.func.sig.span() =>
let #result_ident = pgx::PgMemoryContexts::For(#funcctx_ident.multi_call_memory_ctx).switch_to(|_| { #func_name(#(#arg_pats),*) });
Some(#func_name(#(#arg_pats),*))
}
};

Expand Down Expand Up @@ -589,8 +587,19 @@ impl PgExtern {
});
iterator_holder = pgx::PgBox::from_pg(funcctx.user_fctx as *mut IteratorHolder<#retval_tys_tuple>);

#( #arg_fetches )*
#result_handler
// function arguments need to be "fetched" while in the function call's
// multi-call-memory-context to ensure that any detoasted datums will
// live long enough for the SRF to use them over each call
let #result_ident = match pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| {
#( #arg_fetches )*
#result_handler
}) {
Some(result) => result,
None => {
pgx::srf_return_done(#fcinfo_ident, &mut funcctx);
return pgx::pg_return_null(#fcinfo_ident)
}
};

iterator_holder.iter = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).leak_and_drop_on_delete(result);
}
Expand Down