pg_onnx: ONNX Runtime integrated with PostgreSQL

[kibae]

Hi, all 🤗

I'd like to introduce another personal project.
https://github.com/kibae/pg_onnx

I developed an extension for PostgreSQL that allows you to perform ML inference using data in the DB. Inside, the onnxruntime-server I introduced before is running.

You can infer directly from the data in the DB, or you can use triggers to update some columns with ML inference results when new data is added.

If you have any questions, please feel free to let me know.
Thank you.

---

Simple Usage

• Import an ONNX file and get the inference results.

SELECT pg_onnx_import_model(
        'sample_model', --------------- model name
        'v20230101', ------------------ model version 
        PG_READ_BINARY_FILE('/your_model_path/model.onnx')::bytea, -- model binary data
        '{"cuda": true}'::jsonb, ------ options
        'sample model' ---------------- description
    );

SELECT pg_onnx_execute_session(
        'sample_model', -- model name
        'v20230101', ----- model version
        '{
          "x": [[1], [2], [3]],
          "y": [[3], [4], [5]],
          "z": [[5], [6], [7]]
        }' --------------- inputs
    );

• Depending on the type and shape of the inputs and outputs of the ML model, you can see different results. Below is an
  example of the result.

                       pg_onnx_execute_session                             
--------------------------------------------------------------------------------
 {"output": [[0.7488641738891602], [0.8607008457183838], [0.9725375175476074]]}

Using inference results with triggers

• When data is added, use the BEFORE INSERT trigger to update some columns with ML inference results.
• Depending on your ML model, this can have a significant performance impact, so be careful when using it.

-- Create a test table
CREATE TABLE trigger_test
(
    id         SERIAL PRIMARY KEY,
    value1     INT,
    value2     INT,
    value3     INT,
    prediction FLOAT
);

-- Create a trigger function
CREATE OR REPLACE FUNCTION trigger_test_insert()
    RETURNS TRIGGER AS
$$
DECLARE
    result jsonb;
BEGIN
    result := pg_onnx_execute_session(
            'sample_model', 'v20230101',
            JSONB_BUILD_OBJECT(
                    'x', ARRAY [[NEW.value1]],
                    'y', ARRAY [[NEW.value2]],
                    'z', ARRAY [[NEW.value3]]));

    -- output shape: float[-1,1]
    -- eg: {"output": [[0.6492120623588562]]}
    NEW.prediction := result -> 'output' -> 0 -> 0;
    RETURN NEW;
END;
$$
    LANGUAGE plpgsql;

-- Create a trigger
CREATE TRIGGER trigger_test_insert
    BEFORE INSERT
    ON trigger_test
    FOR EACH ROW
EXECUTE PROCEDURE trigger_test_insert();