PostgreSQL C Extension function: Table as argument and as Result

Question

I need to write an extension function in C for PostgreSQL that:

Has as argument a table of 3 columns of type double
Has as result a table of 3 columns of type double

Note that the mode should ideally be in materialize mode, since the function will generate a new table by transforming the data on the entire input table.

I did find information related to having the result as sets (multiple rows). But the test of the function does not result into the values as columns - I just get a record with (val,val,val).

Can anyone please direct me to information or preferably a simple example function that has this basic typology?

Thank you!

I tried to find information in the PostgreSQL documentation but it is not clear, specially when talking about tables as input arguments and returns in materialize mode.

Look in the source code for functions that return tables and learn from that. — Laurenz Albe, Apr 20 '23 at 03:02
Thanks @LaurenzAlbe. Great! Can you point to one function that I can look at? — Rafael Scudelari de Macedo, Apr 20 '23 at 09:27
I think `pg_logdir_ls_internal` in `contrib/adminpack/adminpack.c` is a simple enough example, and it uses `SFRM_Materialize`. — Laurenz Albe, Apr 20 '23 at 09:46
Thank you @LaurenzAlbe - I think it will allow me to work my way from there :) — Rafael Scudelari de Macedo, Apr 20 '23 at 10:44

score 1 · Answer 1 · answered Apr 20 '23 at 15:26

Following the suggestion from LaurenzAlbe, I created something like the following.

Note that I did not exactly find how to receive a table as input for a function. The idea here would be to wrap this into a SQL function so that it would create a temp table, fill this temp table, and then pass the name of this temp table to the C function.

For my application, this is enough. But I would love to debate this and improve if anyone else would be willing to help.

Thanks,

#include "postgres.h"
#include <string.h>
#include "fmgr.h"
#include "utils/geo_decls.h"
#include "funcapi.h"
#include "executor/executor.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/tuplestore.h"
#include "catalog/pg_type.h"
#include "executor/spi.h"
#include "access/heapam.h"
#include "access/htup_details.h"
#include "miscadmin.h"

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

PGDLLEXPORT Datum add_table(PG_FUNCTION_ARGS);
PG_FUNCTION_INFO_V1(add_table);

Datum add_table(PG_FUNCTION_ARGS) {

    ReturnSetInfo* rsinfo = (ReturnSetInfo*)fcinfo->resultinfo;

    /* check to see if caller supports us returning a tuplestore */
    if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
        ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("set-valued function called in context that cannot accept a set")));
    if (!(rsinfo->allowedModes & SFRM_Materialize))
        ereport(ERROR, (errcode(ERRCODE_SYNTAX_ERROR), errmsg("materialize mode required, but it is not allowed in this context")));


    /* The name of the table must be given as string 'tablename'::regclass */
    char* inputTableOid = PG_GETARG_OID(0);

    /* Executes the select statement to get input table */
    StringInfoData commandBuffer;
    initStringInfo(&commandBuffer);
    appendStringInfo(&commandBuffer, "SELECT * FROM %s", get_rel_name(inputTableOid));

    /* Starts SPI */
    int connRes = SPI_connect();
    if (connRes != SPI_OK_CONNECT) {
        elog(ERROR, "SPI_connect failed: %d", connRes);
    }
    int selectRes = SPI_exec(commandBuffer.data, 0);
    if (selectRes != SPI_OK_SELECT) {
        elog(ERROR, "SELECT failed: %d", selectRes);
    }

    /* Saves the select results */
    SPITupleTable* inputTupTable = SPI_tuptable; // The pointer to the input table - copies the global var that is filled by SPI_exec
    int inputNRows = SPI_processed; // The actual number of returned rows - copies the global var that is filled by SPI_exec
    TupleDesc inputTupleDesc = inputTupTable->tupdesc; // The TupleDesc of the input table
    int natts = inputTupleDesc->natts; // The number of attributes in the input tuple (cols)

    /* The tupdesc and tuplestore must be created in ecxt_per_query_memory */
    MemoryContext oldcontext = MemoryContextSwitchTo(rsinfo->econtext->ecxt_per_query_memory);

    /* Makes the output TupleDesc */
    TupleDesc tupdesc = CreateTemplateTupleDesc(3);
    TupleDescInitEntry(tupdesc, (AttrNumber)1, "x", FLOAT4OID, -1, 0);
    TupleDescInitEntry(tupdesc, (AttrNumber)2, "y", FLOAT4OID, -1, 0);
    TupleDescInitEntry(tupdesc, (AttrNumber)3, "z", FLOAT4OID, -1, 0);
    tupdesc = BlessTupleDesc(tupdesc);

    /* Checks if random access is allowed */
    bool randomAccess = (rsinfo->allowedModes & SFRM_Materialize_Random) != 0;
    /* Starts the tuplestore */
    Tuplestorestate* tupstore = tuplestore_begin_heap(randomAccess, false, work_mem);

    /* Set the output */
    rsinfo->returnMode = SFRM_Materialize;
    rsinfo->setResult = tupstore;
    rsinfo->setDesc = tupdesc;

    /* Returns to the old context */
    MemoryContextSwitchTo(oldcontext);


    /* Fills the output table */
    for (size_t i = 0; i < inputNRows; i++) {
        /* Getting the values of the current row */
        HeapTuple inputTuple = inputTupTable->vals[i]; // Gets the current input tuple (row)

        Datum* inputValues = palloc(natts * sizeof(Datum));
        bool* inputNulls = palloc(natts * sizeof(bool));
        heap_deform_tuple(inputTuple, inputTupleDesc, inputValues, inputNulls); // Gets the value of the current row

        /* Adds 1 to each column of the input row */
        Datum* outputValues = palloc(natts * sizeof(Datum));
        for (size_t j = 0; j < natts; j++) {
            outputValues[j] = Float4GetDatum(DatumGetFloat4(inputValues[j]) + 1);
        }

        /* Builds the output tuple (row) */
        HeapTuple outputTuple = heap_form_tuple(tupdesc, outputValues, inputNulls);
        /* Puts in the output tuplestore */
        tuplestore_puttuple(tupstore, outputTuple);
    }

    return (Datum)0;

}

Then, to add link this extension into your postgresql, you use the following code:

-- FUNCTION: public.add_table(oid)

DROP FUNCTION IF EXISTS public.add_table(oid);

CREATE OR REPLACE FUNCTION public.add_table(
    oid,
    OUT x real,
    OUT y real,
    OUT z real)
    RETURNS SETOF record 
    LANGUAGE 'c'
    COST 1
    IMMUTABLE STRICT PARALLEL UNSAFE
    ROWS 1000

AS '...Path.../<dllName>.dll', 'add_table'
;

ALTER FUNCTION public.add_table(oid)
    OWNER TO postgres;

To call the function:

-- The table contains 3 real columns
select * from add_table('test_table'::regclass);

This will return the table where all values are incremented.

PostgreSQL C Extension function: Table as argument and as Result

1 Answers1