如何在PostgreSQL C函数中遍历表实现循环生成AGE邻接矩阵

需求概述 #

要为Apache AGE开发C函数get_adjacency_matrix，返回存储两节点间边数的邻接矩阵：

• 查询边存在性复杂度O(1)，空间复杂度O(v²)（v为顶点数）
• 顶点数据存储在graph_name._ag_label_vertex表，边数据存储在graph_name._ag_label_edge表
• 顶点表包含id（整数）和properties（JSON）列；边表包含id、start_id、end_id和properties列

已完成准备工作：

• 在age--1.3.0.sql中声明函数：

CREATE FUNCTION ag_catalog.get_adjacency_matrix(graph_name name)
RETURNS integer[][]
LANGUAGE c
AS 'MODULE_PATHNAME';

• 在src/backend/commands/graph_commands.c中定义函数框架：

PG_FUNCTION_INFO_V1(get_adjacency_matrix);
/*
* Function for returning the graph's adjacency matrix. This adjacency matrix
* stores the number of edges between two nodes. The complexity to check if an
* edge exists within an adjacency matrix is O(1), while the space complexity
* is O(v^2) with v = number of vertices.
*/
Datum get_adjacency_matrix(PG_FUNCTION_ARGS)
{
    PG_RETURN_VOID();
}

实现方案 #

两次使用SPI查询是合理方案：先获取所有顶点建立ID到矩阵索引的映射，再遍历边填充矩阵。以下是具体实现步骤：

1. 初始化SPI与变量 #

连接SPI，获取传入的图名参数，定义顶点ID映射、矩阵、顶点数量等核心变量。

2. 查询所有顶点ID #

构造顶点查询SQL，执行SPI查询后遍历结果集，用数组存储顶点ID，同时用哈希表建立ID到连续矩阵索引的映射（解决AGE顶点ID为大整数无法直接作为索引的问题）。

3. 初始化邻接矩阵 #

根据顶点数量v创建v×v的整数矩阵，初始值全部设为0。

4. 查询所有边并更新矩阵 #

构造边查询SQL，执行SPI查询后遍历每条边，通过哈希表找到顶点对应的矩阵索引，累加边数。

5. 构造并返回矩阵结果 #

使用PostgreSQL的construct_md_array函数将C二维数组转换为PostgreSQL支持的integer[][]类型返回。

6. 清理资源 #

关闭SPI连接，释放哈希表、数组等分配的内存。

完整实现代码 #

替换graph_commands.c中的函数框架为以下代码：

PG_FUNCTION_INFO_V1(get_adjacency_matrix);

Datum get_adjacency_matrix(PG_FUNCTION_ARGS)
{
    Name graph_name = PG_GETARG_NAME(0);
    int spi_rc;
    TupleDesc tupdesc;
    SPITupleTable *spi_tuptable;
    int num_vertices;
    uint64 *vertex_ids;
    HTAB *vertex_id_map;
    int **matrix;
    int i, j;
    Datum result;

    // 初始化SPI连接
    spi_rc = SPI_connect();
    if (spi_rc != SPI_OK_CONNECT)
        elog(ERROR, "SPI_connect failed: %d", spi_rc);

    // 查询所有顶点ID
    char vertex_query[256];
    snprintf(vertex_query, sizeof(vertex_query),
             "SELECT id FROM %I._ag_label_vertex", graph_name);

    spi_rc = SPI_exec(vertex_query, 0);
    if (spi_rc != SPI_OK_SELECT)
        elog(ERROR, "SPI_exec failed for vertex query: %d", spi_rc);

    num_vertices = SPI_processed;
    if (num_vertices == 0)
    {
        SPI_finish();
        PG_RETURN_NULL();
    }

    spi_tuptable = SPI_tuptable;
    tupdesc = spi_tuptable->tupdesc;

    // 分配存储顶点ID的数组
    vertex_ids = palloc(num_vertices * sizeof(uint64));
    // 创建哈希表，映射vertex_id到矩阵索引
    HASHCTL hash_ctl;
    memset(&hash_ctl, 0, sizeof(hash_ctl));
    hash_ctl.keysize = sizeof(uint64);
    hash_ctl.entrysize = sizeof(int);
    hash_ctl.hash = uint64_hash;

    vertex_id_map = hash_create("vertex_id_map", num_vertices, &hash_ctl, HASH_ELEM | HASH_FUNCTION);

    // 遍历顶点结果，填充数组和哈希表
    for (i = 0; i < num_vertices; i++)
    {
        HeapTuple tuple = spi_tuptable->vals[i];
        Datum id_datum = SPI_getbinval(tuple, tupdesc, 1, NULL);
        uint64 vertex_id = DatumGetUInt64(id_datum);

        vertex_ids[i] = vertex_id;
        int *idx_ptr = hash_search(vertex_id_map, &vertex_id, HASH_ENTER, NULL);
        *idx_ptr = i;
    }

    // 初始化邻接矩阵
    matrix = palloc(num_vertices * sizeof(int *));
    for (i = 0; i < num_vertices; i++)
    {
        matrix[i] = palloc0(num_vertices * sizeof(int));
    }

    // 查询所有边并更新矩阵
    char edge_query[256];
    snprintf(edge_query, sizeof(edge_query),
             "SELECT start_id, end_id FROM %I._ag_label_edge", graph_name);

    spi_rc = SPI_exec(edge_query, 0);
    if (spi_rc != SPI_OK_SELECT)
        elog(ERROR, "SPI_exec failed for edge query: %d", spi_rc);

    spi_tuptable = SPI_tuptable;
    tupdesc = spi_tuptable->tupdesc;

    int num_edges = SPI_processed;
    for (i = 0; i < num_edges; i++)
    {
        HeapTuple tuple = spi_tuptable->vals[i];
        Datum start_datum = SPI_getbinval(tuple, tupdesc, 1, NULL);
        Datum end_datum = SPI_getbinval(tuple, tupdesc, 2, NULL);

        uint64 start_id = DatumGetUInt64(start_datum);
        uint64 end_id = DatumGetUInt64(end_datum);

        // 查找顶点对应的矩阵索引
        int *start_idx_ptr = hash_search(vertex_id_map, &start_id, HASH_FIND, NULL);
        int *end_idx_ptr = hash_search(vertex_id_map, &end_id, HASH_FIND, NULL);

        if (start_idx_ptr != NULL && end_idx_ptr != NULL)
        {
            int start_idx = *start_idx_ptr;
            int end_idx = *end_idx_ptr;
            matrix[start_idx][end_idx]++;
        }
        // 处理边引用不存在顶点的异常情况
        else
        {
            elog(WARNING, "Edge references non-existent vertex (start_id: %lu, end_id: %lu)", start_id, end_id);
        }
    }

    // 构造PostgreSQL二维数组返回
    int dims[2] = {num_vertices, num_vertices};
    int lbs[2] = {1, 1}; // PostgreSQL数组默认从1开始索引
    Datum *matrix_data = palloc(num_vertices * num_vertices * sizeof(Datum));

    int pos = 0;
    for (i = 0; i < num_vertices; i++)
    {
        for (j = 0; j < num_vertices; j++)
        {
            matrix_data[pos++] = Int32GetDatum(matrix[i][j]);
        }
    }

    result = construct_md_array(matrix_data, NULL, 2, dims, lbs, INT4OID, sizeof(int), true, 'i');

    // 清理资源
    hash_destroy(vertex_id_map);
    pfree(vertex_ids);
    for (i = 0; i < num_vertices; i++)
    {
        pfree(matrix[i]);
    }
    pfree(matrix);
    pfree(matrix_data);

    SPI_finish();

    PG_RETURN_DATUM(result);
}

关键说明 #

• SPI使用：两次SPI查询是PostgreSQL扩展中访问表数据的标准方式，适配AGE的存储结构，是当前场景的最优选择之一
• 顶点ID映射：通过哈希表将大整数顶点ID转换为连续矩阵索引，保证矩阵的紧凑性和索引的有效性
• 内存管理：使用palloc分配的内存会由PostgreSQL自动回收，但显式清理哈希表和数组更符合规范
• 数组构造：遵循PostgreSQL数组从1开始索引的特性，用construct_md_array完成C数组到PostgreSQL数组的转换

内容的提问来源于stack exchange，提问作者Matheus Farias