pg_array_tools.c

#include <postgres.h>
#include <utils/array.h>
#include "funcapi.h"
#include "fmgr.h"
#include "stdbool.h"
#include "utils/builtins.h"
#include "executor/spi.h"
#include "access/tupdesc.h"
#include <stdlib.h>
#undef HASH_FUNCTION
#define HASH_FUNCTION HASH_JEN
#include "uthash.h"

#define SQL_CACHE_SIZE 6

PG_MODULE_MAGIC;

typedef struct {
    int32 key;
    int32 val;
    UT_hash_handle hh;  /* Pole wymagane przez UTHASH */
} MappingEntry;

typedef struct {
    char *sql;             /* Zapytanie SQL */
    MappingEntry *hash;    /* Wskaźnik do tabeli hash */
    int count;             /* Liczba elementów */
} SqlCacheEntry;

static SqlCacheEntry cache_list[SQL_CACHE_SIZE];

/* Bufory globalne do cache'owania */
static char *cached_sql = NULL;
static MappingEntry *cached_hash = NULL;  /* Hash table zamiast tablicy */

static MappingEntry* ensure_loaded(char *sql_query) {
    int ret;
    SPITupleTable *tuptable;
    TupleDesc tupdesc;
    int ntups;
    int i;

    /* 1. Sprawdź czy to SQL nie jest już w cache */
    for (i = 0; i < SQL_CACHE_SIZE; i++) {
        if (cache_list[i].sql != NULL && strcmp(cache_list[i].sql, sql_query) == 0) {
            /* HIT: Znajdźmy ten wpis na samej górze (strategia LRU) */
            /*
            if (i > 0) {
                / * Przesuń elementy w dół, aby ten był na 0 * /
                SqlCacheEntry temp = cache_list[i];
                memmove(&cache_list[1], &cache_list[0], i * sizeof(SqlCacheEntry));
                cache_list[0] = temp;
            }*/
            return cache_list[i].hash;
        }
    }

    /* 2. MISS: Musimy wczytać dane z bazy */
    ret = SPI_connect();
    if (ret < 0) {
        ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
                        errmsg("SPI_connect failed")));
    }

    ret = SPI_exec(sql_query, 0);
    if (ret < 0 || ret != SPI_OK_SELECT) {
        SPI_finish();
        ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
                        errmsg("SPI_exec failed for query: %s", sql_query)));
    }

    tuptable = SPI_tuptable;
    tupdesc = tuptable->tupdesc;
    ntups = SPI_processed;
    if (tupdesc->natts < 2) {
        SPI_finish();
        ereport(ERROR, (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
                        errmsg("Subquery error")));
    }

    /* 
       WAŻNE: Alokujemy pamięć w TopMemoryContext.
       Dzięki temu dane przetrwają nawet jeśli funkcja jest wywoływana 
       w pętli (gdzie CurrentMemoryContext może zostać resetowany).
    */
    MemoryContext oldcontext = MemoryContextSwitchTo(TopMemoryContext);// to jest dobrze w TopMemoryContext

    /* Jeśli cache był pełny, zwolnij najstarszy element (na samym dole) */
    if (SQL_CACHE_SIZE > 0 && cache_list[SQL_CACHE_SIZE - 1].sql != NULL) {
        pfree(cache_list[SQL_CACHE_SIZE - 1].sql);
        MappingEntry *elem, *tmp;
        HASH_ITER(hh, cache_list[SQL_CACHE_SIZE - 1].hash, elem, tmp) {
            HASH_DEL(cache_list[SQL_CACHE_SIZE - 1].hash, elem);
            pfree(elem);
        }
    }

    /* Przesuń istniejące elementy w dół (strategia LRU) */
    if (SQL_CACHE_SIZE > 1) {
        memmove(&cache_list[1], &cache_list[0], (SQL_CACHE_SIZE - 1) * sizeof(SqlCacheEntry));
    }

    /* Przygotuj nowy element na samej górze (indeks 0) */
    MappingEntry *new_hash = NULL;
    
    if (ntups > 0) {
        bool isnullKey;
        bool isnullVal;
        for (int j = 0; j < ntups; j++) {
            HeapTuple tuple = tuptable->vals[j];
            MappingEntry *entry = palloc(sizeof(MappingEntry));
            entry->key = DatumGetInt32(SPI_getbinval(tuple, tupdesc, 1, &isnullKey));
            entry->val = DatumGetInt32(SPI_getbinval(tuple, tupdesc, 2, &isnullVal));
            if (!isnullKey && !isnullVal) {
                HASH_ADD_INT(new_hash, key, entry);
            }
        }
        cache_list[0].hash = new_hash;
        cache_list[0].count = ntups;
    } else {
        cache_list[0].hash = NULL;
        cache_list[0].count = 0;
    }

    /* Zapisz zapytanie SQL */
    cache_list[0].sql = pstrdup(sql_query);
    MemoryContextSwitchTo(oldcontext);
    SPI_finish();
    return cache_list[0].hash;
}


PG_FUNCTION_INFO_V1(translate_int);

Datum translate_int(PG_FUNCTION_ARGS)
{
    ArrayType    *input_array = PG_GETARG_ARRAYTYPE_P(0);
    Datum        sql_datum   = PG_GETARG_DATUM(1);
    char         *sql_query  = TextDatumGetCString(sql_datum);
    bool         filter_mode = PG_GETARG_BOOL(2);
    
    ArrayType        *output_array;
    int32       *input_vec;
    int          n_input;
    MappingEntry *map_hash;
    MappingEntry *res;
    int          i;
    int32       *new_values;

    /* 1. Sprawdź cache i wczytaj dane */
    map_hash = ensure_loaded(sql_query);
    int ndims = ARR_NDIM(input_array);

    if (ndims == 0)
    {
        output_array = construct_empty_array(INT4ARRAYOID);
        PG_RETURN_ARRAYTYPE_P(output_array);
    }

    /* 2. Przygotuj dane wejściowe */
    if (ndims != 1)
    {
        /* Pobierz informacje o rozmiarach */
        int32 dims[ARR_NDIM(input_array)];
        int i;
        for (i = 0; i < ndims; i++)
        {
            dims[i] = ARR_DIMS(input_array)[i];
        }

        ereport(ERROR,
                (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
                errmsg("Only 1D arrays are supported"),
                errdetail("Received %d-dimensional array with dimensions: %d",
                        ndims,
                        ndims > 0 ? ARR_DIMS(input_array)[0] : 0),
                errhint("Use a single-dimensional array like {1,2,3}")));
        PG_RETURN_NULL();
    }

    if (ARR_HASNULL(input_array)){
        ereport(ERROR, (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
                        errmsg("Input array must not contain nulls")));
        PG_RETURN_NULL();
    }

    input_vec = (int32 *)ARR_DATA_PTR(input_array);
    n_input = ARR_DIMS(input_array)[0];
    if (n_input == 0)
    {
        output_array = construct_empty_array(INT4ARRAYOID);
        PG_RETURN_ARRAYTYPE_P(output_array);
    }
    /* 3. Przygotuj bufor wynikowy */
    new_values = palloc(n_input * sizeof(int32));

    /* 4. Pętla transformacji z UTHASH lookup */
    int out_idx = 0;

    for (i = 0; i < n_input; i++) {
        int32 current_val = input_vec[i];
        bool found = false;

        if (map_hash != NULL) {
            /* O(1) lookup */
            HASH_FIND_INT(map_hash, &current_val, res);
            
            if (res != NULL) {
                new_values[out_idx] = res->val;
                found = true;
            }
        }

        if (found) {
            out_idx++;
        } else {
            if (!filter_mode) {
                new_values[out_idx] = current_val;
                out_idx++;
            }
        }
    }

    /* 1. Przygotuj tablicę Datumów */
    Datum *elem_dats = palloc(out_idx * sizeof(Datum));
    if (elem_dats == NULL) {
        pfree(new_values);  // Zwolnij pamięć
        ereport(ERROR, (errcode(ERRCODE_OUT_OF_MEMORY),
                        errmsg("Allocation error.")));
        PG_RETURN_NULL();
    }

    for(int j = 0; j < out_idx; j++) {
        elem_dats[j] = Int32GetDatum(new_values[j]);
    }
    pfree(new_values);  // Zwolnij pamięć

    /* 2. Zbuduj tablicę SQL używając standardowego API */
    output_array = construct_array(elem_dats, out_idx, INT4OID, 4, true, 'i');

    /* 3. Zwolnij tymczasową pamięć */
    pfree(elem_dats);

    PG_RETURN_ARRAYTYPE_P(output_array);
}