add a draft schema to support snapshots and multiple training runs for a project

Montana Low · Montana Low · commit 829b62e2a50f · 2022-04-12T16:59:27.000-07:00
diff --git a/README.md b/README.md
@@ -5,7 +5,7 @@ PostgresML aims to be the easiest way to gain value from machine learning. Anyon
 Getting started is as easy as creating a `table` or `view` that holds the training data, and then registering that with PostgresML. 
 
 ```sql
-SELECT pgml.create_regression('Red Wine Quality', training_data_table_or_view_name, label_column_name);
+SELECT pgml.model_regression('Red Wine Quality', training_data_table_or_view_name, label_column_name);
 ```
 
 And predict novel datapoints:
@@ -23,7 +23,7 @@ LIMIT 3;
 (3 rows)
 ```
 
-PostgresML similarly supports classification to predict numeric scores rather than classes for novel data.
+PostgresML similarly supports classification to predict discrete classes rather than numeric scores for novel data.
 
 ```sql
 SELECT pgml.create_classification('Handwritten Digit Classifier', pgml.mnist_training_data, label_column_name);
diff --git a/benchmarks.sql b/benchmarks.sql
@@ -0,0 +1,23 @@
+--
+-- CREATE EXTENSION
+--
+CREATE EXTENSION IF NOT EXISTS plpython3u;
+
+CREATE OR REPLACE FUNCTION pg_call()
+RETURNS INT
+AS $$
+BEGIN
+    RETURN 1;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION py_call()
+RETURNS INT
+AS $$
+    return 1;
+$$ LANGUAGE plpython3u;
+
+\timing on
+SELECT generate_series(1, 50000), pg_call(); -- Time: 20.679 ms
+SELECT generate_series(1, 50000), py_call(); -- Time: 67.355 ms
+
diff --git a/pgml/pgml/model.py b/pgml/pgml/model.py
@@ -0,0 +1,95 @@
+import plpy
+
+class Regression:
+    """Provides continuous real number predictions learned from the training data.
+    """    
+    def __init__(
+        model_name: str, 
+        relation_name: str, 
+        y_column_name: str, 
+        implementation: str = "sklearn.linear_model"
+    ) -> None:
+        """Create a regression model from a table or view filled with training data.
+
+        Args:
+            model_name (str): a human friendly identifier
+            relation_name (str): the table or view that stores the training data
+            y_column_name (str): the column in the training data that acts as the label
+            implementation (str, optional): the algorithm used to implement the regression. Defaults to "sklearn.linear_model".
+        """
+
+        data_source = f"SELECT * FROM {table_name}"
+
+        # Start training.
+        start = plpy.execute(f"""
+            INSERT INTO pgml.model_versions
+                (name, data_source, y_column)
+            VALUES
+                ('{table_name}', '{data_source}', '{y}')
+            RETURNING *""", 1)
+
+        id_ = start[0]["id"]
+        name = f"{table_name}_{id_}"
+
+        destination = models_directory(plpy)
+
+        # Train!
+        pickle, msq, r2 = train(plpy.cursor(data_source), y_column=y, name=name, destination=destination)
+        X = []
+        y = []
+        columns = []
+
+        for row in all_rows(cursor):
+            row = row.copy()
+
+            if y_column not in row:
+                PgMLException(
+                    f"Column `{y}` not found. Did you name your `y_column` correctly?"
+                )
+
+            y_ = row.pop(y_column)
+            x_ = []
+
+            # Always pull the columns in the same order from the row.
+            # Python dict iteration is not always in the same order (hash table).
+            if not columns:
+                for col in row:
+                    columns.append(col)
+
+            for column in columns:
+                x_.append(row[column])
+            X.append(x_)
+            y.append(y_)
+
+        X_train, X_test, y_train, y_test = train_test_split(X, y)
+
+        # Just linear regression for now, but can add many more later.
+        lr = LinearRegression()
+        lr.fit(X_train, y_train)
+
+        # Test
+        y_pred = lr.predict(X_test)
+        msq = mean_squared_error(y_test, y_pred)
+        r2 = r2_score(y_test, y_pred)
+
+        path = os.path.join(destination, name)
+
+        if save:
+            with open(path, "wb") as f:
+                pickle.dump(lr, f)
+
+        return path, msq, r2
+
+
+        plpy.execute(f"""
+            UPDATE pgml.model_versions
+            SET pickle = '{pickle}',
+                successful = true,
+                mean_squared_error = '{msq}',
+                r2_score = '{r2}',
+                ended_at = clock_timestamp()
+            WHERE id = {id_}""")
+
+        return name
+
+            model
diff --git a/pgml/pgml/sql.py b/pgml/pgml/sql.py
@@ -1,6 +1,7 @@
 """Tools to run SQL.
 """
 import os
+import plpy
 
 
 def all_rows(cursor):
@@ -14,7 +15,7 @@ def all_rows(cursor):
             yield row
 
 
-def models_directory(plpy):
+def models_directory():
     """Get the directory where we store our models."""
     data_directory = plpy.execute(
         """
diff --git a/sql/install.sql b/sql/install.sql
@@ -1,10 +1,87 @@
+SET client_min_messages TO WARNING;
 
 -- Create the PL/Python3 extension.
 CREATE EXTENSION IF NOT EXISTS plpython3u;
 
+---
+--- Create schema for models.
+---
 DROP SCHEMA pgml CASCADE;
 CREATE SCHEMA IF NOT EXISTS pgml;
 
+CREATE OR REPLACE FUNCTION pgml.auto_updated_at(tbl regclass) 
+RETURNS VOID 
+AS $$
+    DECLARE name_parts TEXT[];
+    DECLARE name TEXT; 
+BEGIN
+    name_parts := string_to_array(tbl::TEXT, '.');
+    name := name_parts[array_upper(name_parts, 1)];
+
+    EXECUTE format('DROP TRIGGER IF EXISTS %s_auto_updated_at ON %s', name, tbl);
+    EXECUTE format('CREATE TRIGGER %s_auto_updated_at BEFORE UPDATE ON %s
+                    FOR EACH ROW EXECUTE PROCEDURE pgml.set_updated_at()', name, tbl);
+END;
+$$
+LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION pgml.set_updated_at() 
+RETURNS TRIGGER 
+AS $$
+BEGIN
+    IF (
+        NEW IS DISTINCT FROM OLD
+        AND NEW.updated_at IS NOT DISTINCT FROM OLD.updated_at
+    ) THEN
+        NEW.updated_at := CURRENT_TIMESTAMP;
+    END IF;
+    RETURN new;
+END;
+$$
+LANGUAGE plpgsql;
+
+CREATE TABLE pgml.projects(
+	id BIGSERIAL PRIMARY KEY,
+	name TEXT NOT NULL,
+	created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+	updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP
+);
+SELECT pgml.auto_updated_at('pgml.projects');
+
+CREATE TABLE pgml.snapshots(
+	id BIGSERIAL PRIMARY KEY,
+	relation TEXT NOT NULL,
+	y TEXT NOT NULL,
+	validation_ratio FLOAT4 NOT NULL,
+	validation_strategy TEXT NOT NULL,
+	created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+	updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP
+);
+SELECT pgml.auto_updated_at('pgml.snapshots');
+
+CREATE TABLE pgml.models(
+	id BIGSERIAL PRIMARY KEY,
+	project_id BIGINT,
+	snapshot_id BIGINT,
+	created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+	updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+	pickle BYTEA,
+	CONSTRAINT project_id_fk FOREIGN KEY(project_id) REFERENCES pgml.projects(id),
+	CONSTRAINT snapshot_id_fk FOREIGN KEY(snapshot_id) REFERENCES pgml.snapshots(id)
+);
+SELECT pgml.auto_updated_at('pgml.models');
+
+CREATE TABLE pgml.promotions(
+	project_id BIGINT NOT NULL,
+	model_id BIGINT NOT NULL,
+	created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+	CONSTRAINT project_id_fk FOREIGN KEY(project_id) REFERENCES pgml.projects(id),
+	CONSTRAINT model_id_fk FOREIGN KEY(model_id) REFERENCES pgml.models(id)
+);
+CREATE INDEX promotions_project_id_created_at_idx ON pgml.promotions(project_id, created_at);
+SELECT pgml.auto_updated_at('pgml.promotions');
+
+
 ---
 --- Extension version.
 ---
@@ -15,20 +92,28 @@ AS $$
 	return pgml.version()
 $$ LANGUAGE plpython3u;
 
+CREATE OR REPLACE FUNCTION pgml.model_regression(model_name TEXT, relation_name TEXT, y_column_name TEXT, algorithm TEXT)
+RETURNS VOID
+AS $$
+	import pgml
+	pgml.model.regression(model_name, relation_name, y_column_name, algorithm)
+$$ LANGUAGE plpython3u;
+
+
 ---
 --- Track table versions.
 ---
 CREATE TABLE pgml.model_versions(
 	id BIGSERIAL PRIMARY KEY,
-	name VARCHAR,
-	location VARCHAR NULL,
+	name VARCHAR NOT NULL,
 	data_source TEXT,
 	y_column VARCHAR,
 	started_at TIMESTAMP WITHOUT TIME ZONE DEFAULT CURRENT_TIMESTAMP,
 	ended_at TIMESTAMP WITHOUT TIME ZONE NULL,
 	mean_squared_error DOUBLE PRECISION,
 	r2_score DOUBLE PRECISION,
-	successful BOOL NULL
+	successful BOOL NULL,
+	pickle BYTEA 
 );
 
 ---
@@ -54,14 +139,14 @@ AS $$
 	id_ = start[0]["id"]
 	name = f"{table_name}_{id_}"
 
-	destination = models_directory(plpy)
+	destination = models_directory()
 
 	# Train!
-	location, msq, r2 = train(plpy.cursor(data_source), y_column=y, name=name, destination=destination)
+	pickle, msq, r2 = train(plpy.cursor(data_source), y_column=y, name=name, destination=destination)
 
 	plpy.execute(f"""
 		UPDATE pgml.model_versions
-		SET location = '{location}',
+		SET pickle = '{pickle}',
 			successful = true,
 			mean_squared_error = '{msq}',
 			r2_score = '{r2}',
@@ -85,7 +170,7 @@ AS $$
 	if model_name in SD:
 		model = SD[model_name]
 	else:
-		SD[model_name] = load(model_name, models_directory(plpy))
+		SD[model_name] = load(model_name, models_directory())
 		model = SD[model_name]
 
 	scores = model.predict([features,])
diff --git a/sql/test.sql b/sql/test.sql
@@ -6,9 +6,6 @@
 
 SELECT pgml.version();
 
--- Valiate our wine data.
-SELECT pgml.validate('wine_quality_red');
-
 -- Train twice
 SELECT pgml.train('wine_quality_red', 'quality');
 
