Introduction

PostgreSQL is far more than a simple relational database. With its advanced analytical functions, flexible data types, and extensibility, it serves as a powerful backbone for machine learning workflows. Many ML engineers underestimate what can be done directly in the database layer, often pulling raw data into Python for transformations that PostgreSQL handles more efficiently.

This tutorial covers advanced PostgreSQL techniques that are directly applicable to ML pipelines: feature engineering with window functions, query organization with CTEs, caching computed features with materialized views, storing experiment metadata with JSONB, handling large datasets with partitioning, monitoring query performance, and integrating everything with Python using psycopg2 and SQLAlchemy.

Prerequisites

• PostgreSQL 14+ installed and running
• Python 3.9+ with pip
• Basic SQL knowledge (SELECT, JOIN, GROUP BY)
• Familiarity with ML concepts (features, training data, experiments)
• Install required Python packages:

# Install dependencies
pip install psycopg2-binary sqlalchemy pandas

import psycopg2
import sqlalchemy
import pandas as pd

Window Functions for Feature Engineering

Window functions are one of the most powerful tools for ML feature engineering. They allow you to compute aggregates across rows related to the current row without collapsing the result set — perfect for creating rolling statistics, lag features, and ranking features.

Rolling Statistics

-- Create a sample transactions table
CREATE TABLE transactions (
    id SERIAL PRIMARY KEY,
    customerid INTEGER NOT NULL,

    amount DECIMAL(10, 2) NOT NULL,
    category VARCHAR(50),
    createdat TIMESTAMP NOT NULL DEFAULT NOW()
);

-- Rolling average transaction amount (last 7 days) per customer
SELECT
    id,
    customerid,

    amount,
    createdat,
    AVG(amount) OVER (
        PARTITION BY customerid

        ORDER BY createdat
        RANGE BETWEEN INTERVAL '7 days' PRECEDING AND CURRENT ROW
    ) AS rollingavg7d,
    STDDEV(amount) OVER (
        PARTITION BY customerid

        ORDER BY createdat
        RANGE BETWEEN INTERVAL '7 days' PRECEDING AND CURRENT ROW
    ) AS rollingstddev7d,
    COUNT() OVER (

        PARTITION BY customerid

        ORDER BY createdat
        RANGE BETWEEN INTERVAL '7 days' PRECEDING AND CURRENT ROW
    ) AS txncount7d
FROM transactions
ORDER BY customerid, createdat;

-- Lag features: previous transaction amounts and time deltas
SELECT
    id,
    customerid,

    amount,
    createdat,
    LAG(amount, 1) OVER w AS prevamount1,
    LAG(amount, 2) OVER w AS prevamount2,
    LAG(amount, 3) OVER w AS prevamount3,
    amount - LAG(amount, 1) OVER w AS amountdiff,

    EXTRACT(EPOCH FROM (
        createdat - LAG(createdat, 1) OVER w

    )) / 3600.0 AS hourssincelasttxn,
    NTILE(10) OVER (
        PARTITION BY customerid

        ORDER BY amount
    ) AS amountdecile
FROM transactions
WINDOW w AS (PARTITION BY customerid ORDER BY createdat);

-- Rank customers by spending within each category
SELECT
    customerid,

    category,
    SUM(amount) AS totalspent,
    RANK() OVER (
        PARTITION BY category
        ORDER BY SUM(amount) DESC
    ) AS spendingrank,

    PERCENTRANK() OVER (
        PARTITION BY category
        ORDER BY SUM(amount) DESC
    ) AS spendingpercentile

FROM transactions
GROUP BY customerid, category;

-- Complete feature engineering pipeline using CTEs
WITH customerstats AS (

    -- Step 1: Basic aggregations per customer
    SELECT
        customerid,
        COUNT() AS totaltransactions,

        SUM(amount) AS totalspent,
        AVG(amount) AS avgamount,

        STDDEV(amount) AS stddevamount,
        MIN(amount) AS minamount,

        MAX(amount) AS maxamount,
        MIN(createdat) AS firsttransaction,
        MAX(createdat) AS lasttransaction
    FROM transactions
    GROUP BY customerid

),
categoryfeatures AS (
    -- Step 2: Spending distribution across categories
    SELECT
        customerid,

        COUNT(DISTINCT category) AS uniquecategories,
        MODE() WITHIN GROUP (ORDER BY category) AS mostcommoncategory,
        MAX(CASE WHEN category = 'electronics' THEN amount ELSE 0 END) AS maxelectronics

    FROM transactions
    GROUP BY customerid
),
recencyfeatures AS (

    -- Step 3: RFM-style recency features
    SELECT
        customerid,
        EXTRACT(EPOCH FROM (NOW() - MAX(createdat))) / 86400.0 AS dayssincelasttxn,
        EXTRACT(EPOCH FROM (MAX(createdat) - MIN(createdat))) / 86400.0 AS customertenuredays
    FROM transactions
    GROUP BY customerid

)
-- Step 4: Combine all features
SELECT
    cs.customerid,
    cs.totaltransactions,

    cs.totalspent,
    cs.avgamount,

    cs.stddevamount,
    cs.maxamount - cs.minamount AS amountrange,

    COALESCE(cs.stddevamount / NULLIF(cs.avgamount, 0), 0) AS cvamount,
    cf.uniquecategories,

    cf.mostcommoncategory,

    rf.dayssincelasttxn,
    rf.customertenuredays,
    cs.totaltransactions / NULLIF(rf.customertenuredays, 0) AS txnfrequency
FROM customerstats cs

JOIN categoryfeatures cf ON cs.customerid = cf.customerid
JOIN recencyfeatures rf ON cs.customerid = rf.customerid;

-- Find referral chain depth (graph-like feature)
WITH RECURSIVE referralchain AS (
    -- Base case: direct referrals
    SELECT
        customerid,

        referredby,
        1 AS chaindepth

    FROM customers
    WHERE referredby IS NOT NULL

    UNION ALL

    -- Recursive step
    SELECT
        rc.customerid,

        c.referredby,
        rc.chaindepth + 1

    FROM referralchain rc
    JOIN customers c ON rc.referredby = c.customerid
    WHERE c.referredby IS NOT NULL

    AND rc.chaindepth < 10  -- Safety limit
)
SELECT customerid, MAX(chaindepth) AS referraldepth

FROM referralchain
GROUP BY customerid;

-- Create a materialized view for customer features
CREATE MATERIALIZED VIEW mvcustomerfeatures AS

SELECT
    t.customerid,
    COUNT() AS totaltransactions,

    SUM(t.amount) AS totalspent,
    AVG(t.amount) AS avgtransaction,

    STDDEV(t.amount) AS stddevtransaction,
    COUNT(DISTINCT t.category) AS categorydiversity,

    COUNT(DISTINCT DATETRUNC('day', t.createdat)) AS activedays,
    MAX(t.createdat) AS lastactivity,
    EXTRACT(EPOCH FROM (NOW() - MAX(t.createdat))) / 86400.0 AS daysinactive
FROM transactions t
GROUP BY t.customerid

WITH DATA;

-- Create an index for fast lookups
CREATE UNIQUE INDEX idxmvcustomerfeaturesid

    ON mvcustomerfeatures (customerid);

-- Refresh the view (run on a schedule, e.g., daily via cron)
REFRESH MATERIALIZED VIEW CONCURRENTLY mvcustomerfeatures;

-- Track when views were last refreshed
CREATE TABLE mvrefreshlog (
    viewname TEXT PRIMARY KEY,

    lastrefreshed TIMESTAMP NOT NULL DEFAULT NOW(),
    durationms INTEGER,

    rowcount BIGINT
);

-- Refresh function with logging
CREATE OR REPLACE FUNCTION refreshfeatureview(viewname TEXT)

RETURNS VOID AS $
DECLARE
    starttime TIMESTAMP;
    endtime TIMESTAMP;

    rcount BIGINT;
BEGIN
    starttime := clocktimestamp();

    EXECUTE format('REFRESH MATERIALIZED VIEW CONCURRENTLY %I', viewname);
    endtime := clocktimestamp();

    EXECUTE format('SELECT COUNT() FROM %I', viewname) INTO rcount;


    INSERT INTO mvrefreshlog (viewname, lastrefreshed, durationms, rowcount)

    VALUES (
        viewname,
        NOW(),
        EXTRACT(MILLISECONDS FROM (endtime - starttime))::INTEGER,
        rcount
    )
    ON CONFLICT (viewname)

    DO UPDATE SET
        lastrefreshed = EXCLUDED.lastrefreshed,

        durationms = EXCLUDED.durationms,

        rowcount = EXCLUDED.rowcount;

END;
$ LANGUAGE plpgsql;

-- Usage
SELECT refreshfeatureview('mvcustomerfeatures');

-- Experiment tracking table
CREATE TABLE mlexperiments (
    id SERIAL PRIMARY KEY,
    experimentname VARCHAR(200) NOT NULL,

    modeltype VARCHAR(100) NOT NULL,
    hyperparameters JSONB NOT NULL DEFAULT '{}',
    metrics JSONB NOT NULL DEFAULT '{}',
    featureconfig JSONB,

    tags TEXT[],
    createdat TIMESTAMP NOT NULL DEFAULT NOW(),
    trainingdurationseconds FLOAT
);

-- Insert an experiment
INSERT INTO mlexperiments (experimentname, modeltype, hyperparameters, metrics, tags)

VALUES (
    'churnpredictionv3',

    'xgboost',
    '{"nestimators": 500, "maxdepth": 6, "learningrate": 0.01,
      "subsample": 0.8, "colsamplebytree": 0.8}',

    '{"accuracy": 0.923, "precision": 0.891, "recall": 0.876,
      "f1": 0.883, "aucroc": 0.954}',
    ARRAY['production', 'churn', 'v3']
);

-- Query experiments by hyperparameter values
SELECT
    experimentname,

    hyperparameters->>'learningrate' AS lr,
    metrics->>'aucroc' AS auc,

    metrics->>'f1' AS f1
FROM mlexperiments
WHERE modeltype = 'xgboost'

  AND (hyperparameters->>'maxdepth')::int >= 5
  AND (metrics->>'aucroc')::float > 0.90

ORDER BY (metrics->>'aucroc')::float DESC;

-- Aggregate metrics across experiments
SELECT
    modeltype,

    AVG((metrics->>'aucroc')::float) AS avgauc,

    MAX((metrics->>'aucroc')::float) AS bestauc,

    COUNT() AS numexperiments
FROM mlexperiments

GROUP BY modeltype
ORDER BY avgauc DESC;


-- GIN index for fast JSONB queries
CREATE INDEX idxexperimentshyperparams ON mlexperiments
    USING GIN (hyperparameters);
CREATE INDEX idxexperimentsmetrics ON mlexperiments

    USING GIN (metrics);

-- Create a partitioned table for ML training data
CREATE TABLE mltrainingdata (

    id BIGSERIAL,
    featurevector FLOAT8[] NOT NULL,
    label INTEGER NOT NULL,
    datasplit VARCHAR(10) NOT NULL CHECK (datasplit IN ('train', 'val', 'test')),
    createddate DATE NOT NULL DEFAULT CURRENTDATE,
    source VARCHAR(50),
    PRIMARY KEY (id, createddate)

) PARTITION BY RANGE (createddate);

-- Create monthly partitions
CREATE TABLE mltrainingdata202601
    PARTITION OF mltrainingdata
    FOR VALUES FROM ('2026-01-01') TO ('2026-02-01');

CREATE TABLE mltrainingdata202602
    PARTITION OF mltrainingdata
    FOR VALUES FROM ('2026-02-01') TO ('2026-03-01');

CREATE TABLE mltrainingdata202603
    PARTITION OF mltrainingdata
    FOR VALUES FROM ('2026-03-01') TO ('2026-04-01');

-- Automate partition creation
CREATE OR REPLACE FUNCTION createmonthlypartition(
    tablename TEXT,

    year INT,
    month INT
) RETURNS VOID AS $
DECLARE
    partitionname TEXT;
    startdate DATE;

    enddate DATE;
BEGIN
    partitionname := format('%s%s%s', tablename, year,
                            LPAD(month::TEXT, 2, '0'));
    startdate := makedate(year, month, 1);
    enddate := startdate + INTERVAL '1 month';

    EXECUTE format(
        'CREATE TABLE IF NOT EXISTS %I PARTITION OF %I
         FOR VALUES FROM (%L) TO (%L)',
        partitionname, tablename, startdate, enddate
    );

    RAISE NOTICE 'Created partition: %', partitionname;

END;
$ LANGUAGE plpgsql;

-- Create partitions for the next 12 months
DO $
BEGIN
    FOR m IN 1..12 LOOP
        PERFORM createmonthlypartition('mltrainingdata', 2026, m);

    END LOOP;
END $;

-- PostgreSQL automatically prunes irrelevant partitions
EXPLAIN ANALYZE
SELECT COUNT(), AVG(label::float)
FROM mltrainingdata

WHERE createddate BETWEEN '2026-01-01' AND '2026-01-31'
  AND datasplit = 'train';

-- Only scans mltrainingdata202601

Monitoring with pgstat

Understanding query performance is critical when your ML pipeline depends on database queries. PostgreSQL provides extensive statistics through pgstat views.

-- Enable pgstatstatements (add to postgresql.conf)
-- sharedpreloadlibraries = 'pgstatstatements'


-- Top slow queries hitting your ML tables
SELECT
    calls,
    meanexectime::numeric(10,2) AS avgms,
    totalexectime::numeric(10,2) AS totalms,

    rows,
    LEFT(query, 120) AS querypreview
FROM pgstatstatements
WHERE query ILIKE '%transaction%' OR query ILIKE '%ml%'

ORDER BY meanexectime DESC

LIMIT 20;

-- Table-level statistics
SELECT
    schemaname,
    relname AS tablename,
    seqscan,

    seqtupread,

    idxscan,
    idxtupfetch,
    ntupins AS inserts,
    ntupupd AS updates,
    ntupdel AS deletes,
    nlivetup AS liverows,

    ndeadtup AS deadrows,
    lastvacuum,

    lastautovacuum,
    lastanalyze

FROM pgstatusertables
WHERE relname LIKE 'ml%' OR relname LIKE 'mv%'
ORDER BY seqtupread DESC;

-- Index usage analysis
SELECT
    indexrelname AS indexname,

    relname AS tablename,
    idxscan AS timesused,
    pgsizepretty(pgrelationsize(indexrelid)) AS indexsize

FROM pgstatuserindexes
WHERE schemaname = 'public'
ORDER BY idxscan DESC;


-- Cache hit ratio (should be > 99%)
SELECT
    sum(heapblksread) AS heapread,
    sum(heapblkshit)  AS heaphit,

    ROUND(sum(heapblkshit)  100.0 /

          NULLIF(sum(heapblkshit) + sum(heapblksread), 0), 2

    ) AS cachehitratio

FROM pgstatiousertables;

import psycopg2
from psycopg2.extras import RealDictCursor, executevalues

import numpy as np
import pandas as pd

class MLDatabaseClient:
    """Database client optimized for ML workflows."""

    def init(self, host, port, dbname, user, password):
        self.connparams = {
            "host": host,
            "port": port,
            "dbname": dbname,
            "user": user,
            "password": password,
        }
        self.conn = None

    def connect(self):
        self.conn = psycopg2.connect(self.connparams)

        self.conn.autocommit = False
        return self

    def close(self):
        if self.conn:
            self.conn.close()

    def enter(self):
        return self.connect()

    def exit(self, exctype, excval, exctb):
        if exctype:

            self.conn.rollback()
        else:
            self.conn.commit()
        self.close()

    def fetchfeatures(self, customerids: list) -> pd.DataFrame:

        """Fetch pre-computed features for a list of customers."""
        query = """
            SELECT  FROM mvcustomerfeatures

            WHERE customerid = ANY(%s)
        """
        with self.conn.cursor(cursorfactory=RealDictCursor) as cur:

            cur.execute(query, (customerids,))
            rows = cur.fetchall()
        return pd.DataFrame(rows)

    def fetchtrainingdata(self, split: str, datefrom: str,

                            dateto: str, limit: int = None) -> pd.DataFrame:
        """Fetch training data with partition pruning."""
        query = """
            SELECT featurevector, label

            FROM mltrainingdata

            WHERE datasplit = %s
              AND createddate BETWEEN %s AND %s

        """
        params = [split, datefrom, dateto]

        if limit:
            query += " LIMIT %s"
            params.append(limit)

        with self.conn.cursor() as cur:
            cur.execute(query, params)
            rows = cur.fetchall()

        features = np.array([row[0] for row in rows])
        labels = np.array([row[1] for row in rows])
        return features, labels

    def logexperiment(self, name: str, modeltype: str,

                       hyperparams: dict, metrics: dict, tags: list):
        """Log an ML experiment to the database."""
        import json
        query = """
            INSERT INTO mlexperiments
                (experimentname, modeltype, hyperparameters, metrics, tags)
            VALUES (%s, %s, %s, %s, %s)
            RETURNING id
        """
        with self.conn.cursor() as cur:
            cur.execute(query, (
                name, modeltype,

                json.dumps(hyperparams),
                json.dumps(metrics),
                tags
            ))
            experimentid = cur.fetchone()[0]
        self.conn.commit()
        return experimentid


    def bulkinserttrainingdata(self, features: np.ndarray,
                                   labels: np.ndarray, split: str):
        """Efficiently bulk insert training data."""
        data = [
            (features[i].tolist(), int(labels[i]), split)
            for i in range(len(labels))
        ]
        query = """
            INSERT INTO mltrainingdata (featurevector, label, datasplit)
            VALUES %s
        """
        with self.conn.cursor() as cur:
            executevalues(cur, query, data, pagesize=1000)
        self.conn.commit()


Usage example
if name == "main":
    with MLDatabaseClient(
        host="localhost", port=5432,
        dbname="mldb", user="mluser", password="secret"
    ) as db:
        # Fetch features
        df = db.fetchfeatures([1, 2, 3, 4, 5])

        print(f"Fetched {len(df)} customer feature rows")

        # Log experiment
        expid = db.logexperiment(

            name="churnv4",
            modeltype="lightgbm",

            hyperparams={"nestimators": 1000, "learningrate": 0.05},

            metrics={"aucroc": 0.961, "f1": 0.892},
            tags=["production", "churn"]
        )
        print(f"Logged experiment ID: {expid}")

from sqlalchemy import (
    createengine, Column, Integer, String, Float,
    DateTime, ARRAY, JSON, func, text
)
from sqlalchemy.orm import declarativebase, Session

from sqlalchemy.dialects.postgresql import JSONB
from datetime import datetime

Base = declarativebase()

class MLExperiment(Base):
    tablename = "mlexperiments"


    id = Column(Integer, primarykey=True, autoincrement=True)
    experimentname = Column(String(200), nullable=False)

    modeltype = Column(String(100), nullable=False)
    hyperparameters = Column(JSONB, default={})
    metrics = Column(JSONB, default={})
    tags = Column(ARRAY(String))
    createdat = Column(DateTime, default=datetime.utcnow)

    trainingdurationseconds = Column(Float)


    def repr(self):
        return f"name} ({self.modeltype})>"



engine = createengine(
    "postgresql://mluser:secret@localhost:5432/mldb",
    poolsize=10,

    maxoverflow=20,
    echo=False
)

Query best experiments using SQLAlchemy
with Session(engine) as session:
    bestexperiments = (

        session.query(MLExperiment)
        .filter(MLExperiment.modeltype == "xgboost")
        .filter(
            MLExperiment.metrics["aucroc"].astext.cast(Float) > 0.90

        )
        .orderby(
            MLExperiment.metrics["aucroc"].astext.cast(Float).desc()

        )
        .limit(10)
        .all()
    )

    for exp in bestexperiments:
        print(f"{exp.experimentname}: AUC={exp.metrics['aucroc']}")

    # Use raw SQL for complex feature queries
    featuredf = pd.readsql(
        text("SELECT * FROM mvcustomerfeatures WHERE daysinactive < 30"),

        engine
    )